Skip to content

Port asg017's DiskANN index - #11

Open
MayCXC wants to merge 5 commits into
photostructure:mainfrom
MayCXC:diskann-index-photostructure
Open

MayCXC wants to merge 5 commits into
photostructure:mainfrom
MayCXC:diskann-index-photostructure

Conversation

@MayCXC

@MayCXC MayCXC commented Oct 23, 2026 •

Copy link
Copy Markdown

Ports asg017's DiskANN index, the second of its approximate-nearest-neighbour indexes, as asg017's main has it at 04d28bd. This builds on #8: its three commits come first here, and the last two are this PR's.

What it adds. A float or int8 vector column can be declared indexed by diskann(neighbor_quantizer=binary) or indexed by diskann(neighbor_quantizer=int8). The column keeps a graph of its rows, after LM-DiskANN: each row links to up to n_neighbors others (72 by default) and keeps their vectors quantized. A KNN query walks the graph from its medoid, keeping the search_list_size nearest rows it has seen (128 by default), and ranks each row by its exact distance once it reaches it. With buffer_threshold=N, inserted rows wait in a buffer and join the graph N at a time, and a query reads the buffer in full.

create virtual table v using vec0(emb float[768] indexed by diskann(neighbor_quantizer=binary, n_neighbors=32));
select rowid, distance from v where emb match ? and k = 10;
insert into v(v) values ('search_list_size=256');  -- for this connection's queries

The graph is in _diskann_nodesNN, the vectors by rowid in _vectorsNN, the buffer in _diskann_bufferNN and the medoid in _info, declared as upstream declares them, so a database opens with either. A table whose every vector column is indexed by diskann() keeps no chunks. SQLITE_VEC_ENABLE_DISKANN=0 builds without it.

What it ports. asg017's sqlite-vec-diskann.c as of 575371d, with the refusal of a vector UPDATE (9df59b4), the confirmed flag that keeps deleted nodes out of results (b008654), the delete repair's error (c36a995) and scrub (01b4b2a), the NULL blob checks (82f4eb0), auxiliary columns (b7fc459), the rename of its shadow tables (c560710, 8105eee), the finalizing of every cached statement on commit (fe94171), and the text primary key tests (f2c9fb8).

Where it differs from upstream.

  • diskann_knn() takes rescore_knn()'s arguments and results: a rowid filter ranks the rows it names by their exact distances, distance constraints filter the rows the search found, and MMR reranks them. Upstream reads only the query vector and k, while xBestIndex has told SQLite the other constraints are handled, so it returns rows outside them:
    create virtual table t using vec0(emb float[8] indexed by diskann(neighbor_quantizer=binary, n_neighbors=8));
    -- rows 1 to 20, row i holding ((i*7 + j*3) % 11) / 10.0 - 0.5 at index j; ? is row 3's vector
    select rowid from t where emb match ? and k = 3 and rowid in (1, 2);
    -- asg017 main: 3, 14, 17; here: 1, 2
    select rowid, distance from t where emb match ? and k = 3 and distance > 1.0;
    -- asg017 main: 3 at 0.0, 14 at 0.0, 17 at 1.034; here: 17, 11 and 6 at 1.034
  • An int8 column's vectors are quantized as int8 values: under the int8 quantizer a vector is its own quantized form, and under binary each value above 0 sets its bit. Upstream reads an int8 query as floats, 4 times its bytes, so AddressSanitizer reports a heap-buffer-overflow in diskann_quantize_vector() on inserts into an int8[2048] column, and it takes an int8 vector's first bytes as its binary form.
  • Re-ranking replaces a row's quantized distance with its exact one whether it is larger or smaller, so every result's distance is exact; upstream keeps the quantized distance when the exact one is larger, as it often is under the int8 quantizer with cosine.
  • A deleted medoid is replaced by a node of the graph. Upstream picks a row of _vectorsNN, which can be one still in the buffer, with no node, so no search reaches the graph until the buffer is flushed.
  • The visited set grows as a search fills it and holds rowid 0; upstream's has 4 * search_list_size slots, answers "not visited" once they are full, and cannot hold rowid 0.
  • Errors reach the statement, so an interrupt stays an interrupt, through the search, the prune, the reverse edges, a delete's repair and scrub, the buffer and its flush, and the KNN query's read of the buffer; upstream skips the failed step or reports SQLITE_ERROR. A node blob whose size the column's options do not give fails with a message naming the node.
  • A cached statement is reset as soon as it has run, so VACUUM works after a KNN query; upstream's stay in progress until the next commit, and VACUUM fails with cannot VACUUM - SQL statements in progress.
  • diskann_robust_prune() runs diskann_prune_select(), the selection the unit tests and diskann-prune-direct.c exercise, which reads the distance between two candidates through a callback; upstream's insert runs a copy of the selection inline, and its diskann_prune_select() takes a matrix of distances only the tests build.
  • The int8 quantizer clamps to [-128, 127] and takes NaN to 127, as the rescore int8 quantizer does; upstream casts out-of-range floats to int8.
  • CREATE options and their values must be spelled in full and their numbers fit an int; upstream accepts any prefix, as in n_neigh=16 or neighbor_quantizer=b. The commands take digits only, as oversample=N does, and set every DiskANN column of the table; upstream reads them with atoi() and sets only the first DiskANN column.
  • optimize, which upstream lacks, leaves a DiskANN column's rows, which are keyed by rowid, as they are.

