Select each chunk's k nearest rows with a heap

sqlite-vec.c

@@ -6218,7 +6218,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++;
@@ -6285,44 +6286,105 @@ 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.
+ *
+ * out holds a heap of the k candidates kept so far, whose root is the one that
+ * ranks last, so n candidates cost n log k comparisons; the heap is then
+ * sorted in place. This is upstream's SQLITE_VEC_EXPERIMENTAL_MIN_IDX variant
+ * (https://github.com/asg017/sqlite-vec/blob/04d28bd21773981e2d266bbf6aa4efbd011eb4f6/sqlite-vec.c#L6516)
+ * with min_idx_after() as its order, so that 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);
 
-  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;
-      }
+  int heap_size = 0;
+
+#define HEAP_SIFT_UP(pos)                                                      \
+  do {                                                                         \
+    int _c = (pos);                                                            \
+    while (_c > 0) {                                                           \
+      int _p = (_c - 1) / 2;                                                   \
+      if (min_idx_after(distances, out[_c], out[_p])) {                        \
+        i32 _tmp = out[_p];                                                    \
+        out[_p] = out[_c];                                                     \
+        out[_c] = _tmp;                                                        \
+        _c = _p;                                                               \
+      } else                                                                   \
+        break;                                                                 \
+    }                                                                          \
+  } while (0)
+
+#define HEAP_SIFT_DOWN(pos, sz)                                                \
+  do {                                                                         \
+    int _p = (pos);                                                            \
+    for (;;) {                                                                 \
+      int _l = 2 * _p + 1, _r = 2 * _p + 2, _largest = _p;                     \
+      if (_l < (sz) && min_idx_after(distances, out[_l], out[_largest]))       \
+        _largest = _l;                                                         \
+      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;                                                    \
+      _p = _largest;                                                           \
+    }                                                                          \
+  } while (0)
+
+  for (int i = 0; i < n; i++) {
+    if (!bitmap_get(candidates, i))
+      continue;
+    if (heap_size < k) {
+      out[heap_size] = i;
+      heap_size++;
+      HEAP_SIFT_UP(heap_size - 1);
+    } else if (min_idx_after(distances, out[0], i)) {
+      out[0] = i;
+      HEAP_SIFT_DOWN(0, heap_size);
     }
+  }
 
-    out[ik] = min_idx;
-    bitmap_set(bTaken, min_idx, 1);
+  // Heapsort to produce ascending order.
+  for (int i = heap_size - 1; i > 0; i--) {
+    i32 tmp = out[0];
+    out[0] = out[i];
+    out[i] = tmp;
+    HEAP_SIFT_DOWN(0, i);
   }
-  *k_used = k;
+
+#undef HEAP_SIFT_UP
+#undef HEAP_SIFT_DOWN
+
+  *k_used = heap_size;
   return SQLITE_OK;
 }
 
@@ -7307,7 +7369,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
@@ -7363,12 +7424,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;
@@ -7643,7 +7698,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,
@@ -7676,7 +7731,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);
@@ -11282,7 +11336,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
     i32 *topk_rowids = NULL;
     f32 *distances = NULL;
     u8 *candidates = NULL;
-    u8 *taken = NULL;
     int rc = SQLITE_OK;
 
     knn_data = sqlite3_malloc(sizeof(*knn_data));
@@ -11351,18 +11404,12 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
       goto knn_cleanup;
     }
 
-    taken = bitmap_new(bsize);
-    if (!taken) {
-      rc = SQLITE_NOMEM;
-      goto knn_cleanup;
-    }
-
     bitmap_fill(candidates, bsize);
     for (size_t i = bsize; i >= p->blob->nvectors; i--) {
       bitmap_set(candidates, i, 0);
     }
     i32 k_used = 0;
-    min_idx(distances, bsize, candidates, topk_rowids, k, taken, &k_used);
+    min_idx(distances, bsize, candidates, topk_rowids, k, &k_used);
     knn_data->current_idx = 0;
     knn_data->distances = distances;
     knn_data->k = k;
@@ -11372,7 +11419,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
 
     // Cleanup temporary allocations (not owned by knn_data)
     sqlite3_free(candidates);
-    sqlite3_free(taken);
     cleanup(queryVector);
     return SQLITE_OK;
 
@@ -11381,7 +11427,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
     sqlite3_free(topk_rowids);
     sqlite3_free(distances);
     sqlite3_free(candidates);
-    sqlite3_free(taken);
     cleanup(queryVector);
     return rc;
   } else {

tests/sqlite-vec-internal.h

@@ -11,7 +11,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,116 @@
+import random
+import struct
+
+import numpy as np
+
+
+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):
+    # vec0 refuses NaN elements, but a static blob reads its vectors from the
+    # application's memory as they are, so a distance to one can be NaN.
+    db.enable_load_extension(True)
+    db.load_extension("dist/vec0", entrypoint="sqlite3_vec_static_blobs_init")
+    vectors = np.array([[np.nan, 0], [1, 0], [2, 0], [3, 0]], dtype=np.float32)
+    db.execute(
+        "insert into temp.vec_static_blobs(name, data)"
+        " select 'b', vec_static_blob_from_raw(?, '<f4', 2, 4)",
+        [vectors.__array_interface__["data"][0]],
+    )
+    db.execute("create virtual table b using vec_static_blob_entries(b)")
+    rows = db.execute(
+        "select rowid from b where vector match ? and k = 4 order by distance",
+        [_f32([0, 0])],
+    )
+    assert [r[0] for r in rows] == [1, 2, 3, 0]
+
+
+def test_knn_ranks_a_stored_nan_distance_last_across_chunks(db):
+    # vec0 has rejected NaN elements since v0.2.6-alpha, but a table written by
+    # an earlier release can still hold one. Patching a NaN into the stored
+    # chunk stands in for such a table.
+    db.execute("create virtual table v using vec0(e float[1], chunk_size=8)")
+    for i in range(1, 9):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([10 + i])])
+    for i in range(9, 17):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([1])])
+    chunk_id, offset = db.execute(
+        "select chunk_id, chunk_offset from v_rowids where rowid = 16"
+    ).fetchone()
+    vectors = bytearray(
+        db.execute(
+            "select vectors from v_vector_chunks00 where rowid = ?", [chunk_id]
+        ).fetchone()[0]
+    )
+    vectors[offset * 4 : offset * 4 + 4] = _f32([float("nan")])
+    db.execute(
+        "update v_vector_chunks00 set vectors = ? where rowid = ?",
+        [bytes(vectors), chunk_id],
+    )
+
+    def knn(k):
+        return [
+            r[0]
+            for r in db.execute(
+                "select rowid from v where e match ? and k = ?", [_f32([0]), k]
+            )
+        ]
+
+    # The NaN, in the second chunk, ranks after the first chunk's rows too.
+    assert knn(16) == [15, 14, 13, 12, 11, 10, 9, 1, 2, 3, 4, 5, 6, 7, 8, 16]
+    assert knn(9) == [15, 14, 13, 12, 11, 10, 9, 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;