feat: overloading for IVF nprobe

Author: Seco1024

build/zenann.cpython-311-darwin.so

@@ -11,7 +11,9 @@ class IVFFlatIndex : public IndexBase {
     ~IVFFlatIndex() override;
     void train() override;
     SearchResult search(const Vector& query, size_t k) const override;
+    SearchResult search(const Vector& query, size_t k, size_t nprobe) const;
     std::vector<SearchResult> search_batch(const Dataset& queries, size_t k) const;
+    std::vector<SearchResult> search_batch(const Dataset& queries, size_t k, size_t nprobe) const;
     void write_index(const std::string& filename) const;
     static std::shared_ptr<IVFFlatIndex> read_index(const std::string& filename);
 private:

include/zenann/IVFFlatIndex.h

@@ -37,21 +37,33 @@ PYBIND11_MODULE(zenann, m) {
 
     // Bind IVFFlatIndex
     py::class_<IVFFlatIndex, IndexBase, std::shared_ptr<IVFFlatIndex>>(m, "IVFFlatIndex")
-        .def(py::init<size_t, size_t, size_t>(),
-             py::arg("dim"), py::arg("nlist"), py::arg("nprobe") = 1)
-        .def("build", &IVFFlatIndex::build, py::arg("data"),
-             "Add data and train the IVF index")
-        .def("train", &IVFFlatIndex::train,
-             "Train the IVF index (run K-means and build inverted lists)")
-        .def("search", &IVFFlatIndex::search, py::arg("query"), py::arg("k"),
-             "Search top-k nearest neighbors")
-        .def("search_batch", &IVFFlatIndex::search_batch, py::arg("queries"), py::arg("k"),
-             "Search top-k for a batch of queries")
-        // Persistence API
-        .def("write_index", &IVFFlatIndex::write_index, py::arg("filename"),
-             "Serialize the index to a file")
-        .def_static("read_index", &IVFFlatIndex::read_index, py::arg("filename"),
-             "Load an index from a file");
+    .def(py::init<size_t, size_t, size_t>(),
+         py::arg("dim"), py::arg("nlist"), py::arg("nprobe") = 1)
+    .def("build", &IVFFlatIndex::build, py::arg("data"))
+    .def("train", &IVFFlatIndex::train)
+    .def("search",
+         py::overload_cast<const Vector&, size_t>(
+             &IVFFlatIndex::search,
+             py::const_),
+         py::arg("query"), py::arg("k"))
+    .def("search",
+         py::overload_cast<const Vector&, size_t, size_t>(
+             &IVFFlatIndex::search,
+             py::const_),
+         py::arg("query"), py::arg("k"), py::arg("nprobe"))
+    .def("search_batch",
+         py::overload_cast<const Dataset&, size_t>(
+             &IVFFlatIndex::search_batch,
+             py::const_),
+         py::arg("queries"), py::arg("k"))
+    .def("search_batch",
+         py::overload_cast<const Dataset&, size_t, size_t>(
+             &IVFFlatIndex::search_batch,
+             py::const_),
+         py::arg("queries"), py::arg("k"), py::arg("nprobe"))
+
+    .def("write_index", &IVFFlatIndex::write_index, py::arg("filename"))
+    .def_static("read_index", &IVFFlatIndex::read_index, py::arg("filename"));
 
     py::class_<KDTreeIndex, IndexBase, std::shared_ptr<KDTreeIndex>>(m, "KDTreeIndex")
         .def(py::init<size_t>(), py::arg("dim"))

python/zenann_pybind.cpp

@@ -95,6 +95,14 @@ SearchResult IVFFlatIndex::search(const Vector& query, size_t k) const {
     return result;
 }
 
+SearchResult IVFFlatIndex::search(const Vector& query, size_t k, size_t nprobe) const {
+    size_t old = this->nprobe_;
+    const_cast<IVFFlatIndex*>(this)->nprobe_ = nprobe;      
+    SearchResult res = this->search(query, k);    
+    const_cast<IVFFlatIndex*>(this)->nprobe_ = old;    
+    return res;
+}
+
 std::vector<SearchResult> IVFFlatIndex::search_batch(const Dataset& queries, size_t k) const {
     std::vector<SearchResult> results;
     results.reserve(queries.size());
@@ -104,6 +112,14 @@ std::vector<SearchResult> IVFFlatIndex::search_batch(const Dataset& queries, siz
     return results;
 }
 
+std::vector<SearchResult> IVFFlatIndex::search_batch(const Dataset& queries, size_t k, size_t nprobe) const {
+    size_t old = this->nprobe_;
+    const_cast<IVFFlatIndex*>(this)->nprobe_ = nprobe;
+    auto results = this->search_batch(queries, k);  
+    const_cast<IVFFlatIndex*>(this)->nprobe_ = old;
+    return results;
+}
+
 void IVFFlatIndex::kmeans(const Dataset& data, size_t iterations) {
     size_t n = data.size();
     std::mt19937 rng(123);