Skip to content

Port asg017's experimental IVF index - #12

Open
MayCXC wants to merge 7 commits into
photostructure:mainfrom
MayCXC:ivf-index-photostructure
Open

MayCXC wants to merge 7 commits into
photostructure:mainfrom
MayCXC:ivf-index-photostructure

Conversation

@MayCXC

@MayCXC MayCXC commented Oct 23, 2026 •

Copy link
Copy Markdown

Ports asg017's experimental IVF index, the third of its approximate-nearest-neighbour indexes, as asg017's main has it at 04d28bd. This builds on #11 (DiskANN), which builds on #8: their five commits come first here, and the last two are this PR's.

What it adds. A build compiled with -DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1 can declare a float vector column indexed by ivf(), which clusters its rows around nlist centroids (128 by default) with k-means and keeps them in cells of up to 64 rows, each cell holding rows of one centroid. A KNN query reads the cells of the nprobe centroids nearest its vector (10 by default) and the cells of the rows no clustering has placed yet, so until compute-centroids runs it reads every row. With quantizer=int8 or quantizer=binary, the cells and centroids hold the vectors quantized, the float vectors are kept by rowid besides, and a query ranks the k * oversample rows nearest by quantized distance by their exact distances.

create virtual table v using vec0(emb float[768] indexed by ivf(nlist=16, nprobe=4));
-- after the inserts
insert into v(v) values ('compute-centroids');
select rowid, distance from v where emb match ? and k = 10;
insert into v(v) values ('nprobe=8');  -- for this connection's queries

The commands compute-centroids, compute-centroids:{"nlist": N, "max_iterations": N, "seed": N}, set-centroid:N, assign-vectors, clear-centroids and nprobe=N build and tune the index. Its tables are _ivf_centroidsNN, _ivf_cellsNN, _ivf_rowid_mapNN and _ivf_vectorsNN, with the ivf_trained_N row of _info, named as upstream names them. The prebuilt binaries are built without it.

What it ports. asg017's sqlite-vec-ivf.c and sqlite-vec-ivf-kmeans.c as its main has them at 04d28bd (merged in e2c38f3), with the refusal of a binary quantizer on dimensions not divisible by 8 (4bee883), the refusal of a vector UPDATE (9df59b4), auxiliary columns (b7fc459), the command column (6e2c4c6), the rename of its shadow tables (8b81f40), the finalizing of every cached statement on commit (fe94171), and the fixes of the unmerged asg017#313: an insert takes the first free slot of a cell, where upstream takes the slot after its count and overwrites a live row once a delete leaves a hole; assign-vectors reads the centroids at the right stride and by their ids, which set-centroid can leave apart; clear-centroids and set-centroid keep a quantized column's cells and centroids quantized; and the int8 quantizer ranks by the column's metric, where upstream ranks by L2.

Where it differs from upstream.

  • ivf_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 query read, and MMR reranks them. Upstream reads only the query vector and k, so it returns rows outside a rowid filter or a distance constraint.
  • A quantized column's point query and full scan read its float vectors from _ivf_vectorsNN. Upstream returns the quantized bytes, and its point query reads a float vector's length past them.
  • Whether a column has centroids is read from _info at each use. Upstream caches it per connection, so a connection that queried before another one ran compute-centroids reads only the cells of no centroid, which are gone, and finds no rows. Two connections to one file, 100 rows of float[8] in ivf(nlist=4, nprobe=4): a KNN on the first returns rows 7, 97 and 40; the second runs compute-centroids; the same KNN on the first then returns nothing on asg017's main, and 7, 97 and 40 here.
  • The cells are found by centroid through the index of a UNIQUE (centroid_id, cell_id) constraint on a cell_id INTEGER PRIMARY KEY, which ALTER TABLE RENAME renames along with its table. Upstream creates an index named after the table, which keeps its name through a rename, so a table created later under the old name fails to create its own; nothing inside xRename can drop an index, since OP_Destroy refuses while the ALTER TABLE statement reads. Every statement on the cells names its columns and addresses a cell by rowid, so either implementation opens the other's tables.
  • Only float columns can be indexed by ivf(). Upstream accepts int8 and bit columns and sizes their cells for float vectors, so it reads past an int8 vector and ranks a row's own vector at a distance such as 6.6e10.
  • Errors reach the statement, so an interrupt stays an interrupt. Upstream ignores the results of its writes to the cells and of its blob reads and writes, and runs compute-centroids in a SAVEPOINT it opens inside xUpdate, which SQLite refuses while the statement writes. An interrupted command fails with interrupted and the statement's rollback undoes it; upstream reports SQL logic error, and the connection's KNN queries then find no rows. A cell or centroid blob whose size the column's options do not give fails with a message naming it.
  • A cached statement is reset as soon as it has run, so VACUUM works after a KNN query.
  • The commands apply to every IVF column of the table; upstream applies them to the first. compute-centroids:{...} is read with SQLite's json_each(), and its keys nlist, max_iterations and seed take integers in range; upstream finds them with strstr() and reads them with atoi(). The commands take digits only, and CREATE options and their values must be spelled in full; upstream accepts any prefix, as in quant=int8.
  • k-means indexes its arrays with size_t; upstream multiplies int indexes, which overflow past 2^31 floats.
  • optimize, which upstream lacks, leaves an IVF column's rows, which are keyed by rowid, as they are.

