Merge branch 'main' into npapa/support-ingestion-with-updates

Author: Nikos Papailiou

.github/workflows/quarto-render.yml

@@ -1,7 +1,7 @@
 # Cloned from https://github.com/TileDB-Inc/tiledb-quarto-template
 
 name: Render and deploy Quarto files
-on: 
+on:
   push:
   pull_request:
 
@@ -36,7 +36,7 @@ jobs:
       - name: "Quarto render"
         shell: bash
         run: |
-          pip install quartodoc PyYAML click griffe==0.32.3
+          pip install quartodoc PyYAML click "griffe<0.33" # incompatible
           # create a symlink to the tiledbvcf python package, so it doesn't have to be installed
           #ln -s apis/python/src/tiledb/vector_search
           quartodoc build

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

@@ -75,6 +75,6 @@ def query_internal(
         assert queries.dtype == np.float32
 
         queries_m = array_to_matrix(np.transpose(queries))
-        r = query_vq_heap(self._db, queries_m, self._ids, k, nthreads)
+        d, i = query_vq_heap(self._db, queries_m, self._ids, k, nthreads)
 
-        return np.transpose(np.array(r))
+        return np.transpose(np.array(d)), np.transpose(np.array(i))

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

@@ -1,10 +1,12 @@
 import numpy as np
+import sys
 
 from tiledb.vector_search.module import *
 from tiledb.vector_search.storage_formats import storage_formats
 from typing import Any, Mapping, Optional
 
-MAX_UINT64 = 2 ** 63 - 1
+MAX_UINT64 = np.iinfo(np.dtype("uint64")).max
+MAX_FLOAT_32 = np.finfo(np.dtype("float32")).max
 
 
 class Index:
@@ -37,43 +39,44 @@ def __init__(
 
 
     def query(self, queries: np.ndarray, k, **kwargs):
-        # TODO merge results based on scores and use higher k to improve retrieval
         updated_ids = set(self.read_updated_ids())
-        internal_results = self.query_internal(queries, k, **kwargs)
+        internal_results_d, internal_results_i = self.query_internal(queries, k, **kwargs)
         if self.update_arrays_uri is None:
-            return internal_results
-        addition_results = self.query_additions(queries, k)
-        merged_results = np.zeros((queries.shape[0], k), dtype=np.uint64)
+            return internal_results_d, internal_results_i
+
+        addition_results_d, addition_results_i = self.query_additions(queries, k)
+        # Filter updated vectors
         query_id = 0
-        for query in internal_results:
+        for query in internal_results_i:
             res_id = 0
-            additional_res_id = 0
             for res in query:
                 if res in updated_ids:
-                    merged_results[query_id, res_id] = addition_results[query_id, additional_res_id]
-                    additional_res_id += 1
-                else:
-                    merged_results[query_id, res_id] = res
+                    internal_results_d[query_id, res_id] = MAX_FLOAT_32
+                    internal_results_i[query_id, res_id] = 0
                 res_id += 1
             query_id += 1
-        return merged_results
+        # Merge update results
+        results_d = np.hstack((internal_results_d, addition_results_d))
+        results_i = np.hstack((internal_results_i, addition_results_i))
+        sort_index = np.argsort(results_d, axis=1)
+        results_d = np.take_along_axis(results_d, sort_index, axis=1)
+        results_i = np.take_along_axis(results_i, sort_index, axis=1)
+        return results_d[:, 0:k], results_i[:, 0:k]
 
     def query_internal(self, queries: np.ndarray, k, **kwargs):
         raise NotImplementedError
 
     def query_additions(self, queries: np.ndarray, k):
         assert queries.dtype == np.float32
-
         additions_vectors, additions_external_ids = self.read_additions()
         queries_m = array_to_matrix(np.transpose(queries))
-        r = query_vq_heap_pyarray(
+        d, i = query_vq_heap_pyarray(
             array_to_matrix(np.transpose(additions_vectors).astype(self.dtype)),
             queries_m,
             StdVector_u64(additions_external_ids),
             k,
             8)
-
-        return np.transpose(np.array(r))
+        return np.transpose(np.array(d)), np.transpose(np.array(i))
 
     def update(self, vector: np.array, external_id: np.uint64):
         updates_array = self.open_updates_array()
@@ -129,7 +132,7 @@ def open_updates_array(self):
             if tiledb.array_exists(updates_array_uri):
                 raise RuntimeError(f"Array {updates_array_uri} already exists.")
             external_id_dim = tiledb.Dim(
-                name="external_id", domain=(0, MAX_UINT64), dtype=np.dtype(np.uint64)
+                name="external_id", domain=(0, MAX_UINT64-1), dtype=np.dtype(np.uint64)
             )
             dom = tiledb.Domain(external_id_dim)
             vector_attr = tiledb.Attr(name="vector", dtype=self.dtype, var=True)

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

@@ -130,7 +130,7 @@ def query_internal(
         if mode is None:
             queries_m = array_to_matrix(np.transpose(queries))
             if self.memory_budget == -1:
-                r = ivf_query_ram(
+                d, i = ivf_query_ram(
                     self.dtype,
                     self._db,
                     self._centroids,
@@ -139,13 +139,12 @@ def query_internal(
                     self._ids,
                     nprobe=nprobe,
                     k_nn=k,
-                    nth=True,  # ??
                     nthreads=nthreads,
                     ctx=self.ctx,
                     use_nuv_implementation=use_nuv_implementation,
                 )
             else:
-                r = ivf_query(
+                d, i = ivf_query(
                     self.dtype,
                     self.db_uri,
                     self._centroids,
@@ -155,13 +154,12 @@ def query_internal(
                     nprobe=nprobe,
                     k_nn=k,
                     memory_budget=self.memory_budget,
-                    nth=True,  # ??
                     nthreads=nthreads,
                     ctx=self.ctx,
                     use_nuv_implementation=use_nuv_implementation,
                 )
 
-            return np.transpose(np.array(r))
+            return np.transpose(np.array(d)), np.transpose(np.array(i))
         else:
             return self.taskgraph_query(
                 queries=queries,
@@ -322,7 +320,8 @@ def dist_qv_udf(
             res = node.result()
             results.append(res)
 
-        results_per_query = []
+        results_per_query_d = []
+        results_per_query_i = []
         for q in range(queries.shape[0]):
             tmp_results = []
             for j in range(k):
@@ -333,5 +332,6 @@ def dist_qv_udf(
             tmp = sorted(tmp_results, key=lambda t: t[0])[0:k]
             for j in range(len(tmp), k):
                 tmp.append((float(0.0), int(0)))
-            results_per_query.append(np.array(tmp, dtype=np.dtype("float,int"))["f1"])
-        return results_per_query
+            results_per_query_d.append(np.array(tmp, dtype=np.dtype("float,uint64"))["f0"])
+            results_per_query_i.append(np.array(tmp, dtype=np.dtype("float,uint64"))["f1"])
+        return results_per_query_d, results_per_query_i

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

@@ -5,15 +5,13 @@
 #include <pybind11/stl.h>
 
 #include "linalg.h"
-#include "ivf_index.h"
 #include "ivf_query.h"
 #include "flat_query.h"
 
 namespace py = pybind11;
 using Ctx = tiledb::Context;
 
 bool global_debug = false;
-double global_time_of_interest;
 
 bool enable_stats = false;
 std::vector<json> core_stats;
@@ -103,6 +101,19 @@ static void declare_pyarray_to_matrix(py::module& m, const std::string& suffix)
         });
 }
 
+namespace {
+ template <typename ...TArgs>
+ py::tuple make_python_pair(std::tuple<TArgs...>&& arg) {
+    static_assert(sizeof...(TArgs) == 2, "Must have exactly two arguments");
+
+   return py::make_tuple<py::return_value_policy::automatic>(
+      py::cast(std::get<0>(arg), py::return_value_policy::move),
+      py::cast(std::get<1>(arg), py::return_value_policy::move)
+   );
+ }
+
+}
+
 template <typename T, typename Id_Type = uint64_t>
 static void declare_qv_query_heap_infinite_ram(py::module& m, const std::string& suffix) {
   m.def(("qv_query_heap_infinite_ram_" + suffix).c_str(),
@@ -113,8 +124,7 @@ static void declare_qv_query_heap_infinite_ram(py::module& m, const std::string&
          std::vector<Id_Type>& ids,
          size_t nprobe,
          size_t k_nn,
-         bool nth,
-         size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
+         size_t nthreads) -> py::tuple { //std::pair<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::qv_query_heap_infinite_ram(
             parts,
@@ -124,9 +134,8 @@ static void declare_qv_query_heap_infinite_ram(py::module& m, const std::string&
             ids,
             nprobe,
             k_nn,
-            nth,
             nthreads);
-        return r;
+        return make_python_pair(std::move(r));
         }, py::keep_alive<1,2>());
 }
 
@@ -142,8 +151,7 @@ static void declare_qv_query_heap_finite_ram(py::module& m, const std::string& s
          size_t nprobe,
          size_t k_nn,
          size_t upper_bound,
-         bool nth,
-         size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
+         size_t nthreads) -> py::tuple { //std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::qv_query_heap_finite_ram<T, Id_Type>(
             ctx,
@@ -155,9 +163,8 @@ static void declare_qv_query_heap_finite_ram(py::module& m, const std::string& s
             nprobe,
             k_nn,
             upper_bound,
-            nth,
             nthreads);
-        return r;
+        return make_python_pair(std::move(r));
         }, py::keep_alive<1,2>());
 }
 
@@ -171,8 +178,7 @@ static void declare_nuv_query_heap_infinite_ram(py::module& m, const std::string
          std::vector<Id_Type>& ids,
          size_t nprobe,
          size_t k_nn,
-         bool nth,
-         size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
+         size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::nuv_query_heap_infinite_ram_reg_blocked(
             parts,
@@ -182,7 +188,6 @@ static void declare_nuv_query_heap_infinite_ram(py::module& m, const std::string
             ids,
             nprobe,
             k_nn,
-            nth,
             nthreads);
         return r;
         }, py::keep_alive<1,2>());
@@ -200,8 +205,7 @@ static void declare_nuv_query_heap_finite_ram(py::module& m, const std::string&
          size_t nprobe,
          size_t k_nn,
          size_t upper_bound,
-         bool nth,
-         size_t nthreads) -> ColMajorMatrix<size_t> { // TODO change return type
+         size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::nuv_query_heap_finite_ram_reg_blocked<T, Id_Type>(
             ctx,
@@ -213,7 +217,6 @@ static void declare_nuv_query_heap_finite_ram(py::module& m, const std::string&
             nprobe,
             k_nn,
             upper_bound,
-            nth,
             nthreads);
         return r;
         }, py::keep_alive<1,2>());
@@ -398,7 +401,7 @@ static void declare_vq_query_heap(py::module& m, const std::string& suffix) {
            ColMajorMatrix<float>& query_vectors,
            const std::vector<uint64_t> &ids,
            int k,
-           size_t nthreads) -> ColMajorMatrix<size_t> {
+           size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
           auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
           return r;
         });
@@ -411,7 +414,7 @@ static void declare_vq_query_heap_pyarray(py::module& m, const std::string& suff
            ColMajorMatrix<float>& query_vectors,
            const std::vector<uint64_t> &ids,
            int k,
-           size_t nthreads) -> ColMajorMatrix<size_t> {
+           size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
           auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
           return r;
         });
@@ -494,17 +497,17 @@ PYBIND11_MODULE(_tiledbvspy, m) {
         [](ColMajorMatrix<float>& data,
            ColMajorMatrix<float>& query_vectors,
            int k,
-           size_t nthreads) -> ColMajorMatrix<size_t> {
-          auto r = detail::flat::vq_query_nth(data, query_vectors, k, true, nthreads);
+           size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
+          auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
           return r;
         });
 
   m.def("query_vq_u8",
         [](tdbColMajorMatrix<uint8_t>& data,
            ColMajorMatrix<float>& query_vectors,
            int k,
-           size_t nthreads) -> ColMajorMatrix<size_t> {
-          auto r = detail::flat::vq_query_nth(data, query_vectors, k, true, nthreads);
+           size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
+          auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
           return r;
         });
 

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

@@ -207,7 +207,6 @@ def ivf_query_ram(
     ids: "Vector",
     nprobe: int,
     k_nn: int,
-    nth: bool,
     nthreads: int,
     ctx: "Ctx" = None,
     use_nuv_implementation: bool = False,
@@ -233,8 +232,6 @@ def ivf_query_ram(
         Number of probs
     k_nn: int
         Number of nn
-    nth: bool
-        Return nth records
     nthreads: int
         Number of theads
     ctx: Ctx
@@ -252,7 +249,6 @@ def ivf_query_ram(
             ids,
             nprobe,
             k_nn,
-            nth,
             nthreads,
         ]
     )
@@ -281,7 +277,6 @@ def ivf_query(
     nprobe: int,
     k_nn: int,
     memory_budget: int,
-    nth: bool,
     nthreads: int,
     ctx: "Ctx" = None,
     use_nuv_implementation: bool = False,
@@ -309,8 +304,6 @@ def ivf_query(
         Number of nn
     memory_budget: int
         Main memory budget
-    nth: bool
-        Return nth records
     nthreads: int
         Number of theads
     ctx: Ctx
@@ -330,7 +323,6 @@ def ivf_query(
             nprobe,
             k_nn,
             memory_budget,
-            nth,
             nthreads,
         ]
     )