Skip to content

Select each chunk's k nearest rows with the min_idx() heap - #337

Open
MayCXC wants to merge 1 commit into
asg017:mainfrom
MayCXC:knn-topk-heap
Open

MayCXC wants to merge 1 commit into
asg017:mainfrom
MayCXC:knn-topk-heap

Conversation

@MayCXC

@MayCXC MayCXC commented Oct 22, 2026 •

Copy link
Copy Markdown

min_idx() has 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, the scan spends about a million comparisons per chunk on a KNN query with k of 1024 or more. On an in-memory 10,000-row float[768] table (medians of 5 interleaved runs, each the median of 7 queries):

k main main with the option this change
1 8.1 ms 8.1 ms 8.1 ms
100 8.6 ms 8.2 ms 8.2 ms
1000 22.3 ms 8.5 ms 8.6 ms
4096 22.8 ms 8.5 ms 8.9 ms

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:

create virtual table v using vec0(e float[2] distance_metric=cosine, chunk_size=8);
insert into v(rowid, e) values (1, '[0,0]'), (2, '[1,0]'), (3, '[0.9,0.1]'), (4, '[0,1]');
select rowid, distance from v where e match '[1,0]' and k = 2;
-- main:                (1, NULL), (2, 0.0)
-- main with the option: (2, 0.0), (1, NULL)
-- this change:          (2, 0.0), (3, 0.0061)

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 photostructure/sqlite-vec@a791836 by mceachen does in that fork. It took a later chunk's NaN ahead of the numbers still unmerged, so with the zero vector in the second chunk:

create virtual table v using vec0(e float[2] distance_metric=cosine, chunk_size=8);
-- rows 1 to 8 in the first chunk: [1, 2 + i/10]; rows 9 to 15 in the second: [1, (i - 8)/100]; row 16: [0, 0]
select rowid from v where e match '[1,0]' and k = 9;
-- main and this PR before a791836: 9 to 15, then 16, then 1
-- this change:                     9 to 15, then 1, 2

To check the order is otherwise unchanged, I ran 4,000 random KNN queries against main and this branch 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), each build in its own process: every result matched 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, and the second fails on this PR's earlier head too. The Python suite passes (243 passed, 66 skipped), and test-knn-topk.py passes against an AddressSanitizer and UndefinedBehaviorSanitizer build; before the merge change it passed 242 with 66 skipped, 304 with SQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1, and the Rust harness (cargo test --manifest-path tests/Cargo.toml), which the change does not touch. make test-unit does not build on main (tests/sqlite-vec-internal.h redeclares the rescore types), and this change adds no error to it.

It merges cleanly with #267, #327 through #331 and #336. The same change, with the merge fix as mceachen's a791836, is on photostructure/sqlite-vec's main (from photostructure#2), and it is open on vlasky/sqlite-vec as vlasky#18, where the heap is ported from this option.

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant