Return distances from query functions

This PR updates the IVF queries to return top k distances in addition
to top k neighbor indices.

All of our search functions now use a fixed size min heap of to keep
distances and indices.  The heap is templated on the type of distance
and type of index, which are kept as a tuple in the heap.  The heap is
sorted on the first element (the score).  To create a matrix of top k
neighbors, the search functions sorted the heap and then copied the
neighbor indices into a matrix.

With the new API, queries can now be called as ``` auto&& [ D, I ] =
query(...) ```

To update the query functions to return distances in addition to
indices, this PR
  * Created a new file scoring.h which includes most of the
  functionality previously in defs.h
  * Deleted defs.h
  * Created a function `consolidate` which consolidates distance and
  index information from a vector of vector of indices into a single
  vector (the 0th vector in the vector of vectors)
  * Factored out the code to copy the index information in the heap
  structure to a matrix into new functions `get_top_k_from_heap` and two
  overloads of `get_top_k`
  * Created augmented functions `get_top_k_from_heap_with_scores` and
  two overloads of `get_top_k_with_scores` that return a tuple of a
  distance matrix and an index matrix.
  * Replaced the final logic in all of the ivf queries with
  `consolidate`  + `get_top_k_with_scores` (not all functions needed
  `consolidate`)
  * Added extensive new tests in unit_fixed_min_heap.cc, unit_linalg.cc,
  and unit_scoring.cc, notably to validate the new `get_top_k` family of
  functions.
  * Updated the ivf_flat C++ CLI program to use the new query API
  * Updated the Python bindings in module.cc to use the new query API

In addition
  * Propagated the BLAS macro to exclude BLAS code from the library so
  that the CLIs can be built without BLAS
  * Added docstrings to the query functions in ivf/qv.h
  * Added an initializer list constructor to `Matrix` to aid in creating
  tests

Author: lums658

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

@@ -100,9 +100,9 @@ def query(
         assert targets.dtype == np.float32
 
         targets_m = array_to_matrix(np.transpose(targets))
-        r = query_vq_heap(self._db, targets_m, self._ids, k, nthreads)
+        _, r = query_vq_heap(self._db, targets_m, self._ids, k, nthreads)
 
-        return np.transpose(np.array(r))
+        return np.transpose(np.array(_)), np.transpose(np.array(r))
 
 
 class IVFFlatIndex(Index):
@@ -222,7 +222,7 @@ def query(
         if mode is None:
             queries_m = array_to_matrix(np.transpose(queries))
             if self.memory_budget == -1:
-                r = ivf_query_ram(
+                _, r = ivf_query_ram(
                     self.dtype,
                     self._db,
                     self._centroids,
@@ -231,13 +231,12 @@ def query(
                     self._ids,
                     nprobe=nprobe,
                     k_nn=k,
-                    nth=True,  # ??
                     nthreads=nthreads,
                     ctx=self.ctx,
                     use_nuv_implementation=use_nuv_implementation,
                 )
             else:
-                r = ivf_query(
+                _, r = ivf_query(
                     self.dtype,
                     self.parts_db_uri,
                     self._centroids,
@@ -247,13 +246,12 @@ def query(
                     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(_)), np.transpose(np.array(r))
         else:
             return self.taskgraph_query(
                 queries=queries,

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;
@@ -113,8 +111,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) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
 
         auto r = detail::ivf::qv_query_heap_infinite_ram(
             parts,
@@ -124,7 +121,6 @@ static void declare_qv_query_heap_infinite_ram(py::module& m, const std::string&
             ids,
             nprobe,
             k_nn,
-            nth,
             nthreads);
         return r;
         }, py::keep_alive<1,2>());
@@ -142,8 +138,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) -> 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,7 +150,6 @@ static void declare_qv_query_heap_finite_ram(py::module& m, const std::string& s
             nprobe,
             k_nn,
             upper_bound,
-            nth,
             nthreads);
         return r;
         }, py::keep_alive<1,2>());
@@ -171,8 +165,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 +175,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 +192,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 +204,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>());
@@ -394,7 +384,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;
         });
@@ -477,17 +467,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

@@ -188,7 +188,6 @@ def ivf_query_ram(
     ids: "Vector",
     nprobe: int,
     k_nn: int,
-    nth: bool,
     nthreads: int,
     ctx: "Ctx" = None,
     use_nuv_implementation: bool = False,
@@ -214,8 +213,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
@@ -233,7 +230,6 @@ def ivf_query_ram(
             ids,
             nprobe,
             k_nn,
-            nth,
             nthreads,
         ]
     )