As with rescore, a table with a DiskANN column refuses metadata and partition key columns, and a DiskANN column's vector cannot be updated. ARCHITECTURE.md lists the new tables with the other backing tables, the README has a section on the index, and the commit adds an ### Added entry under [Unreleased] in CHANGELOG.md.

Tests.

  • Upstream's test-diskann.py, the DiskANN cases of its test-auxiliary.py, test-rename.py and test-cache-finalize.py, its ten DiskANN fuzz targets with their dictionary, renamed diskann-create.dict for the target the fuzz workflow passes it to, and test-unit.c's DiskANN cases, which call the functions through tests/sqlite-vec-internal.h.
  • New: a search covering the table answers as a flat column does; int8 columns under both quantizers; exact distances; the medoid's replacement from the graph; a point query for a rowid the table does not have; a corrupt node blob; the parsing of options and commands; optimize beside a flat column; VACUUM after a KNN query. test-fullscan.py reads a DiskANN column beside a flat one. Built from asg017's main, 6 of the 12 option cases fail, and so does the VACUUM case.
  • The last commit adds two DiskANN tables to test_interrupt_is_reported_as_interrupt, one holding three rows in its buffer: KNN queries on the graph, on a rowid filter and over the buffer, a full scan, a point query, an insert, an insert that flushes the buffer, an auxiliary column update, deletes of a node, of the medoid and of a buffered row, and optimize, each interrupted at every progress handler call. A build whose node write or node read returns SQLITE_ERROR in place of vec0_error_rc()'s code fails 8 or 12 of the cases.

make test-loadable (680 passed, 3 skipped), make test-unit and make test pass, and so does the Python suite against an AddressSanitizer and UndefinedBehaviorSanitizer build. GCC 14.2 and Clang 19.1 with -Wall -Wextra print nothing for the loadable build with any combination of SQLITE_VEC_ENABLE_RESCORE and SQLITE_VEC_ENABLE_DISKANN, and with -DSQLITE_CORE only main's three unused pApi warnings, which #9 removes. clang-format 21 leaves sqlite-vec.c unchanged, and before this rebase each of the ten DiskANN fuzz targets ran 90 seconds without a finding.

Rebased with #8 onto main at f6a3244, it merges cleanly with main and with #4. On that main, a DELETE learns from vec0Update_Delete_ClearValidity() whether it emptied the chunk; a table whose every vector column is indexed by diskann() keeps no chunks, so it clears no validity bit and deletes no chunk.

@mceachen

Copy link
Copy Markdown
Member

We'll review the index ports one at a time once #4 and #7 are in, starting with rescore (#8).

A vector column can name the index that answers its KNN queries, as in
asg017/sqlite-vec@0de765f. flat() is the exhaustive search a column without
the clause gets. The scanner reads '(', ')' and ',' for the index's options,
and the clause's keywords match exactly even when an old table's definition
is parsed leniently, since no old table has one.
A float vector column declared `indexed by rescore(quantizer=bit)` or
`indexed by rescore(quantizer=int8)`, with an optional `oversample=N` from 1
to 128 (8 by default), also keeps its vectors quantized in _rescore_chunksNN,
in the layout of _vector_chunksNN, and its float32 vectors by rowid in
_rescore_vectorsNN. A KNN query scans the quantized vectors for the
k * oversample nearest rows, then ranks those candidates by the distances of
their float32 vectors. The `oversample=N` command sets the oversample of a
connection's queries. A full scan reads a rescore column's vectors row by row
through vec0_get_vector_data(), and a flat column's by chunk as before.

