Select each chunk's k nearest rows with the min_idx() heap

sqlite-vec-rescore.c

@@ -393,12 +393,11 @@ static int rescore_knn(vec0_vtab *p, vec0_cursor *pCur,
   f32 *chunk_distances = sqlite3_malloc(p->chunk_size * sizeof(f32));
   i32 *chunk_topk_idxs = sqlite3_malloc(k_oversample * sizeof(i32));
   u8 *b = sqlite3_malloc(p->chunk_size / CHAR_BIT);
-  u8 *bTaken = sqlite3_malloc(p->chunk_size / CHAR_BIT);
   u8 *bmRowids = NULL;
   void *baseVectors = sqlite3_malloc((i64)p->chunk_size * (i64)qsize);
 
   if (!cand_rowids || !cand_distances || !tmp_rowids || !tmp_distances ||
-      !chunk_distances || !chunk_topk_idxs || !b || !bTaken || !baseVectors) {
+      !chunk_distances || !chunk_topk_idxs || !b || !baseVectors) {
     rc = SQLITE_NOMEM;
     goto cleanup;
   }
@@ -503,7 +502,7 @@ static int rescore_knn(vec0_vtab *p, vec0_cursor *pCur,
 
     int used1;
     min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs,
-            min(k_oversample, p->chunk_size), bTaken, &used1);
+            min(k_oversample, p->chunk_size), &used1);
 
     i64 merged_used;
     merge_sorted_lists(cand_distances, cand_rowids, cand_used, chunk_distances,
@@ -636,7 +635,6 @@ static int rescore_knn(vec0_vtab *p, vec0_cursor *pCur,
   sqlite3_free(chunk_distances);
   sqlite3_free(chunk_topk_idxs);
   sqlite3_free(b);
-  sqlite3_free(bTaken);
   sqlite3_free(bmRowids);
   sqlite3_free(baseVectors);
   return rc;

sqlite-vec.c

@@ -6432,7 +6432,8 @@ void merge_sorted_lists(f32 *a, i64 *a_rowids, i64 a_length, f32 *b,
       out_rowids[i] = a_rowids[ptrA];
       ptrA++;
     } else {
-      if (a[ptrA] <= b[b_top_idxs[ptrB]]) {
+      // A NaN ranks after every number, as in min_idx()
+      if (a[ptrA] <= b[b_top_idxs[ptrB]] || isnan(b[b_top_idxs[ptrB]])) {
         out[i] = a[ptrA];
         out_rowids[i] = a_rowids[ptrA];
         ptrA++;
@@ -6499,32 +6500,49 @@ void bitmap_fill(u8 *bitmap, i32 n) {
 }
 
 /**
- * @brief Finds the minimum k items in distances, and writes the indicies to
- * out.
+ * @brief Whether the candidate at index a ranks after the one at index b in
+ * min_idx()'s order: by distance, a NaN distance after every number, and
+ * between equal distances the smaller index after the larger.
+ */
+static int min_idx_after(const f32 *distances, i32 a, i32 b) {
+  int a_nan = isnan(distances[a]), b_nan = isnan(distances[b]);
+  if (a_nan != b_nan) {
+    return a_nan;
+  }
+  if (!a_nan && distances[a] != distances[b]) {
+    return distances[a] > distances[b];
+  }
+  return a < b;
+}
+
+/**
+ * @brief Finds the k candidates with the smallest distances, and writes their
+ * indices to out, nearest first. Equal distances keep the larger index first,
+ * and a NaN distance ranks after every number.
  *
  * @param distances input f32 array of size n, the items to consider.
  * @param n: size of distances array.
+ * @param candidates: bitmap of size n, the indices that may be chosen.
  * @param out: Output array of size k, will contain at most k element indicies
  * @param k: Size of output array
+ * @param k_used: set to the number of indices written to out
  * @return int
  */
 int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
-            u8 *bTaken, i32 *k_used) {
+            i32 *k_used) {
   assert(k > 0);
   assert(k <= n);
 
-#ifdef SQLITE_VEC_EXPERIMENTAL_MIN_IDX
-  // Max-heap variant: O(n log k) single-pass.
-  // out[0..heap_size-1] stores indices; heap ordered by distances descending
-  // so out[0] is always the index of the LARGEST distance in the top-k.
-  (void)bTaken;
+  // Max-heap: O(n log k) in a single pass.
+  // out[0..heap_size-1] stores indices; heap ordered by min_idx_after()
+  // so out[0] is always the index that ranks LAST in the top-k.
   int heap_size = 0;
 
   #define HEAP_SIFT_UP(pos) do {                          \
     int _c = (pos);                                       \
     while (_c > 0) {                                      \
       int _p = (_c - 1) / 2;                              \
-      if (distances[out[_p]] < distances[out[_c]]) {      \
+      if (min_idx_after(distances, out[_c], out[_p])) {   \
         i32 _tmp = out[_p]; out[_p] = out[_c]; out[_c] = _tmp; \
         _c = _p;                                          \
       } else break;                                       \
@@ -6535,9 +6553,9 @@ int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
     int _p = (pos);                                       \
     for (;;) {                                            \
       int _l = 2*_p + 1, _r = 2*_p + 2, _largest = _p;  \
-      if (_l < (sz) && distances[out[_l]] > distances[out[_largest]]) \
+      if (_l < (sz) && min_idx_after(distances, out[_l], out[_largest])) \
         _largest = _l;                                    \
-      if (_r < (sz) && distances[out[_r]] > distances[out[_largest]]) \
+      if (_r < (sz) && min_idx_after(distances, out[_r], out[_largest])) \
         _largest = _r;                                    \
       if (_largest == _p) break;                          \
       i32 _tmp = out[_p]; out[_p] = out[_largest]; out[_largest] = _tmp; \
@@ -6552,7 +6570,7 @@ int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
       out[heap_size] = i;
       heap_size++;
       HEAP_SIFT_UP(heap_size - 1);
-    } else if (distances[i] < distances[out[0]]) {
+    } else if (min_idx_after(distances, out[0], i)) {
       out[0] = i;
       HEAP_SIFT_DOWN(0, heap_size);
     }
@@ -6569,35 +6587,6 @@ int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
 
   *k_used = heap_size;
   return SQLITE_OK;
-
-#else
-  // Original: O(n*k) repeated linear scan with bitmap.
-  bitmap_clear(bTaken, n);
-
-  for (int ik = 0; ik < k; ik++) {
-    int min_idx = 0;
-    while (min_idx < n &&
-           (bitmap_get(bTaken, min_idx) || !bitmap_get(candidates, min_idx))) {
-      min_idx++;
-    }
-    if (min_idx >= n) {
-      *k_used = ik;
-      return SQLITE_OK;
-    }
-
-    for (int i = 0; i < n; i++) {
-      if (distances[i] <= distances[min_idx] && !bitmap_get(bTaken, i) &&
-          (bitmap_get(candidates, i))) {
-        min_idx = i;
-      }
-    }
-
-    out[ik] = min_idx;
-    bitmap_set(bTaken, min_idx, 1);
-  }
-  *k_used = k;
-  return SQLITE_OK;
-#endif
 }
 
 int vec0_get_metadata_text_long_value(
@@ -7260,7 +7249,6 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
   f32 *tmp_topk_distances = NULL; // memory: k * 4
   f32 *chunk_distances = NULL;    // memory: chunk_size * 4
   u8 *b = NULL;                   // memory: chunk_size / 8
-  u8 *bTaken = NULL;              // memory: chunk_size / 8
   i32 *chunk_topk_idxs = NULL;    // memory: k * 4
   u8 *bmRowids = NULL;            // memory: chunk_size / 8
   u8 *bmMetadata = NULL;            // memory: chunk_size / 8
@@ -7316,12 +7304,6 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
     goto cleanup;
   }
 
-  bTaken = bitmap_new(p->chunk_size);
-  if (!bTaken) {
-    rc = SQLITE_NOMEM;
-    goto cleanup;
-  }
-
   chunk_topk_idxs = sqlite3_malloc(k * sizeof(i32));
   if (!chunk_topk_idxs) {
     rc = SQLITE_NOMEM;
@@ -7593,7 +7575,7 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
 
     int used1;
     min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs,
-            min(k, p->chunk_size), bTaken, &used1);
+            min(k, p->chunk_size), &used1);
 
     i64 used;
     merge_sorted_lists(topk_distances, topk_rowids, k_used, chunk_distances,
@@ -7626,7 +7608,6 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
   sqlite3_free(tmp_topk_rowids);
   sqlite3_free(tmp_topk_distances);
   sqlite3_free(b);
-  sqlite3_free(bTaken);
   sqlite3_free(bmRowids);
   sqlite3_free(baseVectors);
   sqlite3_free(chunk_distances);

tests/sqlite-vec-internal.h

@@ -15,7 +15,6 @@ int min_idx(
   uint8_t *candidates,
   int32_t *out,
   int32_t k,
-  uint8_t *bTaken,
   int32_t *k_used
 );
 

tests/test-knn-topk.py

@@ -0,0 +1,98 @@
+import random
+import struct
+
+
+def _f32(values):
+    return struct.pack(f"{len(values)}f", *values)
+
+
+def test_knn_returns_the_k_nearest_in_order(db):
+    # Distinct distances, so the order is fully determined; k below, at and
+    # above the chunk size, with and without a filter that thins the chunks.
+    rng = random.Random(5)
+    db.execute(
+        "create virtual table v using vec0(e float[3], g integer, chunk_size=16)"
+    )
+    rows = {}
+    for i in range(1, 201):
+        e = [rng.uniform(-1, 1) for _ in range(3)]
+        rows[i] = (e, i % 3)
+        db.execute("insert into v(rowid, e, g) values (?, ?, ?)", [i, _f32(e), i % 3])
+    for _ in range(10):
+        q = [rng.uniform(-1, 1) for _ in range(3)]
+        for k in (1, 2, 7, 16, 17, 64, 200, 300):
+            for g in (None, 1):
+                kept = [i for i, (_, gi) in rows.items() if g is None or gi != g]
+                got = [
+                    r[0]
+                    for r in db.execute(
+                        "select rowid from v where e match ? and k = ?"
+                        + ("" if g is None else " and g != ?"),
+                        [_f32(q), k] + ([] if g is None else [g]),
+                    )
+                ]
+                expected = sorted(
+                    kept,
+                    key=lambda i: db.execute(
+                        "select vec_distance_l2(?, ?)", [_f32(q), _f32(rows[i][0])]
+                    ).fetchone()[0],
+                )[:k]
+                assert got == expected
+
+
+def test_knn_ties_keep_their_order(db):
+    # Equal distances come back chunk by chunk, each chunk's rows from its last
+    # slot to its first, whether or not k cuts the ties short.
+    db.execute("create virtual table v using vec0(e float[2], chunk_size=8)")
+    for i in range(1, 21):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([1, 1])])
+    q = _f32([1, 1])
+
+    def knn(k):
+        return [
+            r[0]
+            for r in db.execute("select rowid from v where e match ? and k = ?", [q, k])
+        ]
+
+    chunks = [range(8, 0, -1), range(16, 8, -1), range(20, 16, -1)]
+    assert knn(20) == [rowid for chunk in chunks for rowid in chunk]
+    assert knn(5) == [8, 7, 6, 5, 4]
+
+
+def test_knn_ranks_a_nan_distance_last(db):
+    # The cosine distance to a zero vector divides zero by zero, so it is NaN.
+    db.execute(
+        "create virtual table v using vec0("
+        "e float[2] distance_metric=cosine, chunk_size=8)"
+    )
+    vectors = {1: [0, 0], 2: [1, 0], 3: [0.9, 0.1], 4: [0, 1]}
+    for rowid, e in vectors.items():
+        db.execute("insert into v(rowid, e) values (?, ?)", [rowid, _f32(e)])
+    rows = db.execute("select rowid from v where e match ? and k = 4", [_f32([1, 0])])
+    assert [r[0] for r in rows] == [2, 3, 4, 1]
+
+
+def test_knn_ranks_a_nan_distance_last_across_chunks(db):
+    # The zero vector's NaN distance is in the second chunk, whose numbers
+    # are all nearer than the first chunk's, so the merge meets the NaN while
+    # the first chunk's rows are still unmerged.
+    db.execute(
+        "create virtual table v using vec0("
+        "e float[2] distance_metric=cosine, chunk_size=8)"
+    )
+    for i in range(1, 9):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([1, 2 + i / 10])])
+    for i in range(9, 16):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([1, (i - 8) / 100])])
+    db.execute("insert into v(rowid, e) values (16, ?)", [_f32([0, 0])])
+
+    def knn(k):
+        return [
+            r[0]
+            for r in db.execute(
+                "select rowid from v where e match ? and k = ?", [_f32([1, 0]), k]
+            )
+        ]
+
+    assert knn(16) == [*range(9, 16), *range(1, 9), 16]
+    assert knn(9) == [*range(9, 16), 1, 2]

tests/unittest.rs

@@ -10,7 +10,6 @@ fn _min_idx(distances: Vec<f32>, k: i32) -> Vec<i32> {
     let mut out: Vec<i32> = vec![0; k as usize];
     let bitmap_bytes = n / 8;
     let mut candidates: Vec<u8> = vec![0xFF; bitmap_bytes];
-    let mut b_taken: Vec<u8> = vec![0; bitmap_bytes];
     let mut k_used: i32 = 0;
 
     unsafe {
@@ -20,7 +19,6 @@ fn _min_idx(distances: Vec<f32>, k: i32) -> Vec<i32> {
             candidates.as_mut_ptr(),
             out.as_mut_ptr(),
             k,
-            b_taken.as_mut_ptr(),
             &mut k_used,
         );
     }
@@ -68,7 +66,6 @@ extern "C" {
         candidates: *mut u8,
         out: *mut i32,
         k: i32,
-        b_taken: *mut u8,
         k_used: *mut i32,
     ) -> i32;