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;