Ported from asg017#276 (ba0db0b) as asg017's main has it at
04d28bd: with the int8 quantizer that 575371d fixed, the fixes 5e4c557,
82f4eb0 and 5522e86, auxiliary columns (b7fc459), the oversample=N command
and its unknown-command message (6e2c4c6), and the distance constraints of
the unmerged asg017#309. SQLITE_VEC_ENABLE_RESCORE=0 builds without it.

Where this port differs from upstream:

- rescore_knn() takes vec0Filter_knn_chunks_iter()'s arguments and results,
  so vec0Filter_knn() picks between the two and the MMR rerank runs after
  either. It reports errors as the chunk scan does, so an interrupt stays an
  interrupt. It picks each chunk's candidates with this fork's min_idx(),
  whose heap needs no bitmap of taken slots.
- Distance constraints with a NULL, TEXT or BLOB value act as they do for a
  flat column.
- optimize, which upstream lacks, writes a moved row's quantized vectors to
  its new slot and leaves its row of _rescore_vectorsNN, whose rowid does not
  change, and deletes the old _rescore_chunksNN rows with the old chunks. A
  DELETE that empties a chunk deletes its _rescore_chunksNN row too.
- A point query for a rowid the table does not have returns no row, since
  vec0_get_vector_data() finds the row's position first, as for a flat
  column; upstream fails with "Could not fetch vector data".
- _rescore_chunksNN declares "rowid INTEGER PRIMARY KEY", as this fork's new
  tables do, and an oversample=N value must be all digits.

The tests are upstream's test-rescore.py and test-rescore-mutations.py, the
rescore cases of its test-auxiliary.py and test-rename.py, its rescore fuzz
targets, and test-unit.c's bit quantizer, byte size and parser cases.
Upstream's int8 quantizer unit tests, and the uniform-vector case of
rescore-quantize-edge.c, test the min/max quantizer that 575371d replaced, so
they are written here for the current one. Modes 1 and 2 of
rescore-quantize-edge.c copy the fuzz bytes out to float alignment before
reading them as floats, as asg017#343 has them; upstream's read floats one
byte past the mode byte, a misaligned load UBSan reports. New tests check
that each row's quantized slot holds the quantized form of its vector
through inserts, updates, deletes that empty chunks, and optimize, and
rescore-optimize.c fuzzes optimize. test-fullscan.py reads a rescore column
beside a flat one, and runs its delete and optimize cases on a rescore
column too; read by chunk, a rescore column fails all three with "Could not
fetch vector data".
test_interrupt_is_reported_as_interrupt runs a statement on a copy of a
database, interrupted at each progress handler call, and expects
SQLITE_INTERRUPT. A rescore table joins its database, with a KNN query, a
full scan, a point query, an insert, an update, a delete, a delete that
empties a chunk, and optimize, so the rescore code's internal statements are
interrupted too.
A build whose rescore insert returns SQLITE_ERROR in place of
vec0_error_rc()'s code fails the insert case.
A float or int8 vector column declared
`indexed by diskann(neighbor_quantizer=binary)` or
`indexed by diskann(neighbor_quantizer=int8)` keeps a graph of its rows,
after the LM-DiskANN paper: each row links to up to n_neighbors others (72
by default) and keeps their vectors quantized. A KNN query walks the graph from
its medoid, keeping the search_list_size nearest rows it has seen (128 by
default), and ranks each row by its exact distance once it reaches it. With
buffer_threshold=N, inserted rows wait in a buffer and join the graph N at a
time, and a query reads the buffer in full. The graph is in _diskann_nodesNN,
the vectors by rowid in _vectorsNN, the buffer in _diskann_bufferNN and the
medoid in _info, declared as upstream declares them, so a database opens with
either. A table whose every vector column is indexed by diskann() keeps no
chunks. The search_list_size=N, search_list_size_search=N and
search_list_size_insert=N commands set the sizes of a connection's searches.

