From 5afdc610e23b974403f44e886b6dbf8c8c13ee42 Mon Sep 17 00:00:00 2001 From: Buqian Zheng Date: Mon, 13 Nov 2023 17:50:46 +0800 Subject: [PATCH 1/3] Add iterator support to pyknowhere interface. Signed-off-by: Buqian Zheng --- python/knowhere/knowhere.i | 41 +++++++++++++++++++++++++++++++++++++- python/setup.py | 2 +- src/common/config.cc | 1 + 3 files changed, 42 insertions(+), 2 deletions(-) diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i index 18a435d9d..246be528d 100644 --- a/python/knowhere/knowhere.i +++ b/python/knowhere/knowhere.i @@ -51,11 +51,13 @@ import_array(); %include %include %include +%include %include %shared_ptr(knowhere::DataSet) %shared_ptr(knowhere::BinarySet) %template(DataSetPtr) std::shared_ptr; %template(BinarySetPtr) std::shared_ptr; +%template(int64_float_pair) std::pair; %include %include %include @@ -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 it = nullptr) : it_(it) { + } + ~AnnIterator() { + } + + bool HasNext() { + return it_ && it_->HasNext(); + } + + std::pair Next() { + return it_->Next(); + } + + private: + std::shared_ptr it_; +}; + class IndexWrap { public: IndexWrap(const std::string& name, const int32_t& version) { @@ -157,6 +178,22 @@ class IndexWrap { } } + std::vector + 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(); + } + status = knowhere::Status::success; + std::vector 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; diff --git a/python/setup.py b/python/setup.py index 814d200a6..8893901fe 100644 --- 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="milvus-team@zilliz.com", license='Apache License 2.0', diff --git a/src/common/config.cc b/src/common/config.cc index 2e238e911..21444d9b3 100644 --- a/src/common/config.cc +++ b/src/common/config.cc @@ -25,6 +25,7 @@ static const std::unordered_set 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", From d761bf61fd638154b22e8d747f496deaf31dc55f Mon Sep 17 00:00:00 2001 From: Buqian Zheng Date: Fri, 17 Nov 2023 17:07:51 +0800 Subject: [PATCH 2/3] Rename AnnIterator to AnnIteratorWrap Signed-off-by: Buqian Zheng --- python/knowhere/knowhere.i | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i index 246be528d..bce832b1d 100644 --- a/python/knowhere/knowhere.i +++ b/python/knowhere/knowhere.i @@ -115,11 +115,11 @@ class GILReleaser { PyThreadState* save; }; -class AnnIterator { +class AnnIteratorWrap { public: - AnnIterator(std::shared_ptr it = nullptr) : it_(it) { + AnnIteratorWrap(std::shared_ptr it = nullptr) : it_(it) { } - ~AnnIterator() { + ~AnnIteratorWrap() { } bool HasNext() { @@ -178,16 +178,16 @@ class IndexWrap { } } - std::vector + std::vector 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(); + return std::vector(); } status = knowhere::Status::success; - std::vector result; + std::vector result; for (auto it : res.value()) { result.emplace_back(it); } @@ -506,4 +506,4 @@ SetSimdType(const std::string type) { %} -%template(AnnIteratorVector) std::vector; +%template(AnnIteratorWrapVector) std::vector; From 53522418266955678c3fc60233b3d092965e3bea Mon Sep 17 00:00:00 2001 From: Buqian Zheng Date: Wed, 22 Nov 2023 10:58:58 +0800 Subject: [PATCH 3/3] throw exception when nullptr iterator is provided Signed-off-by: Buqian Zheng --- python/knowhere/knowhere.i | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i index bce832b1d..70dd7f7ed 100644 --- a/python/knowhere/knowhere.i +++ b/python/knowhere/knowhere.i @@ -118,12 +118,15 @@ class GILReleaser { class AnnIteratorWrap { public: AnnIteratorWrap(std::shared_ptr it = nullptr) : it_(it) { + if (it_ == nullptr) { + throw std::runtime_error("ann iterator must not be nullptr."); + } } ~AnnIteratorWrap() { } bool HasNext() { - return it_ && it_->HasNext(); + return it_->HasNext(); } std::pair Next() { @@ -182,12 +185,12 @@ class IndexWrap { 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); + std::vector result; if (!res.has_value()) { status = res.error(); - return std::vector(); + return result; } status = knowhere::Status::success; - std::vector result; for (auto it : res.value()) { result.emplace_back(it); }