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;