Ported from asg017's sqlite-vec-diskann.c as its main has it at 04d28bd
(575371d), with the refusal of a vector UPDATE (9df59b4), the confirmed flag
that keeps deleted nodes out of results (b008654), the delete repair's error
(c36a995) and scrub (01b4b2a), the NULL blob checks (82f4eb0), auxiliary
columns (b7fc459), the rename of its shadow tables (c560710, 8105eee), the
finalizing of every cached statement on commit (fe94171), and the text
primary key tests (f2c9fb8). SQLITE_VEC_ENABLE_DISKANN=0 builds without it.

Where this port differs from upstream:

- diskann_knn() takes rescore_knn()'s arguments and results: a rowid filter
  ranks the rows it names by their exact distances, distance constraints
  filter the rows the search found, and MMR reranks them. Upstream reads
  only the query vector and k, so it returns rows outside a rowid filter or
  a distance constraint.
- An int8 column's vectors are quantized as int8 values: under the int8
  quantizer a vector is its own quantized form, and under binary each value
  above 0 sets its bit. Upstream reads an int8 query as floats, 4 times its
  bytes, and takes an int8 vector's first bytes as its binary form.
- Re-ranking replaces a row's quantized distance with its exact one whether
  it is larger or smaller, so every result's distance is exact; upstream
  keeps the quantized distance when the exact one is larger, as it often is
  under the int8 quantizer with cosine.
- A deleted medoid is replaced by a node of the graph. Upstream picks a row
  of _vectorsNN, which can be one still in the buffer, with no node, so no
  search reaches the graph until the buffer is flushed.
- The visited set grows as a search fills it and holds rowid 0; upstream's
  has 4 * search_list_size slots, answers "not visited" once they are full,
  and cannot hold rowid 0.
- Errors reach the statement, so an interrupt stays an interrupt, through
  the search, the prune, the reverse edges, a delete's repair and scrub, the
  buffer and its flush, and the KNN query's read of the buffer; upstream
  skips the failed step or reports SQLITE_ERROR. A node blob whose size the
  column's options do not give fails with a message naming the node.
- A cached statement is reset as soon as it has run, so VACUUM works after a
  KNN query; upstream's stay in progress until the next commit, and VACUUM
  fails with "SQL statements in progress".
- diskann_robust_prune() runs diskann_prune_select(), the selection the unit
  tests and diskann-prune-direct.c exercise, which reads the distance
  between two candidates through a callback; upstream's insert runs a copy
  of the selection inline, and its diskann_prune_select() takes a matrix of
  distances only the tests build.
- The int8 quantizer clamps to [-128, 127] and takes NaN to 127, as the
  rescore int8 quantizer does; upstream casts out-of-range floats to int8.
- CREATE options and their values must be spelled in full and their numbers
  fit an int; upstream accepts any prefix, as in `n_neigh=16` or
  `neighbor_quantizer=b`. The commands take digits only, as oversample=N
  does, and set every DiskANN column of the table; upstream reads them with
  atoi() and sets only the first DiskANN column.
- optimize, which upstream lacks, leaves a DiskANN column's rows, which are
  keyed by rowid, as they are.

The tests are upstream's test-diskann.py, the DiskANN cases of its
test-auxiliary.py, test-rename.py and test-cache-finalize.py, its DiskANN
fuzz targets with their dictionary, renamed diskann-create.dict for the
target the fuzz workflow passes it to, and test-unit.c's DiskANN cases, which
call the functions through tests/sqlite-vec-internal.h. New tests check that
a search covering the table answers as a flat column does, int8 columns under
both quantizers, exact distances, the medoid's replacement from the graph, a
point query for a rowid the table does not have, a corrupt node blob, the
parsing of options and commands, optimize beside a flat column, and VACUUM
after a KNN query. test-fullscan.py reads a DiskANN column beside a flat one.
test_interrupt_is_reported_as_interrupt runs a statement on a copy of a
database, interrupted at each progress handler call, and expects
SQLITE_INTERRUPT. Two DiskANN tables join its database, one holding three
rows in its buffer, with KNN queries on the graph, on a rowid filter and over
the buffer, a full scan, a point query, an insert, an insert that flushes the
buffer, an auxiliary column update, deletes of a node, of the medoid and of a
buffered row, and optimize, so the DiskANN code's internal statements are
interrupted too.
A build whose node write or node read returns SQLITE_ERROR in place of
vec0_error_rc()'s code fails 8 or 12 of the cases.

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.

2 participants