Support get vector data from index for Annoy (#201)

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

Hide file tree

Showing 3 changed files with 28 additions and 7 deletions.
diff --git a/knowhere/index/vector_index/IndexAnnoy.cpp b/knowhere/index/vector_index/IndexAnnoy.cpp
@@ -105,6 +105,21 @@ IndexAnnoy::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
     index_->build(GetIndexParamNtrees(config));
 }
 
+DatasetPtr
+IndexAnnoy::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++) {
+        index_->get_item(p_ids[i], p_x + i * dim);
+    }
+    return GenResultDataset(p_x);
+}
+
 DatasetPtr
 IndexAnnoy::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) {
     if (!index_) {

diff --git a/knowhere/index/vector_index/IndexAnnoy.h b/knowhere/index/vector_index/IndexAnnoy.h
@@ -31,16 +31,16 @@ class IndexAnnoy : public VecIndex {
     }
 
     BinarySet
-    Serialize(const Config& config) override;
+    Serialize(const Config&) override;
 
     void
-    Load(const BinarySet& index_binary) override;
+    Load(const BinarySet&) override;
 
     void
-    BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override;
+    BuildAll(const DatasetPtr&, const Config&) override;
 
     void
-    Train(const DatasetPtr& dataset_ptr, const Config& config) override {
+    Train(const DatasetPtr&, const Config&) override {
         KNOWHERE_THROW_MSG("Annoy not support build item dynamically, please invoke BuildAll interface.");
     }
 
@@ -50,7 +50,10 @@ class IndexAnnoy : public VecIndex {
     }
 
     DatasetPtr
-    Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView bitset) override;
+    GetVectorById(const DatasetPtr&, const Config&) override;
+
+    DatasetPtr
+    Query(const DatasetPtr&, const Config&, const faiss::BitsetView) override;
 
     int64_t
     Count() override;

diff --git a/unittest/test_annoy.cpp b/unittest/test_annoy.cpp
@@ -64,11 +64,14 @@ TEST_P(AnnoyTest, annoy_basic) {
     ASSERT_EQ(index_->Dim(), dim);
     ASSERT_GT(index_->Size(), 0);
 
+    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);
 }
 
 TEST_P(AnnoyTest, annoy_delete) {