feat: pybind wrapper for KDTree

Author: Seco1024

Makefile

@@ -9,7 +9,7 @@ PROJECT_INCLUDE := -I./include -I./include/zenann
 ALL_INCLUDES := $(PYBIND11_INCLUDES) $(PYTHON_INCLUDE) $(PROJECT_INCLUDE)
 ALL_LIBS := $(PYTHON_LIB)
 
-SOURCES := src/IndexBase.cpp src/IVFFlatIndex.cpp python/zenann_pybind.cpp
+SOURCES := src/IndexBase.cpp src/IVFFlatIndex.cpp src/KDTreeIndex.cpp python/zenann_pybind.cpp
 EXT_SUFFIX := $(shell python3-config --extension-suffix)
 TARGET := build/zenann$(EXT_SUFFIX)
 

python/zenann_pybind.cpp

@@ -3,6 +3,7 @@
 
 #include "IndexBase.h"
 #include "IVFFlatIndex.h"
+#include "KDTreeIndex.h"
 
 namespace py = pybind11;
 using namespace zenann;
@@ -11,9 +12,7 @@ using namespace zenann;
 struct PyIndexBase : IndexBase {
     using IndexBase::IndexBase;
     void train() override { }
-    SearchResult search(const Vector&, size_t) const override {
-        return {};
-    }
+    SearchResult search(const Vector&, size_t) const override { return {}; }
 };
 
 PYBIND11_MODULE(zenann, m) {
@@ -24,7 +23,7 @@ PYBIND11_MODULE(zenann, m) {
         .def_readonly("indices", &SearchResult::indices)
         .def_readonly("distances", &SearchResult::distances);
 
-    // Bind IndexBase with shared_ptr holder
+    // Bind IndexBase
     py::class_<IndexBase, PyIndexBase, std::shared_ptr<IndexBase>>(m, "IndexBase")
         .def(py::init<size_t>(), py::arg("dim"))
         .def("build", &IndexBase::build, py::arg("data"),
@@ -34,9 +33,9 @@ PYBIND11_MODULE(zenann, m) {
         .def("search", &IndexBase::search, py::arg("query"), py::arg("k"),
              "Search k nearest neighbors")
         .def_property_readonly("dimension", &IndexBase::dimension,
-             "Dimension of vectors in the index");
+             "Dimension of vectors");
 
-    // Bind IVFFlatIndex as a subclass of IndexBase
+    // 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)
@@ -45,7 +44,25 @@ PYBIND11_MODULE(zenann, m) {
         .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 using IVF index")
-        .def("search_batch", &IVFFlatIndex::search_batch, py::arg("queries"), py::arg("k"), 
-        "Search top-k for a batch of queries");
+             "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");
+
+    py::class_<KDTreeIndex, IndexBase, std::shared_ptr<KDTreeIndex>>(m, "KDTreeIndex")
+        .def(py::init<size_t>(), py::arg("dim"))
+        .def("build", &KDTreeIndex::build, py::arg("data"),
+             "Add data and build KDTree index")
+        .def("train", &KDTreeIndex::train,
+             "Rebuild KDTree index")
+        .def("search", &KDTreeIndex::search, py::arg("query"), py::arg("k"),
+             "Exact k-NN search using KDTree")
+        .def("write_index", &KDTreeIndex::write_index, py::arg("filename"),
+             "Serialize KDTree index to file")
+        .def_static("read_index", &KDTreeIndex::read_index, py::arg("filename"),
+             "Load KDTree index from file");
 }