IVF_PQ re-ranking

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

@@ -89,6 +89,7 @@ def query_internal(
         self,
         queries: np.ndarray,
         k: int = 10,
+        k_factor: float = 1.0,
         nprobe: Optional[int] = 100,
         **kwargs,
     ):
@@ -101,6 +102,13 @@ def query_internal(
             2D array of query vectors. This can be used as a batch query interface by passing multiple queries in one call.
         k: int
             Number of results to return per query vector.
+        k_factor: int
+            To improve accuracy, IVF_PQ can search for more vectors than requested and then
+            perform re-ranking using the original non-PQ-encoded vectors. This can be slightly
+            slower, but is more accurate. k_factor is the factor by which to increase the number
+            of vectors searched. 1 means we search for exactly `k` vectors. 10 means we search for
+            `10*k` vectors.
+            Defaults to 1.
         nprobe: int
             Number of partitions to check per query.
             Use this parameter to trade-off accuracy for latency and cost.
@@ -120,11 +128,11 @@ def query_internal(
 
         if self.memory_budget == -1:
             distances, ids = self.index.query_infinite_ram(
-                queries_feature_vector_array, k, nprobe
+                queries_feature_vector_array, k, nprobe, k_factor
             )
         else:
             distances, ids = self.index.query_finite_ram(
-                queries_feature_vector_array, k, nprobe, self.memory_budget
+                queries_feature_vector_array, k, nprobe, self.memory_budget, k_factor
             )
 
         return np.array(distances, copy=False), np.array(ids, copy=False)

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

@@ -498,29 +498,38 @@ void init_type_erased_module(py::module_& m) {
           [](IndexIVFPQ& index,
              const FeatureVectorArray& vectors,
              size_t top_k,
-             size_t nprobe) {
-            auto r =
-                index.query(QueryType::InfiniteRAM, vectors, top_k, nprobe);
+             size_t nprobe,
+             float k_factor) {
+            auto r = index.query(
+                QueryType::InfiniteRAM, vectors, top_k, nprobe, 0, k_factor);
             return make_python_pair(std::move(r));
           },
           py::arg("vectors"),
           py::arg("top_k"),
-          py::arg("nprobe"))
+          py::arg("nprobe"),
+          py::arg("k_factor") = 1.f)
       .def(
           "query_finite_ram",
           [](IndexIVFPQ& index,
              const FeatureVectorArray& vectors,
              size_t top_k,
              size_t nprobe,
-             size_t memory_budget) {
+             size_t memory_budget,
+             float k_factor) {
             auto r = index.query(
-                QueryType::FiniteRAM, vectors, top_k, nprobe, memory_budget);
+                QueryType::FiniteRAM,
+                vectors,
+                top_k,
+                nprobe,
+                memory_budget,
+                k_factor);
             return make_python_pair(std::move(r));
           },
           py::arg("vectors"),
           py::arg("top_k"),
           py::arg("nprobe"),
-          py::arg("memory_budget"))
+          py::arg("memory_budget"),
+          py::arg("k_factor") = 1.f)
       .def(
           "write_index",
           [](IndexIVFPQ& index,

apis/python/test/test_index.py

@@ -329,6 +329,9 @@ def test_ivf_pq_index(tmp_path):
     assert distances[0][0] == MAX_FLOAT32
     assert ids[0][0] == MAX_UINT64
     query_and_check_distances(index, queries, 1, [[MAX_FLOAT32]], [[MAX_UINT64]])
+    query_and_check_distances(
+        index, queries, 1, [[MAX_FLOAT32]], [[MAX_UINT64]], k_factor=2.0
+    )
     check_default_metadata(uri, vector_type, STORAGE_VERSION, "IVF_PQ")
 
     update_vectors = np.empty([5], dtype=object)
@@ -344,6 +347,14 @@ def test_ivf_pq_index(tmp_path):
     query_and_check_distances(
         index, np.array([[2, 2, 2]], dtype=np.float32), 2, [[0, 3]], [[2, 1]]
     )
+    query_and_check_distances(
+        index,
+        np.array([[2, 2, 2]], dtype=np.float32),
+        2,
+        [[0, 3]],
+        [[2, 1]],
+        k_factor=2.0,
+    )
 
     index = index.consolidate_updates()
 
@@ -366,6 +377,14 @@ def test_ivf_pq_index(tmp_path):
             [[0], [0]],
             [[i], [i]],
         )
+        query_and_check_distances(
+            index,
+            np.array([[i, i, i], [i, i, i]], dtype=np.float32),
+            1,
+            [[0], [0]],
+            [[i], [i]],
+            k_factor=2.0,
+        )
 
     # Test that we can query with k > 1.
     query_and_check_distances(
@@ -380,6 +399,14 @@ def test_ivf_pq_index(tmp_path):
         [[0, 3], [0, 3]],
         [[0, 1], [4, 3]],
     )
+    query_and_check_distances(
+        index,
+        np.array([[0, 0, 0], [4, 4, 4]], dtype=np.float32),
+        2,
+        [[0, 3], [0, 3]],
+        [[0, 1], [4, 3]],
+        k_factor=2.0,
+    )
 
     vfs = tiledb.VFS()
     assert vfs.dir_size(uri) > 0

apis/python/test/test_ingestion.py

@@ -245,9 +245,17 @@ def test_ivf_pq_ingestion_u8(tmp_path):
         queries,
         k=k,
         nprobe=nprobe,
-        use_nuv_implementation=True,
     )
