Use std::vector as opaque type to avoid copy (#65)

* Use array_to_matrix to clean up cruft

* Use std::vector as opaque type to avoid copy

Cleanup

* Don't use span, unnecessary

* Add py buffer support

Author: ihnorton

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

@@ -57,16 +57,11 @@ def query(
         """
         # TODO:
         # - typecheck targets
-        # - don't copy the array
         # - add all the options and query strategies
 
         assert targets.dtype == np.float32
 
-        # TODO: make Matrix constructor from py::array. This is ugly (and copies).
-        # Create a Matrix from the input targets
-        targets_m = ColMajorMatrix_f32(*targets.shape)
-        targets_m_a = np.array(targets_m, copy=False)
-        targets_m_a[:] = targets
+        targets_m = array_to_matrix(targets)
 
         r = query_vq(self._db, targets_m, k, nqueries, nthreads)
         return np.array(r)
@@ -116,10 +111,7 @@ def query(self, targets: np.ndarray, k=10, nqueries=10, nthreads=8, nprobe=1):
         """
         assert targets.dtype == np.float32
 
-        # TODO: use Matrix constructor from py::array
-        targets_m = ColMajorMatrix_f32(*targets.shape)
-        targets_m_a = np.array(targets_m, copy=False)
-        targets_m_a[:] = targets
+        targets_m = array_to_matrix(targets)
 
         r = query_kmeans(
             self._db.dtype,

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

@@ -14,6 +14,9 @@ using Ctx = tiledb::Context;
 bool global_debug = true;
 double global_time_of_interest;
 
+PYBIND11_MAKE_OPAQUE(std::vector<uint32_t>);
+PYBIND11_MAKE_OPAQUE(std::vector<uint64_t>);
+
 namespace {
 
 
@@ -91,18 +94,19 @@ static void declare_pyarray_to_matrix(py::module& m, const std::string& suffix)
         });
 }
 
-template <typename T>
+template <typename T, typename Id_Type = uint64_t>
 static void declare_kmeans_query(py::module& m, const std::string& suffix) {
   m.def(("kmeans_query_" + suffix).c_str(),
       [](const ColMajorMatrix<T>& parts,
          const ColMajorMatrix<float>& centroids,
          const ColMajorMatrix<float>& query_vectors,
-         std::vector<uint64_t>& indices,
-         std::vector<uint64_t>& ids,
+         std::vector<Id_Type> indices,
+         std::vector<Id_Type> ids,
          size_t nprobe,
          size_t k_nn,
          bool nth,
          size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
+
         auto r = detail::ivf::qv_query_heap_infinite_ram(
             parts,
             centroids,
@@ -132,8 +136,29 @@ static void declareColMajorMatrixSubclass(py::module& mod,
   cls.def(py::init<const Ctx&, std::string, size_t>(),  py::keep_alive<1,2>());
 }
 
+template <typename T>
+void declareStdVector(py::module& m) {
+
+  auto name = std::string("IntVector") + typeid(T).name();
+  py::class_<std::vector<T>>(m, name.c_str(), py::buffer_protocol())
+    .def(py::init<>())
+    .def("clear", &std::vector<T>::clear)
+    .def("pop_back", &std::vector<T>::pop_back)
+    .def("__len__", [](const std::vector<T> &v) { return v.size(); })
+    .def_buffer([](std::vector<T> &v) -> py::buffer_info {
+        return py::buffer_info(
+            v.data(),                               /* Pointer to buffer */
+            sizeof(T),                          /* Size of one scalar */
+            py::format_descriptor<T>::format(), /* Python struct-style format descriptor */
+            1,                                      /* Number of dimensions */
+            { v.size() },                 /* Buffer dimensions */
+            { sizeof(T) });
+    });
 }
 
+} // anonymous namespace
+
+
 PYBIND11_MODULE(_tiledbvspy, m) {
 
   py::class_<tiledb::Context> (m, "Ctx", py::module_local())
@@ -150,14 +175,19 @@ PYBIND11_MODULE(_tiledbvspy, m) {
 
   /* === Vector === */
 
-  declareVector<uint32_t>(m, "_u32");
-  declareVector<uint64_t>(m, "_u64");
-  declareVector<float>(m, "_f32");
-  declareVector<double>(m, "_f64");
+  // Must have matching PYBIND11_MAKE_OPAQUE declaration at top of file
+  declareStdVector<uint32_t>(m);
+  declareStdVector<uint64_t>(m);
 
   m.def("read_vector_u32", &read_vector<uint32_t>, "Read a vector from TileDB");
   m.def("read_vector_u64", &read_vector<uint64_t>, "Read a vector from TileDB");
 
+  m.def("_create_vector_u64", []() {
+    auto v = std::vector<uint64_t>(10);
+    // fill vector with range 1:10 using std::iota
+    std::iota(v.begin(), v.begin() + 10, 0);
+    return v;
+  });
 
   /* === Matrix === */
 

apis/python/test/test_api.py

@@ -24,6 +24,9 @@ def test_load_matrix(tmpdir):
     assert np.array_equal(m, data)
     assert np.array_equal(orig_matrix[0, 0], data[0, 0])
 
+def test_vector(tmpdir):
+    v = vspy._create_vector_u64()
+    assert np.array_equal(np.array(v), np.arange(10))
 
 @pytest.mark.skipif(
     not os.path.exists(os.path.expanduser("~/work/proj/vector-search/datasets/sift-andrew/")),
@@ -58,4 +61,4 @@ def test_flat_query():
     assert np.array_equal(np.sort(ra[:k], axis=0), np.sort(g[:k, :nqueries], axis=0))
 
     g_m = vs.load_as_matrix(g_uri)
-    assert vspy.validate_top_k_u64(r, g_m)
+    assert vspy.validate_top_k_u64(r, g_m)

src/include/detail/linalg/matrix.h

@@ -292,6 +292,20 @@ std::string matrix_info(const std::vector<T>& A, const std::string& msg = "") {
   return str;
 }
 
+/**
+ * Print information about a std::span -- overload.
+ * @param A
+ */
+template <class T>
+std::string matrix_info(const std::span<T>& A, const std::string& msg = "") {
+  std::string str = "# " + msg;
+  if (!msg.empty()) {
+    str += ": ";
+  }
+  str += "Shape: (" + std::to_string(A.size()) + " )";
+  return str;
+}
+
 template <class Matrix>
 void debug_matrix(const Matrix& A, const std::string& msg = "") {
   if (global_debug) {