From 64b479c1a7a6df3a3a5652f7a3f0def6baae612f Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 13:54:11 -0500 Subject: [PATCH] Select each chunk's k nearest rows with the min_idx() heap min_idx() had two bodies: by default a scan of the whole chunk once per result, k times n comparisons, and under SQLITE_VEC_EXPERIMENTAL_MIN_IDX a binary heap of the k kept so far, n log k. 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 22.3 and 22.8 ms against 8.1 ms for k = 1. The heap compares distances alone, so it returns tied rows in another order than the scan: over 1,000 KNN queries on tables built to tie heavily, a build with the option returns other rows or another order for 779 of them. This keeps the heap as min_idx()'s only body and orders it with min_idx_after(): by distance, between equal distances the larger index first, as the scan's <= chose, and a NaN distance after every number. The cosine distance to a zero vector is 0 / 0, a NaN, and the scan returned such a row first whenever it was the first candidate left in its chunk: with a zero vector in slot 0, a k = 2 query for [1, 0] returned it ahead of an exact match and dropped the next nearest row. The bitmap of taken slots the scan needed goes with it, from vec0's KNN loop, the rescore KNN, 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, so a zero vector in the second chunk came back ahead of the first chunk's rows. Measured on that table (medians of 5 interleaved runs of 7 queries), k = 1000 and k = 4096 take 8.6 and 8.9 ms, k = 100 8.2 ms instead of 8.6, and k = 1 8.1 ms either way; the heap as the option builds it, without the order, takes 8.5 ms at both. 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 default build row for row. 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, and checks that a zero vector in a cosine column comes back last, within its chunk and behind an earlier chunk's rows; on main both of those tests fail, with the zero vector first. --- sqlite-vec-rescore.c | 6 +-- sqlite-vec.c | 83 ++++++++++++------------------- tests/sqlite-vec-internal.h | 1 - tests/test-knn-topk.py | 98 +++++++++++++++++++++++++++++++++++++ tests/unittest.rs | 3 -- 5 files changed, 132 insertions(+), 59 deletions(-) create mode 100644 tests/test-knn-topk.py diff --git a/sqlite-vec-rescore.c b/sqlite-vec-rescore.c index 6a472144..7de1e17d 100644 --- a/sqlite-vec-rescore.c +++ b/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; diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..b359a392 100644 --- a/sqlite-vec.c +++ b/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); diff --git a/tests/sqlite-vec-internal.h b/tests/sqlite-vec-internal.h index 313add4d..af8257bd 100644 --- a/tests/sqlite-vec-internal.h +++ b/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 ); diff --git a/tests/test-knn-topk.py b/tests/test-knn-topk.py new file mode 100644 index 00000000..60b1430a --- /dev/null +++ b/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] diff --git a/tests/unittest.rs b/tests/unittest.rs index 7d098560..213abf1f 100644 --- a/tests/unittest.rs +++ b/tests/unittest.rs @@ -10,7 +10,6 @@ fn _min_idx(distances: Vec, 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;