Update GetVectorByIds interface (#868)

cydrain · web-flow · commit fbf2b6e63c44 · 2023-05-11T14:01:16.000+08:00
Signed-off-by: Yudong Cai &lt;yudong.cai@zilliz.com&gt;
diff --git a/include/knowhere/index.h b/include/knowhere/index.h
@@ -204,21 +204,8 @@ class Index {
     }
 
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Json& json) const {
-        Json json_(json);
-        auto cfg = this->node->CreateConfig();
-        {
-            auto res = Config::FormatAndCheck(*cfg, json_);
-            LOG_KNOWHERE_DEBUG_ << "GetVectorByIds config dump: " << json_.dump();
-            if (res != Status::success) {
-                return unexpected(res);
-            }
-        }
-        auto res = Config::Load(*cfg, json_, knowhere::SEARCH);
-        if (res != Status::success) {
-            return unexpected(res);
-        }
-        return this->node->GetVectorByIds(dataset, *cfg);
+    GetVectorByIds(const DataSet& dataset) const {
+        return this->node->GetVectorByIds(dataset);
     }
 
     bool
diff --git a/include/knowhere/index_node.h b/include/knowhere/index_node.h
@@ -39,7 +39,7 @@ class IndexNode : public Object {
     RangeSearch(const DataSet& dataset, const Config& cfg, const BitsetView& bitset) const = 0;
 
     virtual expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const = 0;
+    GetVectorByIds(const DataSet& dataset) const = 0;
 
     virtual bool
     HasRawData(const std::string& metric_type) const = 0;
diff --git a/include/knowhere/index_node_thread_pool_wrapper.h b/include/knowhere/index_node_thread_pool_wrapper.h
@@ -53,8 +53,8 @@ class IndexNodeThreadPoolWrapper : public IndexNode {
     }
 
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const {
-        return index_node_->GetVectorByIds(dataset, cfg);
+    GetVectorByIds(const DataSet& dataset) const {
+        return index_node_->GetVectorByIds(dataset);
     }
 
     bool
diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i
@@ -141,8 +141,8 @@ class IndexWrap {
     }
 
     knowhere::DataSetPtr
-    GetVectorByIds(knowhere::DataSetPtr dataset, const std::string& json, knowhere::Status& status) {
-        auto res = idx.GetVectorByIds(*dataset, knowhere::Json::parse(json));
+    GetVectorByIds(knowhere::DataSetPtr dataset, knowhere::Status& status) {
+        auto res = idx.GetVectorByIds(*dataset);
         if (res.has_value()) {
             status = knowhere::Status::success;
             return res.value();
diff --git a/src/index/diskann/diskann.cc b/src/index/diskann/diskann.cc
@@ -61,7 +61,7 @@ class DiskANNIndexNode : public IndexNode {
     RangeSearch(const DataSet& dataset, const Config& cfg, const BitsetView& bitset) const override;
 
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override;
+    GetVectorByIds(const DataSet& dataset) const override;
 
     bool
     HasRawData(const std::string& metric_type) const override {
@@ -625,10 +625,7 @@ DiskANNIndexNode<T>::RangeSearch(const DataSet& dataset, const Config& cfg, cons
 
 template <typename T>
 expected<DataSetPtr, Status>
-DiskANNIndexNode<T>::GetVectorByIds(const DataSet& dataset, const Config& cfg) const {
-    if (!is_prepared_.load()) {
-        const_cast<DiskANNIndexNode<T>*>(this)->Prepare(cfg);
-    }
+DiskANNIndexNode<T>::GetVectorByIds(const DataSet& dataset) const {
     if (!is_prepared_.load() || !pq_flash_index_) {
         LOG_KNOWHERE_ERROR_ << "Failed to load diskann.";
         return unexpected(Status::empty_index);
diff --git a/src/index/flat/flat.cc b/src/index/flat/flat.cc
@@ -201,7 +201,7 @@ class FlatIndexNode : public IndexNode {
     }
 
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {
+    GetVectorByIds(const DataSet& dataset) const override {
         auto dim = Dim();
         auto rows = dataset.GetRows();
         auto ids = dataset.GetIds();
diff --git a/src/index/gpu/flat_gpu/flat_gpu.cc b/src/index/gpu/flat_gpu/flat_gpu.cc
@@ -95,7 +95,7 @@ class GpuFlatIndexNode : public IndexNode {
     }
 
     virtual expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {
+    GetVectorByIds(const DataSet& dataset) const override {
         DataSetPtr results = std::make_shared<DataSet>();
         auto nq = dataset.GetRows();
         auto dim = dataset.GetDim();
diff --git a/src/index/gpu/ivf_gpu/ivf_gpu.cc b/src/index/gpu/ivf_gpu/ivf_gpu.cc
@@ -170,7 +170,7 @@ class GpuIvfIndexNode : public IndexNode {
     }
 
     virtual expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {
+    GetVectorByIds(const DataSet& dataset) const override {
         return unexpected(Status::not_implemented);
     }
 
diff --git a/src/index/hnsw/hnsw.cc b/src/index/hnsw/hnsw.cc
@@ -251,7 +251,7 @@ class HnswIndexNode : public IndexNode {
     }
 
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {
+    GetVectorByIds(const DataSet& dataset) const override {
         if (!index_) {
             return unexpected(Status::empty_index);
         }
diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc
@@ -61,7 +61,7 @@ class IvfIndexNode : public IndexNode {
     expected<DataSetPtr, Status>
     RangeSearch(const DataSet& dataset, const Config& cfg, const BitsetView& bitset) const override;
     expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override;
+    GetVectorByIds(const DataSet& dataset) const override;
     bool
     HasRawData(const std::string& metric_type) const override {
         if constexpr (std::is_same<faiss::IndexIVFFlat, T>::value) {
@@ -525,7 +525,7 @@ IvfIndexNode<T>::RangeSearch(const DataSet& dataset, const Config& cfg, const Bi
 
 template <typename T>
 expected<DataSetPtr, Status>
-IvfIndexNode<T>::GetVectorByIds(const DataSet& dataset, const Config& cfg) const {
+IvfIndexNode<T>::GetVectorByIds(const DataSet& dataset) const {
     if (!this->index_) {
         return unexpected(Status::empty_index);
     }
diff --git a/src/index/ivf_raft/ivf_raft.cuh b/src/index/ivf_raft/ivf_raft.cuh
@@ -553,7 +553,7 @@ class RaftIvfIndexNode : public IndexNode {
     }
 
     virtual expected<DataSetPtr, Status>
-    GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {
+    GetVectorByIds(const DataSet& dataset) const override {
         return unexpected(Status::not_implemented);
     }
 
diff --git a/tests/ut/test_diskann.cc b/tests/ut/test_diskann.cc
@@ -274,8 +274,10 @@ TEST_CASE("Test DiskANNIndexNode.", "[diskann]") {
             auto diskann = knowhere::IndexFactory::Instance().Create("DISKANN", diskann_index_pack);
             auto knn_search_json = knn_search_gen().dump();
             knowhere::Json knn_json = knowhere::Json::parse(knn_search_json);
+            auto res = diskann.Search(*query_ds, knn_json, nullptr);
+            REQUIRE(res.has_value());
             auto ids_ds = GenIdsDataSet(kNumRows, kDim);
-            auto results = diskann.GetVectorByIds(*ids_ds, knn_json);
+            auto results = diskann.GetVectorByIds(*ids_ds);
             REQUIRE(results.has_value());
             auto xb = (float*)base_ds->GetTensor();
             auto data = (float*)results.value()->GetTensor();
diff --git a/tests/ut/test_get_vector.cc b/tests/ut/test_get_vector.cc
@@ -61,7 +61,7 @@ TEST_CASE("Test Binary Get Vector By Ids", "[Binary GetVectorByIds]") {
 
         auto idx_new = knowhere::IndexFactory::Instance().Create(name);
         idx_new.Deserialize(bs);
-        auto results = idx_new.GetVectorByIds(*ids_ds, json);
+        auto results = idx_new.GetVectorByIds(*ids_ds);
         REQUIRE(results.has_value());
         auto xb = (uint8_t*)train_ds->GetTensor();
         auto res_rows = results.value()->GetRows();
@@ -164,7 +164,7 @@ TEST_CASE("Test Float Get Vector By Ids", "[Float GetVectorByIds]") {
         if (!idx_new.HasRawData(metric)) {
             return;
         }
-        auto results = idx_new.GetVectorByIds(*ids_ds, json);
+        auto results = idx_new.GetVectorByIds(*ids_ds);
         REQUIRE(results.has_value());
         auto xb = (float*)train_ds_copy->GetTensor();
         auto res_rows = results.value()->GetRows();

Original file line number	Diff line number	Diff line change
`@@ -53,8 +53,8 @@ class IndexNodeThreadPoolWrapper : public IndexNode {`
`53`	`53`	`}`
`54`	`54`
`55`	`55`	`expected<DataSetPtr, Status>`
`56`		`- GetVectorByIds(const DataSet& dataset, const Config& cfg) const {`
`57`		`- return index_node_->GetVectorByIds(dataset, cfg);`
	`56`	`+ GetVectorByIds(const DataSet& dataset) const {`
	`57`	`+ return index_node_->GetVectorByIds(dataset);`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`bool`
Original file line number	Diff line number	Diff line change
`@@ -201,7 +201,7 @@ class FlatIndexNode : public IndexNode {`
`201`	`201`	`}`
`202`	`202`
`203`	`203`	`expected<DataSetPtr, Status>`
`204`		`- GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {`
	`204`	`+ GetVectorByIds(const DataSet& dataset) const override {`
`205`	`205`	`auto dim = Dim();`
`206`	`206`	`auto rows = dataset.GetRows();`
`207`	`207`	`auto ids = dataset.GetIds();`
Original file line number	Diff line number	Diff line change
`@@ -95,7 +95,7 @@ class GpuFlatIndexNode : public IndexNode {`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`virtual expected<DataSetPtr, Status>`
`98`		`- GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {`
	`98`	`+ GetVectorByIds(const DataSet& dataset) const override {`
`99`	`99`	`DataSetPtr results = std::make_shared<DataSet>();`
`100`	`100`	`auto nq = dataset.GetRows();`
`101`	`101`	`auto dim = dataset.GetDim();`
Original file line number	Diff line number	Diff line change
`@@ -170,7 +170,7 @@ class GpuIvfIndexNode : public IndexNode {`
`170`	`170`	`}`
`171`	`171`
`172`	`172`	`virtual expected<DataSetPtr, Status>`
`173`		`- GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {`
	`173`	`+ GetVectorByIds(const DataSet& dataset) const override {`
`174`	`174`	`return unexpected(Status::not_implemented);`
`175`	`175`	`}`
`176`	`176`
Original file line number	Diff line number	Diff line change
`@@ -251,7 +251,7 @@ class HnswIndexNode : public IndexNode {`
`251`	`251`	`}`
`252`	`252`
`253`	`253`	`expected<DataSetPtr, Status>`
`254`		`- GetVectorByIds(const DataSet& dataset, const Config& cfg) const override {`
	`254`	`+ GetVectorByIds(const DataSet& dataset) const override {`
`255`	`255`	`if (!index_) {`
`256`	`256`	`return unexpected(Status::empty_index);`
`257`	`257`	`}`