Use preloaded matrices as input for kmeans query call (#47)

Author: ihnorton

apis/python/src/tiledb/vector_search/index.py

@@ -92,10 +92,10 @@ def __init__(self, uri, dtype: np.dtype):
         self.dtype = dtype
 
         ctx = Ctx({})  # TODO pass in a context
-        # TODO self._db = load_as_matrix(self.db_uri)
+        self._db = load_as_matrix(self.parts_db_uri)
         self._centroids = load_as_matrix(self.centroids_uri)
         self._index = read_vector_u64(ctx, self.index_uri)
-        # self._ids = load_as_matrix(self.ids_uri)
+        self._ids = read_vector_u64(ctx, self.ids_uri)
 
     def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
         """
@@ -122,12 +122,12 @@ def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
         targets_m_a[:] = targets
 
         r = query_kmeans(
-            self.dtype,
-            self.parts_db_uri,
+            self._db.dtype,
+            self._db,
             self._centroids,
             targets_m,
             self._index,
-            self.ids_uri,
+            self._ids,
             nprobe=nprobe,
             k_nn=k,
             nth=True,  # ??

apis/python/src/tiledb/vector_search/module.cc

@@ -74,26 +74,40 @@ static void declareColMajorMatrix(py::module& mod, std::string const& suffix) {
 
 }
 
+template <typename T>
+static void declare_pyarray_to_matrix(py::module& m, const std::string& suffix) {
+  m.def(("pyarray_copyto_matrix" + suffix).c_str(),
+      [](py::array_t<T> arr) -> ColMajorMatrix<T> {
+        py::buffer_info info = arr.request();
+        if (info.ndim != 2)
+          throw std::runtime_error("Number of dimensions must be two");
+        if (info.format != py::format_descriptor<T>::format())
+          throw std::runtime_error("Mismatched buffer format!");
+
+        auto data = std::unique_ptr<T[]>{new T[info.shape[0] * info.shape[1]]};
+        auto r = ColMajorMatrix<T>(std::move(data), info.shape[0], info.shape[1]);
+        return r;
+        });
+}
+
 template <typename T>
 static void declare_kmeans_query(py::module& m, const std::string& suffix) {
   m.def(("kmeans_query_" + suffix).c_str(),
-      [](Ctx ctx,
-         const std::string& part_uri,
+      [](const ColMajorMatrix<T>& parts,
          const ColMajorMatrix<float>& centroids,
          const ColMajorMatrix<float>& query_vectors,
          std::vector<uint64_t>& indices,
-         const std::string& id_uri,
+         std::vector<shuffled_ids_type>& ids,
          size_t nprobe,
          size_t k_nn,
          bool nth,
-         size_t nthreads) -> ColMajorMatrix<size_t> {
+         size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
         auto r = detail::ivf::qv_query_heap_infinite_ram<T>(
-            ctx,
-            part_uri,
+            parts,
             centroids,
             query_vectors,
             indices,
-            id_uri,
+            ids,
             nprobe,
             k_nn,
             nth,
@@ -102,22 +116,6 @@ static void declare_kmeans_query(py::module& m, const std::string& suffix) {
         }, py::keep_alive<1,2>());
 }
 
-template <typename T>
-static void declare_pyarray_to_matrix(py::module& m, const std::string& suffix) {
-  m.def(("pyarray_copyto_matrix" + suffix).c_str(),
-      [](py::array_t<T> arr) -> ColMajorMatrix<T> {
-        py::buffer_info info = arr.request();
-        if (info.ndim != 2)
-          throw std::runtime_error("Number of dimensions must be two");
-        if (info.format != py::format_descriptor<T>::format())
-          throw std::runtime_error("Mismatched buffer format!");
-
-        auto data = std::unique_ptr<T[]>{new T[info.shape[0] * info.shape[1]]};
-        auto r = ColMajorMatrix<T>(std::move(data), info.shape[0], info.shape[1]);
-        return r;
-        });
-}
-
 
 // Declarations for typed subclasses of ColMajorMatrix
 template <typename P>

apis/python/src/tiledb/vector_search/module.py

@@ -84,11 +84,11 @@ def query_vq(db: "colMajorMatrix", *args):
 
 def query_kmeans(
     dtype: np.dtype,
-    parts_uri: str,
+    parts_db: "colMajorMatrix",
     centroids_db: "colMajorMatrix",
     query_vectors: "colMajorMatrix",
-    index_db: "Vector",
-    ids_uri: str,
+    indices: "Vector",
+    ids: "Vector",
     nprobe: int,
     k_nn: int,
     nth: bool,
@@ -128,12 +128,11 @@ def query_kmeans(
 
     args = tuple(
         [
-            ctx,
-            parts_uri,
+            parts_db,
             centroids_db,
             query_vectors,
-            index_db,
-            ids_uri,
+            indices,
+            ids,
             nprobe,
             k_nn,
             nth,