-    assert accuracy(result, gt_i) > MINIMUM_ACCURACY
+    query_accuracy = accuracy(result, gt_i)
+    assert query_accuracy > MINIMUM_ACCURACY
+
+    _, result = index_ram.query(
+        queries,
+        k=k,
+        k_factor=2,
+        nprobe=nprobe,
+    )
+    assert accuracy(result, gt_i) > query_accuracy + 0.1
 
     _, result = index_ram.query(
         queries,

src/include/api/ivf_flat_index.h

@@ -614,7 +614,7 @@ class IndexIVFFlat {
               (float*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] = impl_index_.query_infinite_ram(qspan, k_nn, nprobe);
+          auto [s, t, _] = impl_index_.query_infinite_ram(qspan, k_nn, nprobe);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};
@@ -624,7 +624,7 @@ class IndexIVFFlat {
               (uint8_t*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] = impl_index_.query_infinite_ram(qspan, k_nn, nprobe);
+          auto [s, t, _] = impl_index_.query_infinite_ram(qspan, k_nn, nprobe);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};
@@ -651,7 +651,7 @@ class IndexIVFFlat {
               (float*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] =
+          auto [s, t, _] =
               impl_index_.query_finite_ram(qspan, k_nn, nprobe, upper_bound);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
@@ -662,7 +662,7 @@ class IndexIVFFlat {
               (uint8_t*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] =
+          auto [s, t, _] =
               impl_index_.query_finite_ram(qspan, k_nn, nprobe, upper_bound);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};

src/include/api/ivf_pq_index.h

@@ -248,11 +248,13 @@ class IndexIVFPQ {
       const QueryVectorArray& vectors,
       size_t top_k,
       size_t nprobe,
-      size_t upper_bound = 0) {
+      size_t upper_bound = 0,
+      float k_factor = 1.f) {
     if (!index_) {
       throw std::runtime_error("Cannot query() because there is no index.");
     }
-    return index_->query(queryType, vectors, top_k, nprobe, upper_bound);
+    return index_->query(
+        queryType, vectors, top_k, nprobe, upper_bound, k_factor);
   }
 
   void write_index(
@@ -398,7 +400,8 @@ class IndexIVFPQ {
         const QueryVectorArray& vectors,
         size_t top_k,
         size_t nprobe,
-        size_t upper_bound) = 0;
+        size_t upper_bound,
+        float k_factor) = 0;
 
     virtual void write_index(
         const tiledb::Context& ctx,
@@ -508,7 +511,8 @@ class IndexIVFPQ {
         const QueryVectorArray& vectors,
         size_t top_k,
         size_t nprobe,
-        size_t upper_bound) override {
+        size_t upper_bound,
+        float k_factor) override {
       // @todo using index_type = size_t;
       auto dtype = vectors.feature_type();
 
@@ -519,8 +523,8 @@ class IndexIVFPQ {
               (float*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] =
-              impl_index_.query(queryType, qspan, top_k, nprobe, upper_bound);
+          auto [s, t] = impl_index_.query(
+              queryType, qspan, top_k, nprobe, upper_bound, k_factor);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};
@@ -530,8 +534,8 @@ class IndexIVFPQ {
               (uint8_t*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] =
-              impl_index_.query(queryType, qspan, top_k, nprobe, upper_bound);
+          auto [s, t] = impl_index_.query(
+              queryType, qspan, top_k, nprobe, upper_bound, k_factor);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};

src/include/detail/ivf/dist_qv.h

@@ -124,12 +124,15 @@ auto dist_qv_finite_ram_part(
           fixed_min_pair_heap<score_type, shuffled_ids_type>(k_nn));
 
   size_t part_offset = 0;
+  size_t indices_offset = 0;
   while (partitioned_vectors.load()) {
     _i.start();
     auto current_part_size = ::num_partitions(partitioned_vectors);
     size_t parts_per_thread = (current_part_size + nthreads - 1) / nthreads;
 
-    std::vector<std::future<decltype(min_scores)>> futs;
+    std::vector<std::future<std::vector<
+        fixed_min_triplet_heap<score_type, shuffled_ids_type, size_t>>>>
+        futs;
     futs.reserve(nthreads);
 
     for (size_t n = 0; n < nthreads; ++n) {
@@ -148,7 +151,8 @@ auto dist_qv_finite_ram_part(
              k_nn,
              first_part,
              last_part,
-             part_offset]() {
+             part_offset,
+             indices_offset]() {
               return apply_query(
                   partitioned_vectors,
                   std::optional<std::vector<int>>{},
@@ -159,6 +163,7 @@ auto dist_qv_finite_ram_part(
                   first_part,
                   last_part,
                   part_offset,
+                  indices_offset,
                   distance);
             }));
       }
@@ -167,13 +172,14 @@ auto dist_qv_finite_ram_part(
       auto min_n = futs[n].get();
 
       for (size_t j = 0; j < num_queries; ++j) {
-        for (auto&& [e, f] : min_n[j]) {
+        for (auto&& [e, f, _] : min_n[j]) {
           min_scores[j].insert(e, f);
         }
       }
     }
 
     part_offset += current_part_size;
+    indices_offset += num_vectors(partitioned_vectors);
     _i.stop();
   }
   return min_scores;