Caiyd 1689 fix sq8h fail (#1723)

* #1689 fix SQ8H search fail on SIFT-1B dataset
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix bug
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix unittest
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix gpu search fail
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix db_test
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* revert segment back
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix test_gpuresource
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix unittest
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

* fix test_gpuresource
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

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

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();

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();

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;
 }
 

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>;

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

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);
 

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;
 

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>;

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;
 }
 

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
+size_t
+ConcurrentBitset::capacity() {
+    return capacity_;
+}
+
+size_t
 ConcurrentBitset::size() {
-    return size_;
+    return ((capacity_ + 8 - 1) >> 3);
 }
 
-const unsigned char*
+const uint8_t*
 ConcurrentBitset::bitset() {
-    return reinterpret_cast<const unsigned char*>(bitset_.data());
+    return reinterpret_cast<const uint8_t*>(bitset_.data());
 }
 
 }  // namespace faiss

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_;
 
 };
 

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();

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_));
 

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