Skip to content

Select each chunk's k nearest rows with a heap - #18

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

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

Conversation

@MayCXC

@MayCXC MayCXC commented Oct 22, 2026 •

Copy link
Copy Markdown

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 spends about a million comparisons per chunk. On an in-memory 10,000-row float[768] table (medians of 5 interleaved runs, each the median of 7 queries):

k main this change
1 9.0 ms 9.3 ms
100 9.9 ms 8.9 ms
1000 17.9 ms 9.6 ms
4096 18.1 ms 9.4 ms

asg017/sqlite-vec 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 than its default build.

This ports that heap as 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 after every number instead of wherever its position put it. A vec0 table refuses NaN and infinite elements and zero vectors under cosine distance, but one written before v0.2.6-alpha, which began refusing NaN, can still hold a NaN, and a static blob reads the application's memory as it is. The bitmap of taken slots the callers allocated for the scan goes with it, from the KNN loop, the static blob filter, 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; mceachen's test from that commit, which patches a NaN into a stored chunk to stand in for an older table, fails on main and on this PR's earlier head.

To check the order is 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. 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, checks that a static blob vector with a NaN element comes back last (on main the scan returns it first, from slot 0), and that a stored NaN in the second chunk comes back behind the first chunk's rows. The suite passes (147 passed, 4 skipped), and test-knn-topk.py passes against an AddressSanitizer and UndefinedBehaviorSanitizer build; before the merge change the C unit tests, the Rust harness (cargo test --manifest-path tests/Cargo.toml) and test-loadable.py under AddressSanitizer passed too, and the change touches none of what they test.

It merges cleanly with #9 through #14 and with #17 and #19. 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 upstream as asg017#337, which makes the build option's heap the default.

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 an in-memory 10,000-row float[768]
table, k = 1000 and k = 4096 took 17.9 and 18.1 ms against 9.0 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 after every number instead of wherever its position put it.
vec0 refuses NaN and infinite elements and zero vectors under cosine
distance, but a table written before v0.2.6-alpha, which began
refusing NaN, can still hold one, and a static blob reads the
application's memory as it is. The bitmap of taken slots the callers allocated for
the scan goes with it, from the KNN loop, the static blob filter, 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.

Measured on that table (medians of 5 interleaved runs of 7 queries),
k = 1000 and k = 4096 take 9.6 and 9.4 ms, k = 100 8.9 ms instead of
9.9, and k = 1 9.3 ms against 9.0. 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, checks that a static blob vector with a NaN element comes
back last (the scan returned it first, from slot 0), and, with
mceachen's test from a791836, that a NaN patched into a stored chunk
comes back behind the earlier chunk's rows, which fails on main.

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