Support get vector data from index for HNSW (#200)

Signed-off-by: yudong.cai <[email protected]>
milvus-io · May 31, 2022 · 7165744 · 7165744
1 parent a993cf2
commit 7165744
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/knowhere/index/vector_index/IndexHNSW.cpp b/knowhere/index/vector_index/IndexHNSW.cpp
@@ -25,14 +25,6 @@
 
 namespace knowhere {
 
-// void
-// normalize_vector(float* data, float* norm_array, size_t dim) {
-//     float norm = 0.0f;
-//     for (int i = 0; i < dim; i++) norm += data[i] * data[i];
-//     norm = 1.0f / (sqrtf(norm) + 1e-30f);
-//     for (int i = 0; i < dim; i++) norm_array[i] = data[i] * norm;
-// }
-
 BinarySet
 IndexHNSW::Serialize(const Config& config) {
     if (!index_) {
@@ -127,6 +119,21 @@ IndexHNSW::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
     // LOG_KNOWHERE_DEBUG_ << GetStatistics()->ToString();
 }
 
+DatasetPtr
+IndexHNSW::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
+    if (!index_) {
+        KNOWHERE_THROW_MSG("index not initialize");
+    }
+
+    GET_DATA_WITH_IDS(dataset_ptr)
+
+    float* p_x = (float*)malloc(sizeof(float) * dim * rows);
+    for (int64_t i = 0; i < rows; i++) {
+        memcpy(p_x + i * dim, index_->getDataByInternalId(p_ids[i]), dim * sizeof(float));
+    }
+    return GenResultDataset(p_x);
+}
+
 DatasetPtr
 IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
     if (!index_) {

diff --git a/knowhere/index/vector_index/IndexHNSW.h b/knowhere/index/vector_index/IndexHNSW.h
@@ -39,6 +39,9 @@ class IndexHNSW : public VecIndex {
     void
     AddWithoutIds(const DatasetPtr&, const Config&) override;
 
+    DatasetPtr
+    GetVectorById(const DatasetPtr&, const Config&) override;
+
     DatasetPtr
     Query(const DatasetPtr&, const Config&, const faiss::BitsetView) override;
 

diff --git a/unittest/test_hnsw.cpp b/unittest/test_hnsw.cpp
@@ -74,11 +74,14 @@ TEST_P(HNSWTest, HNSW_basic) {
 
     index_->Load(bs);
 
+    auto result = index_->GetVectorById(id_dataset, conf_);
+    AssertVec(result, base_dataset, id_dataset, nq, dim);
+
     auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_type_);
     ASSERT_TRUE(adapter->CheckSearch(conf_, index_type_, index_mode_));
 
-    auto result = index_->Query(query_dataset, conf_, nullptr);
-    AssertAnns(result, nq, k);
+    auto result1 = index_->Query(query_dataset, conf_, nullptr);
+    AssertAnns(result1, nq, k);
 
     auto result2 = index_->Query(query_dataset, conf_, *bitset);
     AssertAnns(result2, nq, k, CheckMode::CHECK_NOT_EQUAL);