From 22791f5cb22a2edfeac6b46cde2f02b182113b0d Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 13:41:29 -0500 Subject: [PATCH] Select each chunk's k nearest rows with a heap min_idx() chose the k nearest candidates of a chunk by scanning the whole chunk once per result, k times n comparisons. With the default chunk size of 1024, a KNN query with k of 1024 or more spent about a million comparisons per chunk: on an in-memory 10,000-row float[768] table, k = 1000 and k = 4096 took 17.9 and 18.1 ms against 9.0 ms for k = 1. Upstream has the fix behind a build option: asg017/sqlite-vec 0de765f added SQLITE_VEC_EXPERIMENTAL_MIN_IDX, a min_idx() that keeps the k best candidates in a binary heap in the output array, whose root is the one that ranks last, so a chunk costs n log k comparisons, then sorts the heap in place. That heap compares distances alone, so tied rows come back in another order than the scan's: built with the option, 779 of 1,000 tie-heavy queries on asg017's main return other rows or another order. This makes that heap the only min_idx() and gives it one total order, min_idx_after(): by distance, between equal distances the larger index first, as the scan's <= chose, and a NaN distance after every number instead of wherever its position put it. vec0 refuses NaN and infinite elements and zero vectors under cosine distance, but a table written before v0.2.6-alpha, which began refusing NaN, can still hold one, and a static blob reads the application's memory as it is. The bitmap of taken slots the callers allocated for the scan goes with it, from the KNN loop, the static blob filter, and the declarations in tests/sqlite-vec-internal.h and tests/unittest.rs. merge_sorted_lists(), which merges each chunk's nearest rows into the ones kept from the chunks before it, ranks a NaN after every number too, as mceachen's photostructure/sqlite-vec@a791836 does in that fork: it took a later chunk's NaN ahead of the numbers still unmerged. Measured on that table (medians of 5 interleaved runs of 7 queries), k = 1000 and k = 4096 take 9.6 and 9.4 ms, k = 100 8.9 ms instead of 9.9, and k = 1 9.3 ms against 9.0. Over 4,000 random KNN queries on tables built to tie heavily (vectors on a 5 by 3 grid, chunk sizes 8 to 64, deletes, metadata filters, k from 1 to 1000), every result matches the previous selection row for row; with the tie rule reversed, 827 of 1,000 such queries differ. tests/test-knn-topk.py checks the k nearest against a brute-force order for k below, at and above the chunk size, with and without a filter, pins the order of tied rows, checks that a static blob vector with a NaN element comes back last (the scan returned it first, from slot 0), and, with mceachen's test from a791836, that a NaN patched into a stored chunk comes back behind the earlier chunk's rows, which fails on main. --- sqlite-vec.c | 133 ++++++++++++++++++++++++------------ tests/sqlite-vec-internal.h | 1 - tests/test-knn-topk.py | 116 +++++++++++++++++++++++++++++++ tests/unittest.rs | 3 - 4 files changed, 205 insertions(+), 48 deletions(-) create mode 100644 tests/test-knn-topk.py diff --git a/sqlite-vec.c b/sqlite-vec.c index ab77083c..2b2fcecb 100644 --- a/sqlite-vec.c +++ b/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 { diff --git a/tests/sqlite-vec-internal.h b/tests/sqlite-vec-internal.h index 6d4a783a..a237e278 100644 --- a/tests/sqlite-vec-internal.h +++ b/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 ); diff --git a/tests/test-knn-topk.py b/tests/test-knn-topk.py new file mode 100644 index 00000000..52f09247 --- /dev/null +++ b/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(?, ', k: i32) -> Vec { let mut out: Vec = vec![0; k as usize]; let bitmap_bytes = n / 8; let mut candidates: Vec = vec![0xFF; bitmap_bytes]; - let mut b_taken: Vec = vec![0; bitmap_bytes]; let mut k_used: i32 = 0; unsafe { @@ -20,7 +19,6 @@ fn _min_idx(distances: Vec, k: i32) -> Vec { 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;