Add iterator support to pyknowhere interface.

Signed-off-by: Buqian Zheng <[email protected]>
zilliztech · Nov 13, 2023 · 5afdc61 · 5afdc61
1 parent f7fd047
commit 5afdc61
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 2 deletions.
diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i
@@ -51,11 +51,13 @@ import_array();
 %include <std_pair.i>
 %include <std_map.i>
 %include <std_shared_ptr.i>
+%include <std_vector.i>
 %include <exception.i>
 %shared_ptr(knowhere::DataSet)
 %shared_ptr(knowhere::BinarySet)
 %template(DataSetPtr) std::shared_ptr<knowhere::DataSet>;
 %template(BinarySetPtr) std::shared_ptr<knowhere::BinarySet>;
+%template(int64_float_pair) std::pair<long long int, float>;
 %include <knowhere/expected.h>
 %include <knowhere/dataset.h>
 %include <knowhere/binaryset.h>
@@ -104,7 +106,7 @@ del Enum
 %inline %{
 
 class GILReleaser {
-public:
+ public:
     GILReleaser() : save(PyEval_SaveThread()) {
     }
     ~GILReleaser() {
@@ -113,6 +115,25 @@ public:
     PyThreadState* save;
 };
 
+class AnnIterator {
+ public:
+    AnnIterator(std::shared_ptr<IndexNode::iterator> it = nullptr) : it_(it) {
+    }
+    ~AnnIterator() {
+    }
+
+    bool HasNext() {
+        return it_ && it_->HasNext();
+    }
+
+    std::pair<int64_t, float> Next() {
+        return it_->Next();
+    }
+
+ private:
+    std::shared_ptr<IndexNode::iterator> it_;
+};
+
 class IndexWrap {
  public:
     IndexWrap(const std::string& name, const int32_t& version) {
@@ -157,6 +178,22 @@ class IndexWrap {
         }
     }
 
+    std::vector<AnnIterator>
+    GetAnnIterator(knowhere::DataSetPtr dataset, const std::string& json, const knowhere::BitsetView& bitset, knowhere::Status& status) {
+        GILReleaser rel;
+        auto res = idx.AnnIterator(*dataset, knowhere::Json::parse(json), bitset);
+        if (!res.has_value()) {
+            status = res.error();
+            return std::vector<AnnIterator>();
+        }
+        status = knowhere::Status::success;
+        std::vector<AnnIterator> result;
+        for (auto it : res.value()) {
+            result.emplace_back(it);
+        }
+        return result;
+    }
+
     knowhere::DataSetPtr
     RangeSearch(knowhere::DataSetPtr dataset, const std::string& json, const knowhere::BitsetView& bitset, knowhere::Status& status){
         GILReleaser rel;
@@ -468,3 +505,5 @@ SetSimdType(const std::string type) {
 }
 
 %}
+
+%template(AnnIteratorVector) std::vector<AnnIterator>;
diff --git a/python/setup.py b/python/setup.py
@@ -90,7 +90,7 @@ def get_readme():
     description=(
         "A library for efficient similarity search and clustering of vectors."
     ),
-    url="https://github.com/milvus-io/knowhere",
+    url="https://github.com/zilliztech/knowhere",
     author="Milvus Team",
     author_email="[email protected]",
     license='Apache License 2.0',

diff --git a/src/common/config.cc b/src/common/config.cc
@@ -25,6 +25,7 @@ static const std::unordered_set<std::string> ext_legal_json_keys = {"metric_type
                                                                     "M",               // HNSW param
                                                                     "efConstruction",  // HNSW param
                                                                     "ef",              // HNSW param
+                                                                    "seed_ef",         // HNSW iterator param
                                                                     "level",
                                                                     "index_type",
                                                                     "index_mode",