From 696cff596f99ba21d0c6fe302c1714182b4e1319 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 12:36:07 -0500 Subject: [PATCH] perf(sqlite-vec.c): 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 a 10,000-row float[768] table, k = 1000 and k = 4096 took 17.8 and 18.0 ms against 8.7 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, which a vec0 table cannot produce because it refuses NaN and infinite elements but a static blob can, after every number instead of wherever its position put it. The bitmap of taken slots the callers allocated for the scan goes with it, from the KNN loop, the static blob filter and the declaration in tests/sqlite-vec-internal.h. Measured on that table (medians of 5 interleaved runs of 7 queries), k = 1000 and k = 4096 take 9.3 ms; k = 100 takes 8.8 ms instead of 9.2, and k = 1 8.7 ms either way. 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, and checks that a static blob vector with a NaN element comes back last; the scan returned it first, from slot 0. --- CHANGELOG.md | 1 + sqlite-vec.c | 130 ++++++++++++++++++++++++------------ tests/sqlite-vec-internal.h | 1 - tests/test-knn-topk.py | 82 +++++++++++++++++++++++ 4 files changed, 170 insertions(+), 44 deletions(-) create mode 100644 tests/test-knn-topk.py diff --git a/CHANGELOG.md b/CHANGELOG.md index da5f90d1..6465d57f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Changed - Reading a vector column without a KNN query, as in `select embedding from v`, reads each chunk's vectors through one blob handle instead of opening the chunk's blob again for every row: selecting every vector of a 10,000-row `float[768]` table took 25 ms instead of 2,080 ms, and `select vec_distance_cosine(embedding, ?) from v` 41 ms instead of 2,093 ms. +- A KNN query picks each chunk's k nearest rows with the heap of upstream's `SQLITE_VEC_EXPERIMENTAL_MIN_IDX` build option ([asg017/sqlite-vec@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f4570c1d65a0699e8e38d6c5562da9e23e)) instead of one pass over the chunk for every row it keeps, ordered so that it returns the same rows in the same order as before: on a 10,000-row `float[768]` table, `k = 1000` and `k = 4096` take 9.3 ms instead of 17.8 and 18.0 ms, `k = 100` 8.8 ms instead of 9.2, and `k = 1` 8.7 ms either way. A NaN distance, which only a static blob can produce since `vec0` rejects NaN elements, ranks after every number instead of where its position put it. ### Fixed diff --git a/sqlite-vec.c b/sqlite-vec.c index 429b475d..49ba8776 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -6668,44 +6668,105 @@ int bitmap_any(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; } @@ -8013,7 +8074,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 @@ -8069,12 +8129,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; @@ -8371,7 +8425,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, @@ -8404,7 +8458,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); @@ -12277,7 +12330,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, i32 *topk_rowids = NULL; f32 *distances = NULL; u8 *candidates = NULL; - u8 *taken = NULL; knn_data = sqlite3_malloc(sizeof(*knn_data)); if (!knn_data) { @@ -12341,18 +12393,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; @@ -12363,7 +12409,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, // Cleanup temporaries (not owned by knn_data) queryVectorCleanup(queryVector); sqlite3_free(candidates); - sqlite3_free(taken); return SQLITE_OK; knn_cleanup: @@ -12372,7 +12417,6 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, sqlite3_free(topk_rowids); sqlite3_free(distances); sqlite3_free(candidates); - sqlite3_free(taken); return rc; } else { pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN; diff --git a/tests/sqlite-vec-internal.h b/tests/sqlite-vec-internal.h index 28b72267..20a788c5 100644 --- a/tests/sqlite-vec-internal.h +++ b/tests/sqlite-vec-internal.h @@ -12,7 +12,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..25f9b839 --- /dev/null +++ b/tests/test-knn-topk.py @@ -0,0 +1,82 @@ +import random +import struct + +import numpy as np + +from conftest import get_extension_path + + +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(get_extension_path(), 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(?, '