The IVF code shares the DiskANN code's helpers, which take names for both: diskann_stmt() and diskann_exec_rowid() become vec0_cached_stmt() and vec0_exec_i64(), and vec0_parse_diskann_int() becomes vec0_parse_option_int(). The commit adds an ### Added entry under [Unreleased] in CHANGELOG.md, and ARCHITECTURE.md and the README describe the index.

Tests.

  • Fix ivf index bugs asg017/sqlite-vec#313's test-ivf.py, test-ivf-mutations.py, test-ivf-quantization.py and test-ivf-rename.py; upstream's eight IVF fuzz targets with their dictionary, built with the IVF index compiled in (upstream's fuzz Makefile builds them without it, so they create no IVF table); and the IVF cases of Fix ivf index bugs asg017/sqlite-vec#313's test-unit.c, which call the functions through tests/sqlite-vec-internal.h. Fix ivf index bugs asg017/sqlite-vec#313's test-ann-embeddings.py needs an embedding server on localhost and is not ported.
  • New: a query reading every cell answers as a flat column does, constraint for constraint; the float vectors of a quantized column; a second connection after compute-centroids; a rename followed by a table under the old name; the refused columns and options; an interrupted command; the commands on every IVF column; the command errors; the compute-centroids options and seed; corrupt cells; INSERT OR REPLACE; VACUUM after a KNN query; optimize beside a flat column; and k-means in the unit tests.
  • conftest.py skips the test-ivf files when vec_debug() names no IVF build flag, as upstream's does; make test-unit passes CPPFLAGS to the compiler, and a new CI job, build-linux-x86_64-ivf, compiles the IVF index in and runs every suite.
  • The last commit adds three IVF tables to test_interrupt_is_reported_as_interrupt: KNN queries over the centroids, on a rowid filter and on a quantized column, full scans, a point query, inserts into cells of a centroid, of the quantized column and of no centroid, an auxiliary column update, deletes, optimize, compute-centroids with and without options, clear-centroids, set-centroid and assign-vectors, each interrupted at every progress handler call. A build whose writes to the cells return SQLITE_ERROR in place of vec0_error_rc()'s code fails 12 of the 36 cases.

Without the IVF index, make test-loadable passes (680 passed, 182 skipped); with CPPFLAGS=-DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1 it passes 859 with 3 skipped, and so does the suite against an AddressSanitizer and UndefinedBehaviorSanitizer build. make test-unit and make test pass both ways. GCC 14.2 and Clang 19.1 with -Wall -Wextra print nothing for the loadable build with any combination of SQLITE_VEC_ENABLE_RESCORE, SQLITE_VEC_ENABLE_DISKANN and the IVF flag, 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 eight IVF fuzz targets ran 90 seconds without a finding.

Rebased with #8 and #11 onto main at f6a3244, it merges cleanly with main and with #4. build-linux-x86_64-ivf follows #6's build-linux-aarch64-extension job in .github/workflows/test.yaml.

@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.
A build compiled with -DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1 can declare a
float vector column `indexed by ivf()`, which clusters its rows around nlist
centroids (128 by default) with k-means and keeps them in cells of up to 64
rows, each cell holding rows of one centroid. A KNN query reads the cells of
the nprobe centroids nearest its vector (10 by default) and the cells of the
rows no clustering has placed yet, so until compute-centroids runs it reads
every row. With quantizer=int8 or quantizer=binary, the cells and centroids
hold the vectors quantized, the float32 vectors are kept by rowid besides, and
a query ranks the k * oversample rows nearest by quantized distance by their
exact distances. The commands compute-centroids, compute-centroids:{...},
set-centroid:N, assign-vectors, clear-centroids and nprobe=N build and tune
the index. Its tables are _ivf_centroidsNN, _ivf_cellsNN, _ivf_rowid_mapNN
and _ivf_vectorsNN, with the ivf_trained_N row of _info, named as upstream
names them. The prebuilt binaries are built without it.

Ported from asg017's sqlite-vec-ivf.c and sqlite-vec-ivf-kmeans.c as its main
has them at 04d28bd (merged in e2c38f3), with the refusal of a binary
quantizer on dimensions not divisible by 8 (4bee883), the refusal of a vector
UPDATE (9df59b4), auxiliary columns (b7fc459), the command column (6e2c4c6),
the rename of its shadow tables (8b81f40), the finalizing of every cached
statement on commit (fe94171), and the fixes of the unmerged asg017#313: an
insert takes the first free slot of a cell, where upstream takes the slot
after its count and overwrites a live row once a delete leaves a hole;
assign-vectors reads the centroids at the right stride and by their ids, which
set-centroid can leave apart; clear-centroids and set-centroid keep a
quantized column's cells and centroids quantized; and the int8 quantizer ranks
by the column's metric, where upstream ranks by L2.

