diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8398f0cd46..5390d6edf9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ Please mark all change in change log and use the issue from GitHub
- \#1651 Check validity of dimension when collection metric type is binary one
- \#1663 PQ index parameter 'm' validation
- \#1686 API search_in_files cannot work correctly when vectors is stored in certain non-default partition
+- \#1689 Fix SQ8H search fail on SIFT-1B dataset
## Feature
- \#1603 BinaryFlat add 2 Metric: Substructure and Superstructure
diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp
index 5802557f01..0a77aa18a5 100644
--- a/core/src/db/engine/ExecutionEngineImpl.cpp
+++ b/core/src/db/engine/ExecutionEngineImpl.cpp
@@ -352,7 +352,7 @@ ExecutionEngineImpl::Serialize() {
// here we reset index size by file size,
// since some index type(such as SQ8) data size become smaller after serialized
- index_->set_size(server::CommonUtil::GetFileSize(location_));
+ index_->SetIndexSize(server::CommonUtil::GetFileSize(location_));
ENGINE_LOG_DEBUG << "Finish serialize index file: " << location_ << " size: " << index_->Size();
if (index_->Size() == 0) {
@@ -415,36 +415,22 @@ ExecutionEngineImpl::Load(bool to_cache) {
}
}
+ auto dataset = knowhere::GenDataset(vectors->GetCount(), this->dim_, vectors_data.data());
if (index_type_ == EngineType::FAISS_IDMAP) {
auto bf_index = std::static_pointer_cast<knowhere::IDMAP>(index_);
- std::vector<float> float_vectors;
- float_vectors.resize(vectors_data.size() / sizeof(float));
- memcpy(float_vectors.data(), vectors_data.data(), vectors_data.size());
bf_index->Train(knowhere::DatasetPtr(), conf);
- auto dataset = knowhere::GenDataset(vectors->GetCount(), this->dim_, float_vectors.data());
bf_index->AddWithoutIds(dataset, conf);
bf_index->SetBlacklist(concurrent_bitset_ptr);
} else if (index_type_ == EngineType::FAISS_BIN_IDMAP) {
auto bin_bf_index = std::static_pointer_cast<knowhere::BinaryIDMAP>(index_);
bin_bf_index->Train(knowhere::DatasetPtr(), conf);
- auto dataset = knowhere::GenDataset(vectors->GetCount(), this->dim_, vectors_data.data());
bin_bf_index->AddWithoutIds(dataset, conf);
bin_bf_index->SetBlacklist(concurrent_bitset_ptr);
}
- int64_t index_size = vectors->Size(); // vector data size + vector ids size
- int64_t bitset_size = vectors->GetCount(); // delete list size
- index_->set_size(index_size + bitset_size);
-
- if (!status.ok()) {
- return status;
- }
-
ENGINE_LOG_DEBUG << "Finished loading raw data from segment " << segment_dir;
} else {
try {
- // size_t physical_size = PhysicalSize();
- // server::CollectExecutionEngineMetrics metrics((double)physical_size);
index_ = read_index(location_);
if (index_ == nullptr) {
@@ -476,10 +462,6 @@ ExecutionEngineImpl::Load(bool to_cache) {
index_->SetUids(uids);
ENGINE_LOG_DEBUG << "set uids " << index_->GetUids().size() << " for index " << location_;
- int64_t index_size = index_->Size(); // vector data size + vector ids size
- int64_t bitset_size = index_->Count(); // delete list size
- index_->set_size(index_size + bitset_size);
-
ENGINE_LOG_DEBUG << "Finished loading index file from segment " << segment_dir;
}
} catch (std::exception& e) {
@@ -564,15 +546,12 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) {
try {
index_ = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
ENGINE_LOG_DEBUG << "CPU to GPU" << device_id;
+ GpuCache(device_id);
} catch (std::exception& e) {
ENGINE_LOG_ERROR << e.what();
return Status(DB_ERROR, e.what());
}
}
-
- if (!already_in_cache) {
- GpuCache(device_id);
- }
#endif
return Status::OK();
diff --git a/core/src/db/meta/SqliteMetaImpl.cpp b/core/src/db/meta/SqliteMetaImpl.cpp
index fd04e21270..0c9bd7fb72 100644
--- a/core/src/db/meta/SqliteMetaImpl.cpp
+++ b/core/src/db/meta/SqliteMetaImpl.cpp
@@ -1409,7 +1409,7 @@ SqliteMetaImpl::Size(uint64_t& result) {
result += (uint64_t)(*std::get<0>(total_size));
}
} catch (std::exception& e) {
- return HandleException("Encounter exception when calculte db size", e.what());
+ return HandleException("Encounter exception when calculate db size", e.what());
}
return Status::OK();
diff --git a/core/src/index/archive/VecIndex.cpp b/core/src/index/archive/VecIndex.cpp
index 0d0ecf2347..8225452c08 100644
--- a/core/src/index/archive/VecIndex.cpp
+++ b/core/src/index/archive/VecIndex.cpp
@@ -41,7 +41,7 @@ LoadVecIndex(const knowhere::IndexType& type, const knowhere::BinarySet& index_b
return nullptr;
// else
index->Load(index_binary);
- index->set_size(size);
+ index->SetIndexSize(size);
return index;
}
diff --git a/core/src/index/knowhere/knowhere/index/Index.h b/core/src/index/knowhere/knowhere/index/Index.h
index 1ef8918c7d..87ba3198bf 100644
--- a/core/src/index/knowhere/knowhere/index/Index.h
+++ b/core/src/index/knowhere/knowhere/index/Index.h
@@ -27,19 +27,6 @@ class Index : public milvus::cache::DataObj {
virtual void
Load(const BinarySet&) = 0;
-
- int64_t
- Size() override {
- return size_;
- }
-
- void
- set_size(const int64_t& size) {
- size_ = size;
- }
-
- protected:
- int64_t size_ = -1;
};
using IndexPtr = std::shared_ptr<Index>;
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h
index ab3d08b6ef..4fc5fe7605 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h
+++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h
@@ -64,11 +64,8 @@ class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {
}
int64_t
- Size() override {
- if (size_ != -1) {
- return size_;
- }
- return Count() * Dim() * sizeof(uint8_t);
+ IndexSize() override {
+ return Count() * Dim() / 8;
}
DatasetPtr
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h
index 67848cbf0a..7e201db84f 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h
+++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h
@@ -75,14 +75,6 @@ class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {
return index_->d;
}
- int64_t
- Size() override {
- if (size_ != -1) {
- return size_;
- }
- return Count() * Dim() * sizeof(uint8_t);
- }
-
DatasetPtr
GetVectorById(const DatasetPtr& dataset_ptr, const Config& config);
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h
index 6d3aa39531..e80f9bf92d 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h
+++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h
@@ -61,6 +61,11 @@ class IDMAP : public VecIndex, public FaissBaseIndex {
return index_->d;
}
+ int64_t
+ IndexSize() override {
+ return Count() * Dim() * sizeof(FloatType);
+ }
+
DatasetPtr
GetVectorById(const DatasetPtr& dataset, const Config& config) override;
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h
index 3f2d4b7500..71d5673570 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h
+++ b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h
@@ -17,6 +17,7 @@
#include <vector>
#include "knowhere/common/Dataset.h"
+#include "knowhere/common/Exception.h"
#include "knowhere/common/Typedef.h"
#include "knowhere/index/Index.h"
#include "knowhere/index/vector_index/IndexType.h"
@@ -77,42 +78,65 @@ class VecIndex : public Index {
return nullptr;
}
- virtual void
+ void
GetBlacklist(faiss::ConcurrentBitsetPtr& bitset_ptr) {
bitset_ptr = bitset_;
}
- virtual void
+ void
SetBlacklist(faiss::ConcurrentBitsetPtr bitset_ptr) {
bitset_ = std::move(bitset_ptr);
}
- virtual const std::vector<milvus::segment::doc_id_t>&
+ const std::vector<milvus::segment::doc_id_t>&
GetUids() const {
return uids_;
}
- virtual void
+ void
SetUids(std::vector<milvus::segment::doc_id_t>& uids) {
uids_.clear();
uids_.swap(uids);
}
+ size_t
+ BlacklistSize() {
+ if (bitset_) {
+ return bitset_->size() * sizeof(uint8_t);
+ } else {
+ return 0;
+ }
+ }
+
+ size_t
+ UidsSize() {
+ return uids_.size() * sizeof(segment::doc_id_t);
+ }
+
+ virtual int64_t
+ IndexSize() {
+ if (index_size_ == -1) {
+ KNOWHERE_THROW_MSG("Index size not set");
+ }
+ return index_size_;
+ }
+
+ void
+ SetIndexSize(int64_t size) {
+ index_size_ = size;
+ }
+
int64_t
Size() override {
- if (size_ != -1) {
- return size_;
- }
- return Count() * Dim() * sizeof(FloatType);
+ return BlacklistSize() + UidsSize() + IndexSize();
}
protected:
IndexType index_type_ = "";
IndexMode index_mode_ = IndexMode::MODE_CPU;
faiss::ConcurrentBitsetPtr bitset_ = nullptr;
-
- private:
- std::vector<milvus::segment::doc_id_t> uids_;
+ std::vector<segment::doc_id_t> uids_;
+ int64_t index_size_ = -1;
};
using VecIndexPtr = std::shared_ptr<VecIndex>;
diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
index d6e7532d33..62b2e121aa 100644
--- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
+++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
@@ -30,6 +30,7 @@ CopyGpuToCpu(const VecIndexPtr& index, const Config& config) {
VecIndexPtr result = device_index->CopyGpuToCpu(config);
auto uids = index->GetUids();
result->SetUids(uids);
+ result->SetIndexSize(index->IndexSize());
return result;
} else {
KNOWHERE_THROW_MSG("index type is not gpuindex");
@@ -40,15 +41,18 @@ VecIndexPtr
CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config) {
VecIndexPtr result;
auto uids = index->GetUids();
+ int64_t index_size = index->IndexSize();
if (auto device_index = std::dynamic_pointer_cast<IVFSQHybrid>(index)) {
result = device_index->CopyCpuToGpu(device_id, config);
result->SetUids(uids);
+ result->SetIndexSize(index_size);
return result;
}
if (auto device_index = std::dynamic_pointer_cast<GPUIndex>(index)) {
result = device_index->CopyGpuToGpu(device_id, config);
result->SetUids(uids);
+ result->SetIndexSize(index_size);
return result;
}
@@ -65,6 +69,7 @@ CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& co
}
result->SetUids(uids);
+ result->SetIndexSize(index_size);
return result;
}
diff --git a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp
index bee43bd6a0..25f3d6a66c 100644
--- a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp
+++ b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp
@@ -19,7 +19,7 @@
namespace faiss {
-ConcurrentBitset::ConcurrentBitset(id_type_t size) : size_(size), bitset_((size + 7) >> 3) {
+ConcurrentBitset::ConcurrentBitset(id_type_t capacity) : capacity_(capacity), bitset_((capacity + 8 - 1) >> 3) {
}
bool
@@ -37,14 +37,19 @@ ConcurrentBitset::clear(id_type_t id) {
bitset_[id >> 3].fetch_and(~(0x1 << (id & 0x7)));
}
-ConcurrentBitset::id_type_t
-ConcurrentBitset::size() {
- return size_;
+size_t
+ConcurrentBitset::capacity() {
+ return capacity_;
}
-const unsigned char*
+size_t
+ConcurrentBitset::size() {
+ return ((capacity_ + 8 - 1) >> 3);
+}
+
+const uint8_t*
ConcurrentBitset::bitset() {
- return reinterpret_cast<const unsigned char*>(bitset_.data());
+ return reinterpret_cast<const uint8_t*>(bitset_.data());
}
} // namespace faiss
diff --git a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h
index fff7752c6a..a447041834 100644
--- a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h
+++ b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h
@@ -42,15 +42,18 @@ class ConcurrentBitset {
void
clear(id_type_t id);
- id_type_t
+ size_t
+ capacity();
+
+ size_t
size();
- const unsigned char*
+ const uint8_t*
bitset();
private:
- id_type_t size_;
- std::vector<std::atomic<unsigned char>> bitset_;
+ size_t capacity_;
+ std::vector<std::atomic<uint8_t>> bitset_;
};
diff --git a/core/src/index/unittest/test_gpuresource.cpp b/core/src/index/unittest/test_gpuresource.cpp
index 15fb35e1e5..e70404e176 100644
--- a/core/src/index/unittest/test_gpuresource.cpp
+++ b/core/src/index/unittest/test_gpuresource.cpp
@@ -77,6 +77,7 @@ TEST_F(GPURESTEST, copyandsearch) {
auto result = index_->Query(query_dataset, conf);
AssertAnns(result, nq, k);
+ index_->SetIndexSize(nb * dim * sizeof(float));
auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config());
milvus::knowhere::IVFPtr ivf_idx = std::dynamic_pointer_cast<milvus::knowhere::IVF>(cpu_idx);
ivf_idx->Seal();
@@ -128,6 +129,7 @@ TEST_F(GPURESTEST, trainandsearch) {
auto conf = ParamGenerator::GetInstance().Gen(index_type_);
index_->Train(base_dataset, conf);
index_->Add(base_dataset, conf);
+ index_->SetIndexSize(nb * dim * sizeof(float));
auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config());
milvus::knowhere::IVFPtr ivf_idx = std::dynamic_pointer_cast<milvus::knowhere::IVF>(cpu_idx);
ivf_idx->Seal();
diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp
index 7d85155a63..f318952725 100644
--- a/core/src/index/unittest/test_ivf.cpp
+++ b/core/src/index/unittest/test_ivf.cpp
@@ -186,6 +186,10 @@ TEST_P(IVFTest, clone_test) {
index_->Add(base_dataset, conf_);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dim(), dim);
+
+ /* set peseodo index size, avoid throw exception */
+ index_->SetIndexSize(nq * dim * sizeof(float));
+
auto result = index_->Query(query_dataset, conf_);
AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]);
// PrintResult(result, nq, k);
@@ -269,6 +273,10 @@ TEST_P(IVFTest, gpu_seal_test) {
index_->Add(base_dataset, conf_);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dim(), dim);
+
+ /* set peseodo index size, avoid throw exception */
+ index_->SetIndexSize(nq * dim * sizeof(float));
+
auto result = index_->Query(query_dataset, conf_);
AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]);
@@ -306,6 +314,7 @@ TEST_P(IVFTest, invalid_gpu_source) {
if (index_type_ == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) {
// null faiss index
+ index_->SetIndexSize(0);
milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config());
}
@@ -337,6 +346,7 @@ TEST_P(IVFTest, IVFSQHybrid_test) {
}
fiu_init(0);
+ index_->SetIndexSize(0);
milvus::knowhere::cloner::CopyGpuToCpu(index_, conf_);
ASSERT_ANY_THROW(milvus::knowhere::cloner::CopyCpuToGpu(index_, -1, conf_));
diff --git a/core/unittest/server/test_cache.cpp b/core/unittest/server/test_cache.cpp
index 8e022cb38e..643b95b32b 100644
--- a/core/unittest/server/test_cache.cpp
+++ b/core/unittest/server/test_cache.cpp
@@ -37,7 +37,7 @@ class MockVecIndex : public milvus::knowhere::VecIndex {
public:
MockVecIndex(int64_t dim, int64_t total) : dim_(dim), ntotal_(total) {
int64_t data_size = Dim() * Count() * sizeof(float);
- set_size(data_size);
+ SetIndexSize(data_size);
}
virtual void