@@ -262,7 +258,6 @@ def ivf_query(
     nprobe: int,
     k_nn: int,
     memory_budget: int,
-    nth: bool,
     nthreads: int,
     ctx: "Ctx" = None,
     use_nuv_implementation: bool = False,
@@ -290,8 +285,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
@@ -311,7 +304,6 @@ def ivf_query(
             nprobe,
             k_nn,
             memory_budget,
-            nth,
             nthreads,
         ]
     )

apis/python/test/test_ingestion.py

@@ -25,7 +25,7 @@ def test_flat_ingestion_u8(tmp_path):
         index_uri=index_uri,
         source_uri=os.path.join(dataset_dir, "data.u8bin"),
     )
-    result = index.query(query_vectors, k=k)
+    _, result = index.query(query_vectors, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
@@ -44,11 +44,11 @@ def test_flat_ingestion_f32(tmp_path):
         index_uri=index_uri,
         source_uri=os.path.join(dataset_dir, "data.f32bin"),
     )
-    result = index.query(query_vectors, k=k)
+    _, result = index.query(query_vectors, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = FlatIndex(uri=index_uri)
-    result = index_ram.query(query_vectors, k=k)
+    _, result = index_ram.query(query_vectors, k=k)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
@@ -71,7 +71,7 @@ def test_flat_ingestion_external_id_u8(tmp_path):
         source_uri=os.path.join(dataset_dir, "data.u8bin"),
         external_ids=external_ids
     )
-    result = index.query(query_vectors, k=k)
+    _, result = index.query(query_vectors, k=k)
     assert accuracy(result, gt_i, external_ids_offset=external_ids_offset) > MINIMUM_ACCURACY
 
 
@@ -96,22 +96,22 @@ def test_ivf_flat_ingestion_u8(tmp_path):
         partitions=partitions,
         input_vectors_per_work_item=int(size / 10),
     )
-    result = index.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=index_uri, memory_budget=int(size / 10))
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
+    _, result = index_ram.query(
         query_vectors,
         k=k,
         nprobe=nprobe,
         use_nuv_implementation=True,
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
+    _, result = index_ram.query(
         query_vectors,
         k=k,
         nprobe=nprobe,
@@ -144,26 +144,26 @@ def test_ivf_flat_ingestion_f32(tmp_path):
         input_vectors_per_work_item=int(size / 10),
     )
 
-    result = index.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=index_uri, memory_budget=int(size / 10))
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=index_uri, memory_budget=int(size / 10))
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
+    _, result = index_ram.query(
         query_vectors,
         k=k,
         nprobe=nprobe,
         use_nuv_implementation=True,
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
@@ -186,26 +186,26 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
         source_uri=source_uri,
         partitions=partitions,
     )
-    result = index.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     # Test single query vector handling
-    result1 = index.query(query_vectors[10], k=k, nprobe=nprobe)
+    _, result1 = index.query(query_vectors[10], k=k, nprobe=nprobe)
     assert accuracy(result1, np.array([gt_i[10]])) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=index_uri)
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
+    _, result = index_ram.query(
         query_vectors,
         k=k,
         nprobe=nprobe,
         use_nuv_implementation=True,
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
 
@@ -230,26 +230,26 @@ def test_ivf_flat_ingestion_numpy(tmp_path):
         input_vectors=input_vectors,
         partitions=partitions,
     )
-    result = index.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
     # Test single query vector handling
-    result1 = index.query(query_vectors[10], k=k, nprobe=nprobe)
+    _, result1 = index.query(query_vectors[10], k=k, nprobe=nprobe)
     assert accuracy(result1, np.array([gt_i[10]])) > MINIMUM_ACCURACY
 
     index_ram = IVFFlatIndex(uri=index_uri)
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(
+    _, result = index_ram.query(
         query_vectors,
         k=k,
         nprobe=nprobe,
         use_nuv_implementation=True,
     )
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 
-    result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
+    _, result = index_ram.query(query_vectors, k=k, nprobe=nprobe, mode=Mode.LOCAL)
     assert accuracy(result, gt_i) > MINIMUM_ACCURACY
 def test_ivf_flat_ingestion_external_ids_numpy(tmp_path):
     source_uri = "test/data/siftsmall/siftsmall_base.fvecs"
@@ -275,5 +275,5 @@ def test_ivf_flat_ingestion_external_ids_numpy(tmp_path):
         partitions=partitions,
         external_ids=external_ids
     )
-    result = index.query(query_vectors, k=k, nprobe=nprobe)
+    _, result = index.query(query_vectors, k=k, nprobe=nprobe)
     assert accuracy(result, gt_i, external_ids_offset) > MINIMUM_ACCURACY