Port asg017's experimental IVF index

.github/workflows/test.yaml

@@ -67,6 +67,25 @@ jobs:
           path: dist/*
       # after the upload, so dist/test-unit stays out of the artifact
       - run: make test-unit
+  # The IVF index is compiled only on request, so the other builds skip its
+  # tests; this one compiles it in and runs every suite.
+  build-linux-x86_64-ivf:
+    needs: memory-gate
+    runs-on: ubuntu-latest
+    env:
+      CPPFLAGS: -DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1
+    steps:
+      - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+        with:
+          persist-credentials: false
+      - uses: astral-sh/setup-uv@681c641aba71e4a1c380be3ab5e12ad51f415867 # v7.1.6
+        with:
+          enable-cache: true
+      - run: ./scripts/vendor.sh
+      - run: make loadable static
+      - run: uv sync --locked --managed-python --directory tests
+      - run: make test-loadable python=./tests/.venv/bin/python
+      - run: make test-unit
   build-macos-x86_64-extension:
     runs-on: macos-15-intel
     steps:

ARCHITECTURE.md

@@ -28,6 +28,99 @@ Very much a WIP.
 
 #### `xyz_vector_chunksNN`
 
+- `rowid INTEGER`
+- `vectors BLOB`
+
+#### `xyz_rescore_chunksNN`
+
+The vectors of a column indexed by `rescore()`, quantized, in the layout of
+`xyz_vector_chunksNN`.
+
+- `rowid INTEGER`
+- `vectors BLOB`
+
+#### `xyz_rescore_vectorsNN`
+
+The `float32` vectors of a column indexed by `rescore()`, one row per row of
+the table.
+
+- `rowid INTEGER`
+- `vector BLOB`
+
+#### `xyz_vectorsNN`
+
+The vectors of a column indexed by `diskann()`, one row per row of the table.
+
+- `rowid INTEGER`
+- `vector BLOB`
+
+#### `xyz_diskann_nodesNN`
+
+The graph of a column indexed by `diskann()`, one row per node. A node has
+`n_neighbors` slots: `neighbors_validity` holds a bit per slot, set when the
+slot holds a neighbor, `neighbor_ids` the neighbor's rowid as a 64-bit integer
+in the machine's byte order, and `neighbor_quantized_vectors` its vector
+quantized by the column's `neighbor_quantizer`, a bit or a byte per dimension.
+The `diskann_medoid_NN` row of `xyz_info` names the node every search starts
+at, and is `NULL` while the graph has no node.
+
+- `rowid INTEGER`
+- `neighbors_validity BLOB`
+- `neighbor_ids BLOB`
+- `neighbor_quantized_vectors BLOB`
+
+#### `xyz_diskann_bufferNN`
+
+The vectors of a column indexed by `diskann(..., buffer_threshold=N)` that have
+not joined the graph yet.
+
+- `rowid INTEGER`
+- `vector BLOB`
+
+A table whose every vector column is indexed by `diskann()` keeps no chunks:
+its `xyz_chunks` stays empty, and its `xyz_rowids` rows have no `chunk_id` or
+`chunk_offset`.
+
+#### `xyz_ivf_centroidsNN`
+
+The centroids of a column indexed by `ivf()`, held as its cells hold vectors:
+`float32`, or quantized by the column's `quantizer`. The `ivf_trained_N` row of
+`xyz_info` is 1 while the column has centroids.
+
+- `centroid_id INTEGER`
+- `centroid BLOB`
+
+#### `xyz_ivf_cellsNN`
+
+The rows of a column indexed by `ivf()`, in cells of up to 64. A cell holds
+rows of one centroid, or, with `centroid_id` -1, rows no clustering has placed.
+`validity` holds a bit per slot, set when the slot holds a row, `rowids` the
+row's rowid as a 64-bit integer in the machine's byte order, and `vectors` its
+vector as the centroids hold them. `n_vectors` counts the set bits. A query
+finds the cells of a centroid through the index of the
+`UNIQUE (centroid_id, cell_id)` constraint.
+
+- `cell_id INTEGER`
+- `centroid_id INTEGER`
+- `n_vectors INTEGER`
+- `validity BLOB`
+- `rowids BLOB`
+- `vectors BLOB`
+
+#### `xyz_ivf_rowid_mapNN`
+
+The cell and slot of each row of a column indexed by `ivf()`.
+
+- `rowid INTEGER`
+- `cell_id INTEGER`
+- `slot INTEGER`
+
+#### `xyz_ivf_vectorsNN`
+
+The `float32` vectors of a column indexed by `ivf()` with a `quantizer`, one
+row per row of the table. Without a quantizer, the cells hold the vectors and
+the table does not exist.
+
 - `rowid INTEGER`
 - `vector BLOB`
 
@@ -52,7 +145,15 @@ The supported way to tell which tables belong to a `vec0` table is by name. A
 `vec0` table `xyz` owns these tables and no others (`NN` is two digits):
 
 - `xyz_info`, `xyz_chunks`, `xyz_rowids`: always
-- `xyz_vector_chunksNN`: one per vector column
+- `xyz_vector_chunksNN`: one per vector column, except those indexed by
+  `rescore()`, `diskann()` or `ivf()`
+- `xyz_rescore_chunksNN`, `xyz_rescore_vectorsNN`: one each per vector column
+  indexed by `rescore()`
+- `xyz_vectorsNN`, `xyz_diskann_nodesNN`, `xyz_diskann_bufferNN`: one each per
+  vector column indexed by `diskann()`
+- `xyz_ivf_centroidsNN`, `xyz_ivf_cellsNN`, `xyz_ivf_rowid_mapNN`: one each per
+  vector column indexed by `ivf()`, and `xyz_ivf_vectorsNN` for each of those
+  with a `quantizer`
 - `xyz_auxiliary`: only if there are auxiliary columns
 - `xyz_metadatachunksNN`: one per metadata column
 - `xyz_metadatatextNN`: one per text metadata column
@@ -69,6 +170,15 @@ join sqlite_master as s
  and substr(s.name, 1, length(v.name) + 1) = v.name || '_'
  and (substr(s.name, length(v.name) + 2) in ('rowids', 'chunks', 'info', 'auxiliary')
    or substr(s.name, length(v.name) + 2) glob 'vector_chunks[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'rescore_chunks[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'rescore_vectors[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'vectors[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'diskann_nodes[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'diskann_buffer[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'ivf_centroids[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'ivf_cells[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'ivf_rowid_map[0-9][0-9]'
+   or substr(s.name, length(v.name) + 2) glob 'ivf_vectors[0-9][0-9]'
    or substr(s.name, length(v.name) + 2) glob 'metadatachunks[0-9][0-9]'
    or substr(s.name, length(v.name) + 2) glob 'metadatatext[0-9][0-9]')
 where v.type = 'table' and v.rootpage = 0
@@ -84,24 +194,26 @@ relied on:
 - Up to SQLite 3.53, so does a connection that parsed the schema before loading
   the extension.
 - Before SQLite 3.37 there is no `PRAGMA table_list` and no prefix marking (see
-  below), so the vector chunk tables are never shadow tables.
+  below), so the tables whose suffix has an underscore (`vector_chunksNN`,
+  `rescore_chunksNN`, `rescore_vectorsNN`, `diskann_nodesNN`,
+  `diskann_bufferNN`, and the four `ivf_` tables) are never shadow tables.
 - SQLite finds the virtual table that owns a newly created or newly parsed
   table by splitting its name at the **last** underscore
   ([documented](https://www.sqlite.org/vtab.html#the_xshadowname_method)). For
   `xyz_vector_chunks00` it looks for a virtual table called `xyz_vector`, so
   `xShadowName` is never asked about it on that path.
-- `xShadowName` does report `vector_chunksNN`, and SQLite also marks shadow
+- `xShadowName` does report those suffixes, and SQLite also marks shadow
   tables by name prefix when it parses the virtual table's own schema entry.
   That only sees tables parsed earlier, which in practice means after a
   `VACUUM` (it rewrites `sqlite_master` with virtual tables last). A `vec0`
   table created after that `VACUUM` is not covered until the next one.
-- SQLite 3.54 additionally marks by prefix after every `xConnect`. There the
-  vector chunk tables become `shadow` once the `vec0` table has been used in
-  the connection, except in the connection that created it. `PRAGMA
-  table_list` connects every virtual table first, so it always reports them as
-  `shadow`, but under `SQLITE_DBCONFIG_DEFENSIVE` a direct write to a chunk
-  table is still accepted until the `vec0` table is first touched. This prefix
-  matching is not part of the documented contract.
+- SQLite 3.54 additionally marks by prefix after every `xConnect`. There those
+  tables become `shadow` once the `vec0` table has been used in the
+  connection, except in the connection that created it. `PRAGMA table_list`
+  connects every virtual table first, so it always reports them as `shadow`,
+  but under `SQLITE_DBCONFIG_DEFENSIVE` a direct write to one of them is still
+  accepted until the `vec0` table is first touched. This prefix matching is not
+  part of the documented contract.
 
 Rules that follow from this:
 

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## [Unreleased]
 
+### Added
+
+- A vector column definition accepts `indexed by flat()`, as in `embedding float[768] indexed by flat()`. `flat()` is the exhaustive search that a column without the clause gets. Ported from [asg017@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f) ([asg017#273](https://github.com/asg017/sqlite-vec/pull/273)).
+- A `float` vector column can be declared `indexed by rescore(quantizer=bit)` or `indexed by rescore(quantizer=int8)`, with an optional `oversample=N` from 1 to 128 (8 by default). It also keeps its vectors quantized, in the new `_rescore_chunksNN` shadow table, and its `float` vectors by rowid in `_rescore_vectorsNN`. A KNN query scans the quantized vectors for the `k * oversample` nearest rows, then ranks those by the exact distances of their `float` vectors, to which distance constraints apply. The `oversample=N` command, as in `insert into t(t) values ('oversample=16')`, sets the oversample of a connection's queries. Rowid filters, MMR, auxiliary columns, `UPDATE`, `optimize`, and `ALTER TABLE ... RENAME` work on rescore columns; a table with one cannot have metadata or partition key columns. Ported from [asg017#276](https://github.com/asg017/sqlite-vec/pull/276) ([asg017@ba0db0b](https://github.com/asg017/sqlite-vec/commit/ba0db0b)) with its fixes [asg017@5e4c557](https://github.com/asg017/sqlite-vec/commit/5e4c557), [@82f4eb0](https://github.com/asg017/sqlite-vec/commit/82f4eb0), [@5522e86](https://github.com/asg017/sqlite-vec/commit/5522e86), [@b7fc459](https://github.com/asg017/sqlite-vec/commit/b7fc459), [@6e2c4c6](https://github.com/asg017/sqlite-vec/commit/6e2c4c6), the int8 quantizer of [@575371d](https://github.com/asg017/sqlite-vec/commit/575371d), and the distance constraints of the unmerged [asg017#309](https://github.com/asg017/sqlite-vec/pull/309). Unlike upstream, a point query for a rowid that a rescore table does not have returns no row instead of failing with `Could not fetch vector data`.
+- A `float` or `int8` vector column can be declared `indexed by diskann(neighbor_quantizer=binary)` or `indexed by diskann(neighbor_quantizer=int8)`, with optional `n_neighbors=N` (72 by default), `search_list_size=N` (128) or `search_list_size_search=N` and `search_list_size_insert=N`, and `buffer_threshold=N`. The column keeps a graph of its rows in the new `_diskann_nodesNN` shadow table, its vectors by rowid in `_vectorsNN`, and its buffered inserts in `_diskann_bufferNN`, and a KNN query walks the graph instead of reading every vector. The commands `search_list_size=N`, `search_list_size_search=N` and `search_list_size_insert=N` set those sizes for a connection. Rowid filters, distance constraints, MMR, auxiliary columns, `optimize`, and `ALTER TABLE ... RENAME` work on DiskANN columns; a table with one cannot have metadata or partition key columns, and a DiskANN column's vector cannot be updated. A table whose every vector column is indexed by DiskANN keeps no chunks. Ported from [asg017@575371d](https://github.com/asg017/sqlite-vec/commit/575371d) with [asg017@9df59b4](https://github.com/asg017/sqlite-vec/commit/9df59b4), [@82f4eb0](https://github.com/asg017/sqlite-vec/commit/82f4eb0), [@b008654](https://github.com/asg017/sqlite-vec/commit/b008654), [@c36a995](https://github.com/asg017/sqlite-vec/commit/c36a995), [@01b4b2a](https://github.com/asg017/sqlite-vec/commit/01b4b2a), [@b7fc459](https://github.com/asg017/sqlite-vec/commit/b7fc459) and [@fe94171](https://github.com/asg017/sqlite-vec/commit/fe94171). Unlike upstream, an `int8` column's vectors are quantized as `int8` values (upstream reads an `int8` query as floats, past its end), a result's distance is always its exact one, the row every search starts from is replaced, when deleted, by a row of the graph rather than one still waiting in the buffer, rowid filters, distance constraints and MMR apply, errors and interrupts reach the query instead of being skipped, `VACUUM` works after a KNN query, the commands set every DiskANN column of the table, and an abbreviated option such as `n_neigh=16` or `neighbor_quantizer=bin`, a command value such as `search_list_size=12abc`, and a number past an `int` are refused.
+- A build compiled with `-DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1` can declare a `float` vector column `indexed by ivf()`, with optional `nlist=N` (128 by default, or 0 to leave the centroids to `set-centroid`), `nprobe=N` (10, at most `nlist`), `quantizer=int8` or `quantizer=binary`, and `oversample=N` (1). The column keeps its rows in cells of up to 64 in the new `_ivf_cellsNN` shadow table, around the centroids in `_ivf_centroidsNN`, with each row's cell and slot in `_ivf_rowid_mapNN` and, with a quantizer, its `float` vector in `_ivf_vectorsNN`; a KNN query reads the cells of the `nprobe` centroids nearest its vector and the cells of rows no clustering has placed. 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 it. Rowid filters, distance constraints, MMR, auxiliary columns, `optimize`, and `ALTER TABLE ... RENAME` work on IVF columns; a table with one cannot have metadata or partition key columns, and an IVF column's vector cannot be updated. The prebuilt binaries are built without it. Ported from [asg017@e2c38f3](https://github.com/asg017/sqlite-vec/commit/e2c38f3) with [asg017@4bee883](https://github.com/asg017/sqlite-vec/commit/4bee883), [@9df59b4](https://github.com/asg017/sqlite-vec/commit/9df59b4), [@b7fc459](https://github.com/asg017/sqlite-vec/commit/b7fc459), [@6e2c4c6](https://github.com/asg017/sqlite-vec/commit/6e2c4c6), [@8b81f40](https://github.com/asg017/sqlite-vec/commit/8b81f40) and [@fe94171](https://github.com/asg017/sqlite-vec/commit/fe94171), and the fixes of the unmerged [asg017#313](https://github.com/asg017/sqlite-vec/pull/313). Unlike upstream, a quantized column's point query and full scan return its `float` vectors (upstream returns the quantized bytes, or reads past them), a connection finds the centroids another connection computed (upstream caches the column's state per connection and finds no rows), rowid filters and distance constraints apply, a renamed table's cells are found through an index named after the new table, so a table created under the old name works, an `int8` or `bit` column is refused (upstream reads past an `int8` column's vectors), errors and interrupts reach the statement, so an interrupted command fails with `interrupted` and the connection's queries go on finding its rows (upstream reports `SQL logic error`, and its KNN queries then return no rows), `VACUUM` works after a KNN query, the commands apply to every IVF column of the table, and an abbreviated option, a command value such as `nprobe=5x`, and `compute-centroids` options that are not a JSON object of integers are refused.
+
 ### Changed
 
 - Reading a vector column without a KNN query, as in `select embedding from v`, reads each chunk's vectors through one blob handle instead of opening the chunk's blob again for every row: selecting every vector of a 10,000-row `float[768]` table took 25 ms instead of 2,080 ms, and `select vec_distance_cosine(embedding, ?) from v` 41 ms instead of 2,093 ms.
@@ -10,6 +17,7 @@
 - A KNN query picks each chunk's k nearest rows with the heap of upstream's `SQLITE_VEC_EXPERIMENTAL_MIN_IDX` build option ([asg017/sqlite-vec@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f4570c1d65a0699e8e38d6c5562da9e23e)) instead of one pass over the chunk for every row it keeps, ordered so that it returns the same rows in the same order as before: on a 10,000-row `float[768]` table, `k = 1000` and `k = 4096` take 9.3 ms instead of 17.8 and 18.0 ms, `k = 100` 8.8 ms instead of 9.2, and `k = 1` 8.7 ms either way. A NaN distance ranks after every number instead of where its position put it. A static blob can produce one, and so can a `vec0` table written before 2.1.0, which started rejecting NaN elements. On such a table a NaN row used to push nearer rows out of the result or behind farther ones; queries now return the nearest rows in order, which can leave out a NaN row that a query with `k` smaller than the row count used to return.
 - A `DELETE` that removes the last row of a chunk now deletes the chunk's rows from the `_chunks`, `_vector_chunksNN`, and `_metadatachunksNN` shadow tables, so their pages are free for reuse without `optimize`; the emptied chunk used to stay until `optimize` removed it. Ported from [asg017@cb147c8](https://github.com/asg017/sqlite-vec/commit/cb147c8) via [vlasky@3a64182](https://github.com/vlasky/sqlite-vec/commit/3a64182) ([asg017#268](https://github.com/asg017/sqlite-vec/pull/268)).
 - A `DELETE` of a row whose bit in its chunk's validity bitmap is clear now fails with `vec0 deletion error: validity bit is not set`; it used to delete the row when a later row in the same byte of the bitmap was set. Only a damaged `_chunks` table has such a row.
+- A command that a `vec0` table does not know, as in `insert into t(t) values ('compact')`, now fails with `unknown vec0 command: 'compact'` instead of `SQL logic error`, as upstream [asg017@6e2c4c6](https://github.com/asg017/sqlite-vec/commit/6e2c4c6) does.
 
 ### Fixed
 

Makefile

@@ -211,7 +211,7 @@ test-loadable-watch:
 
 test-unit: | $(prefix)
 	$(CC) tests/test-unit.c sqlite-vec.c vendor/sqlite3.c \
-		-I./ -Ivendor -DSQLITE_CORE -o $(prefix)/test-unit -lm
+		-I./ -Ivendor -DSQLITE_CORE $(EXT_CFLAGS) -o $(prefix)/test-unit -lm
 	$(prefix)/test-unit
 
 test-all: test test-loadable test-unit