Where this port differs from upstream:

- ivf_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 query read, and MMR reranks them. Upstream reads only
  the query vector and k, so it returns rows outside a rowid filter or a
  distance constraint.
- A quantized column's point query and full scan read its float32 vectors
  from _ivf_vectorsNN. Upstream returns the quantized bytes, and its point
  query reads a float32 vector's length past them.
- Whether a column has centroids is read from _info at each use. Upstream
  caches it per connection, so a connection that queried before another one
  ran compute-centroids reads only the cells of no centroid, which are gone,
  and finds no rows.
- The cells are found by centroid through the index of a
  UNIQUE (centroid_id, cell_id) constraint on a cell_id INTEGER PRIMARY KEY,
  which ALTER TABLE RENAME renames along with its table. Upstream creates an
  index named after the table, which keeps its name through a rename, so a
  table created later under the old name fails to create its own; nothing
  inside xRename can drop an index, since OP_Destroy refuses while the ALTER
  TABLE statement reads. Every statement on the cells names its columns and
  addresses a cell by rowid, so either implementation opens the other's
  tables.
- Only float columns can be indexed by ivf(). Upstream accepts int8 and bit
  columns and sizes their cells for float32 vectors, so it reads past an
  int8 vector and ranks a row's own vector at a distance such as 6.6e10.
- Errors reach the statement, so an interrupt stays an interrupt. Upstream
  ignores the results of its writes to the cells and of its blob reads and
  writes, and runs compute-centroids in a SAVEPOINT it opens inside xUpdate,
  which SQLite refuses while the statement writes. An interrupted command
  fails with "interrupted" and the statement's rollback undoes it; upstream
  reports "SQL logic error", and the connection's KNN queries then find no
  rows. A cell or centroid blob whose size the column's options do not give
  fails with a message naming it.
- A cached statement is reset as soon as it has run, so VACUUM works after a
  KNN query.
- The commands apply to every IVF column of the table; upstream applies them
  to the first. compute-centroids:{...} is read with SQLite's json_each(),
  and its keys nlist, max_iterations and seed take integers in range; upstream
  finds them with strstr() and reads them with atoi(). The commands take
  digits only, and CREATE options and their values must be spelled in full;
  upstream accepts any prefix, as in `quant=int8`.
- k-means indexes its arrays with size_t; upstream multiplies int indexes,
  which overflow past 2^31 floats.
- optimize, which upstream lacks, leaves an IVF column's rows, which are
  keyed by rowid, as they are.

The IVF code shares the DiskANN code's helpers, which take names for both:
diskann_stmt() and diskann_exec_rowid() become vec0_cached_stmt() and
vec0_exec_i64(), and vec0_parse_diskann_int() vec0_parse_option_int().

The tests are asg017#313's test-ivf.py, test-ivf-mutations.py,
test-ivf-quantization.py and test-ivf-rename.py, upstream's IVF fuzz targets
with their dictionary, built with the IVF index compiled in (upstream's fuzz
Makefile builds them without it, so they create no IVF table), and the IVF
cases of asg017#313's test-unit.c, which call the functions through
tests/sqlite-vec-internal.h. asg017#313's test-ann-embeddings.py needs an embedding
server on localhost and is not ported. New tests check that a query reading
every cell answers as a flat column does, constraint for constraint, the
float32 vectors of a quantized column, a second connection after
compute-centroids, a rename followed by a table under the old name, the
refused columns and options, an interrupted command, the commands on every
IVF column, the command errors, the compute-centroids options and seed,
corrupt cells, INSERT OR REPLACE, VACUUM after a KNN query, optimize beside a
flat column, and k-means in the unit tests. conftest.py skips the test-ivf
files when vec_debug() names no ivf build flag, as upstream's does; make
test-unit passes CPPFLAGS to the compiler, and a CI job builds the IVF index
in and runs every suite.
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. When the build compiles the IVF index in, three IVF tables
join its database: vi and viq with centroids and a free slot, viq quantized,
and viu with centroids set by hand and its rows in cells of no centroid. The
cases are KNN queries over the centroids, on a rowid filter and on the
quantized column, full scans, a point query, inserts into cells of a
centroid, of the quantized column and of no centroid, an auxiliary column
update, deletes, optimize, compute-centroids with and without options,
clear-centroids, set-centroid and assign-vectors, so the IVF code's internal
statements are interrupted too; a build without the IVF index skips them.
A build whose writes to the cells return SQLITE_ERROR in place of
vec0_error_rc()'s code fails 12 of the 36 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