diff --git a/README.md b/README.md index 7c2651b..33ec50a 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,59 @@ LIMIT 10; Values outside the configured range are clamped silently. Calling `vector_quantize` or `vector_distance_q` without configuring a range returns a SQL error. +## Vector Index + +The `vector_index` virtual table is a faster alternative to `ORDER BY vector_distance(...)` for large tables. It stores vectors in chunks of about 1 MiB and scans them in Go across `GOMAXPROCS` goroutines, instead of calling a SQL function for every row. + +```sql +CREATE VIRTUAL TABLE docs_vec USING vector_index(); -- float32 +CREATE VIRTUAL TABLE docs_q USING vector_index(int8); -- quantized, requires WithQuantRange + +INSERT INTO docs_vec (rowid, embedding) VALUES (1, vector_encode('[0.1, 0.2, 0.3]')); + +-- k nearest neighbors +SELECT rowid, distance +FROM docs_vec +WHERE embedding MATCH vector_encode('[0.15, 0.25, 0.35]') AND k = 5; + +-- LIMIT also works when the query reads only docs_vec +SELECT rowid, distance +FROM docs_vec +WHERE embedding MATCH vector_encode('[0.15, 0.25, 0.35]') +ORDER BY distance +LIMIT 5; + +-- joins need k, because SQLite does not pass LIMIT through a join +SELECT d.content, v.distance +FROM docs_vec v +JOIN documents d ON d.id = v.rowid +WHERE v.embedding MATCH vector_encode('[0.15, 0.25, 0.35]') AND v.k = 5 +ORDER BY v.distance; +``` + +`embedding` is written and matched as a float32 blob. int8 tables quantize on write and on search using the `WithQuantRange` range, and return `vector_distance_q` distances. `UPDATE` and `DELETE` work as on a normal table. + +### Chunk size + +By default each chunk holds about 1 MiB of vector data: 2048 float32 vectors at 128 dimensions, 170 at 1536. `chunk_size=N` sets the number of vectors per chunk when the table is created: + +```sql +CREATE VIRTUAL TABLE docs_vec USING vector_index(float32, chunk_size=512); +``` + +Larger chunks mean fewer SQLite rows to read per search, but writes into a chunk slow down as it grows. Measured on Apple M3 Max at 128, 768, and 1536 dimensions: + +| Vector data per chunk | Insert, vs 1 MiB | Search p50, vs 1 MiB | +|---|---|---| +| 64 KiB | up to 6 µs per row faster; 10 µs slower for float32 at 768 dimensions | 14-133% slower | +| 256 KiB | up to 5 µs per row faster | 10-46% slower for float32; int8 within 10% | +| 1 MiB (default) | | | +| 4 MiB | 4.5-16x slower | at most 4% faster | + +The chunk size cannot be changed after creation; copy the rows into a new table to change it. + +The table keeps its data in ordinary shadow tables (`docs_vec_info`, `docs_vec_chunks`, `docs_vec_data`, `docs_vec_rowids`), so writes follow SQLite transactions. Reopening a database with a different dimension or quantization range than the table was created with returns an error. + ## Embedding Optional `vector_embed` function converts text to embeddings inside SQL. Provide an `Embedder` implementation via `WithEmbedder`: @@ -167,7 +220,7 @@ Each row returned by `vector_chunk` has a `value` (the chunk text) and a `chunk_ ## Design -- **Brute-force scan**: search is a linear scan over all rows, appropriate for SQLite-scale datasets (thousands to low millions of vectors). +- **Brute-force scan**: search is a linear scan over all rows, through the scalar functions or `vector_index`, appropriate for SQLite-scale datasets (thousands to low millions of vectors). - **Pure Go**: all vector math uses `encoding/binary` and `math` from the standard library. - **Single package**: everything lives in package `vector` at the module root. All internals are unexported. @@ -193,7 +246,7 @@ BenchmarkL2DistanceQuantized/dim=1536 3041773 398.2 ns/op 0 B/op 0 a ### SIFT1M -[SIFT1M](http://corpus-texmex.irisa.fr/) is a standard nearest-neighbor benchmark: 1,000,000 base vectors and 10,000 query vectors, 128 dimensions, L2 distance, with published ground-truth neighbors. The first 100 queries run single-threaded with k=100 against an on-disk database with default SQLite settings. +[SIFT1M](http://corpus-texmex.irisa.fr/) is a standard nearest-neighbor benchmark: 1,000,000 base vectors and 10,000 query vectors, 128 dimensions, L2 distance, with published ground-truth neighbors. The first 100 queries run with k=100 against an on-disk database with default SQLite settings. Every implementation searches on one thread except `vector_index`, which uses all 14 cores of the test machine. ``` curl -O ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz @@ -206,8 +259,10 @@ Results on Apple M3 Max, macOS 15.7, Go 1.24.12. Build time is the time to inser | Implementation | Search | Build | p50 | QPS | Recall@10 | Recall@100 | |---|---|---|---|---|---|---| -| go-sqlite-vector `vector_distance` | exact scan | 2.4s | 362 ms | 2.8 | 0.999 | 1.000 | -| go-sqlite-vector `vector_distance_q` (int8) | quantized scan | +3.1s | 339 ms | 3.0 | 0.983 | 0.988 | +| go-sqlite-vector `vector_distance` | exact scan | 2.6s | 385 ms | 2.6 | 0.999 | 1.000 | +| go-sqlite-vector `vector_distance_q` (int8) | quantized scan | +3.3s | 347 ms | 2.9 | 0.983 | 0.988 | +| go-sqlite-vector `vector_index` | exact chunk scan | 13.6s | 145 ms | 6.7 | 0.999 | 1.000 | +| go-sqlite-vector `vector_index(int8)` | quantized chunk scan | 13.6s | 39 ms | 25.0 | 0.983 | 0.988 | | sqlite-vec 0.1.9 `vec_distance_l2` | exact scan | 1.4s | 306 ms | 3.3 | 0.999 | 1.000 | | sqlite-vec 0.1.9 `vec0` | exact scan | 4.4s | 143 ms | 6.9 | 0.999 | 1.000 | | FAISS 1.15 `IndexFlatL2` | exact, in memory | 0.0s | 7.9 ms | 125 | 0.999 | 1.000 | @@ -216,13 +271,13 @@ Results on Apple M3 Max, macOS 15.7, Go 1.24.12. Build time is the time to inser Exact recall@10 is 0.999 rather than 1.000 because the ground truth contains tied distances. The int8 column stores 130 bytes per vector instead of 512. -Scans read the whole table through SQLite's pager. Memory-mapping the database file avoids a `pread` system call per page and speeds up scans of large tables: +Scans read the whole table through SQLite's pager. Memory-mapping the database file avoids a `pread` system call per page and speeds up scans of large tables. Set the size to at least the database file size: ```sql -PRAGMA mmap_size = 1073741824; -- 1 GiB +PRAGMA mmap_size = 4294967296; -- 4 GiB ``` -With this setting, SIFT1M float32 queries take 246 ms (p50) and int8 queries 286 ms. The comparison table above uses default settings for every SQLite implementation. +With this setting, SIFT1M p50 latency drops to 248 ms for `vector_distance`, 287 ms for `vector_distance_q`, 64 ms for `vector_index`, and 17 ms for `vector_index(int8)`. The comparison table above uses default settings for every SQLite implementation. ## License diff --git a/SPEC.md b/SPEC.md index 7526fe3..c1f663c 100644 --- a/SPEC.md +++ b/SPEC.md @@ -106,6 +106,80 @@ Computes squared L2 distance between the dequantized forms of two quantized int8 - **Output**: `REAL` (float64). - **Errors**: returns a SQL error if quantization was not configured. Returns a SQL error if either blob does not have the quantized format magic bytes (0x00, 0x01 prefix). Returns a SQL error if either blob's data length does not equal `dim` (total blob length = 2 + dim). +## vector_index Virtual Table + +`Register` also registers a virtual table module named `vector_index`. It stores vectors in chunks and answers k-nearest neighbor queries by scanning chunks in Go instead of evaluating a SQL function per row. + +```sql +CREATE VIRTUAL TABLE docs_vec USING vector_index(); -- float32 storage +CREATE VIRTUAL TABLE docs_vec USING vector_index(float32); -- same +CREATE VIRTUAL TABLE docs_vec USING vector_index(int8); -- quantized storage +CREATE VIRTUAL TABLE docs_vec USING vector_index(int8, chunk_size=4096); +``` + +Arguments are comma-separated and may appear in any order: + +- `float32` (default) or `int8`: storage type. +- `chunk_size=N`: vectors per chunk, a positive integer. When omitted, `N = max(1, 1048576 / vectorBytes)`, where `vectorBytes` is `dim * 4` for float32 and `dim` for int8, so each chunk holds about 1 MiB of vector data. + +The storage type and chunk size are fixed when the table is created and read from the `name_info` shadow table on later connections. + +Columns: + +| Column | Type | Description | +|---|---|---| +| `rowid` | INTEGER | Row identifier. Assigned automatically when omitted on insert. | +| `embedding` | BLOB | The vector. Written as a float32 blob of `dim * 4` bytes. Read back as a float32 blob, or a quantized blob for int8 tables. | +| `distance` | REAL, hidden | Squared L2 distance to the `MATCH` query. NULL outside a search. | +| `k` | INTEGER, hidden | Number of nearest neighbors to return. | + +### Search + +A search constrains `embedding MATCH` with a float32 query blob and gives k either as a `k` constraint or as `LIMIT`: + +```sql +SELECT rowid, distance FROM docs_vec WHERE embedding MATCH ?1 AND k = 10; + +SELECT rowid, distance FROM docs_vec WHERE embedding MATCH ?1 ORDER BY distance LIMIT 10; +``` + +- Results are ordered by ascending distance, then rowid. +- With `LIMIT n OFFSET m`, the search returns the `n + m` nearest rows and SQLite applies the offset. +- SQLite passes `LIMIT` to the virtual table only when the query reads that table alone. Queries that join, aggregate, or group must use `k`, or put the search in a subquery with its own `LIMIT`. +- For int8 tables, the query vector is quantized with the configured range and distances equal `vector_distance_q`. +- The chunk scan runs on `runtime.GOMAXPROCS(0)` goroutines. + +### Writes + +`INSERT`, `UPDATE`, and `DELETE` are supported. Inserts and updates take a float32 blob; int8 tables quantize it on write, clamping values outside the configured range. + +### Storage + +Each `vector_index` table `name` uses four shadow tables in the same database: + +- `name_info`: dimension, storage type, quantization range, and chunk size at creation. +- `name_chunks`: chunk id and number of occupied slots. +- `name_data`: per chunk, a blob of little-endian int64 rowids and a blob of vectors, each with capacity for the table's chunk size. int8 chunks store raw int8 values without the quantized blob header. +- `name_rowids`: rowid to chunk and slot. + +Deleting a vector moves the chunk's last vector into the freed slot, so occupied slots in every chunk are contiguous. Shadow tables are ordinary tables, so writes follow SQLite transactions. `ALTER TABLE ... RENAME` and `DROP TABLE` rename and drop the shadow tables. + +### Errors + +| Condition | Error message format | +|---|---| +| Unknown argument | `"vector_index: unknown argument %q, expected float32, int8, or chunk_size=N"` | +| Invalid chunk size | `"vector_index: chunk_size must be a positive integer, got %q"` | +| int8 without quantization | `"vector_index: int8 storage requires Register with WithQuantRange"` | +| Dimension differs from creation | `"vector_index: table %s has dimension %d, Register was called with %d"` | +| Storage type differs from creation | `"vector_index: table %s stores %s, declaration says %s"` | +| Quantization range differs from creation | `"vector_index: table %s was created with quantization range [%g, %g], Register was called with [%g, %g]"` | +| Embedding or query not a blob | `"vector_index: %s must be a float32 blob, got %v"` | +| Embedding or query size mismatch | `"vector_index: %s: expected %d bytes (dim=%d), got %d"` | +| Duplicate rowid | `"vector_index: UNIQUE constraint failed: %s.rowid (%d)"` | +| Search without k or LIMIT | `"vector_index: search requires k = ? or LIMIT"` | +| Negative k | `"vector_index: k must be non-negative, got %d"` | + ## Blob Formats ### Float32 Blob @@ -242,6 +316,7 @@ Table-driven integration tests executed against real SQLite connections. Each te - **vector_distance_q**: correct distance after dequantization, format validation, not-configured error, NULL input. - **Round-trip**: encode -> store -> retrieve -> distance pipeline. - **Float32ToBlob / BlobToFloat32**: Go-level encode/decode correctness. +- **vector_index**: search results match `vector_distance` / `vector_distance_q` brute force for `k`, `LIMIT`, `OFFSET`, and joins; results stay correct after deletes, updates, rowid changes, and rollbacks; stored embeddings read back unchanged; schema checks on reconnect; rename and drop. ## Benchmarks @@ -254,6 +329,5 @@ Go benchmarks (`testing.B`) for the core operations at dimensions 384, 768, and ## Design Notes -- **Extensibility**: the package is designed so that a virtual table module (via `sqlite.SetModule` / `SetModule`) could be added in the future without breaking the existing scalar function API. - **No CGo**: all vector math is implemented in pure Go using `encoding/binary` and `math` from the standard library. -- **Nearest neighbor scan**: search is a brute-force scan over all rows. This is appropriate for SQLite-scale datasets (thousands to low millions of vectors). For larger datasets, users should consider a dedicated vector database. +- **Nearest neighbor scan**: search is a brute-force scan over all rows, through either the scalar functions or `vector_index`. This is appropriate for SQLite-scale datasets (thousands to low millions of vectors). For larger datasets, users should consider a dedicated vector database. diff --git a/go.mod b/go.mod index 6202b22..5d71233 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,7 @@ go 1.23.0 toolchain go1.24.12 -require zombiezen.com/go/sqlite v1.4.2 +require zombiezen.com/go/sqlite v1.4.3-0.20260620165750-ca2041ed3dd2 require ( github.com/dustin/go-humanize v1.0.1 // indirect diff --git a/go.sum b/go.sum index 56f4fe4..140968c 100644 --- a/go.sum +++ b/go.sum @@ -47,5 +47,5 @@ modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= -zombiezen.com/go/sqlite v1.4.2 h1:KZXLrBuJ7tKNEm+VJcApLMeQbhmAUOKA5VWS93DfFRo= -zombiezen.com/go/sqlite v1.4.2/go.mod h1:5Kd4taTAD4MkBzT25mQ9uaAlLjyR0rFhsR6iINO70jc= +zombiezen.com/go/sqlite v1.4.3-0.20260620165750-ca2041ed3dd2 h1:gAUUjwioO8gmtw4a/5p8MmzYNWhLOGw32Sz1yXi4KnM= +zombiezen.com/go/sqlite v1.4.3-0.20260620165750-ca2041ed3dd2/go.mod h1:5Kd4taTAD4MkBzT25mQ9uaAlLjyR0rFhsR6iINO70jc= diff --git a/sift_test.go b/sift_test.go index bfaefb1..317e53a 100644 --- a/sift_test.go +++ b/sift_test.go @@ -72,6 +72,16 @@ func TestSIFT1M(t *testing.T) { exec(t, conn, "UPDATE sift SET embedding_q = vector_quantize(embedding)") quantizeTime := time.Since(start) + exec(t, conn, "CREATE VIRTUAL TABLE sift_index USING vector_index(float32)") + start = time.Now() + exec(t, conn, "INSERT INTO sift_index (rowid, embedding) SELECT id, embedding FROM sift") + indexTime := time.Since(start) + + exec(t, conn, "CREATE VIRTUAL TABLE sift_index_q USING vector_index(int8)") + start = time.Now() + exec(t, conn, "INSERT INTO sift_index_q (rowid, embedding) SELECT id, embedding FROM sift") + indexQTime := time.Since(start) + fi, err := os.Stat(dbPath) if err != nil { t.Fatal(err) @@ -80,8 +90,10 @@ func TestSIFT1M(t *testing.T) { t.Logf("SIFT1M: %d base vectors, %d queries, dim=%d, quant range [%g, %g]", len(base), nq, dim, lo, hi) t.Logf("insert float32: %v", insertTime.Round(time.Millisecond)) t.Logf("quantize column: %v", quantizeTime.Round(time.Millisecond)) - t.Logf("database size: %.1f MiB (both columns)", float64(fi.Size())/(1<<20)) - t.Logf("%-10s %10s %10s %10s %8s %8s %8s", "method", "p50", "p99", "QPS", "R@1", "R@10", "R@100") + t.Logf("insert index float32: %v", indexTime.Round(time.Millisecond)) + t.Logf("insert index int8: %v", indexQTime.Round(time.Millisecond)) + t.Logf("database size: %.1f MiB (all tables)", float64(fi.Size())/(1<<20)) + t.Logf("%-14s %10s %10s %10s %8s %8s %8s", "method", "p50", "p99", "QPS", "R@1", "R@10", "R@100") for _, m := range []struct { name string @@ -90,6 +102,8 @@ func TestSIFT1M(t *testing.T) { }{ {"float32", "SELECT id FROM sift ORDER BY vector_distance(embedding, ?1) LIMIT 100", Float32ToBlob}, {"int8", "SELECT id FROM sift ORDER BY vector_distance_q(embedding_q, ?1) LIMIT 100", func(v []float32) []byte { return quantize(v, lo, hi) }}, + {"index float32", "SELECT rowid FROM sift_index WHERE embedding MATCH ?1 AND k = 100", Float32ToBlob}, + {"index int8", "SELECT rowid FROM sift_index_q WHERE embedding MATCH ?1 AND k = 100", Float32ToBlob}, } { lat := make([]time.Duration, nq) var r1, r10, r100 float64 @@ -120,7 +134,7 @@ func TestSIFT1M(t *testing.T) { } slices.Sort(lat) n := float64(nq) - t.Logf("%-10s %10v %10v %10.2f %8.4f %8.4f %8.4f", m.name, + t.Logf("%-14s %10v %10v %10.2f %8.4f %8.4f %8.4f", m.name, lat[nq/2].Round(time.Millisecond), lat[nq*99/100].Round(time.Millisecond), n/total.Seconds(), r1/n, r10/n, r100/n) } diff --git a/vector.go b/vector.go index b179970..0974bde 100644 --- a/vector.go +++ b/vector.go @@ -202,6 +202,14 @@ func Register(conn *sqlite.Conn, dim int, opts ...Option) error { return err } + err = conn.SetModule("vector_index", &sqlite.Module{ + Create: indexConnect(cfg, true), + Connect: indexConnect(cfg, false), + }) + if err != nil { + return err + } + err = conn.SetModule("vector_chunk", &sqlite.Module{ Connect: func(c *sqlite.Conn, opts *sqlite.VTableConnectOptions) (sqlite.VTable, *sqlite.VTableConfig, error) { return &chunkVTable{chunker: cfg.chunker}, &sqlite.VTableConfig{ @@ -269,14 +277,19 @@ func f32At(b []byte, i int) float64 { // Dequantization is affine, so each component difference is // (qa - qb) * (max - min) / 255. func l2SquaredQuantized(a, b []byte, min, max float32) float64 { - a, b = a[2:], b[2:len(a)] + scale := float64(max-min) / 255 + return float64(l2SquaredInt8(a[2:], b[2:])) * scale * scale +} + +// l2SquaredInt8 sums squared differences of raw int8 values of equal length. +func l2SquaredInt8(a, b []byte) int64 { + b = b[:len(a)] var sum int64 for i := range a { d := int64(int8(a[i])) - int64(int8(b[i])) sum += d * d } - scale := float64(max-min) / 255 - return float64(sum) * scale * scale + return sum } func isQuantizedBlob(b []byte) bool { diff --git a/vector_index.go b/vector_index.go new file mode 100644 index 0000000..5ac6068 --- /dev/null +++ b/vector_index.go @@ -0,0 +1,730 @@ +package vector + +import ( + "encoding/binary" + "fmt" + "io" + "runtime" + "slices" + "strconv" + "strings" + "sync" + + "zombiezen.com/go/sqlite" + "zombiezen.com/go/sqlite/sqlitex" +) + +// indexChunkTargetBytes is the default amount of vector data per chunk. +// Searches read whole chunks, so larger chunks mean fewer SQLite rows per +// query, but writes into a chunk get slower as it grows. A sweep over 128 to +// 1536 dimensions found searches fastest near 1 MiB, and inserts 4.5-16x +// slower at 4 MiB than at 1 MiB. +const indexChunkTargetBytes = 1 << 20 + +func defaultChunkCapacity(vectorBytes int) int { + return max(1, indexChunkTargetBytes/vectorBytes) +} + +const ( + indexColEmbedding = iota + indexColDistance + indexColK +) + +const ( + indexPlanScan = iota + indexPlanRowID + indexPlanKNN +) + +type vectorIndex struct { + conn *sqlite.Conn + cfg *config + quantized bool + capacity int + db string + name string +} + +func indexConnect(cfg *config, create bool) sqlite.VTableConnectFunc { + return func(conn *sqlite.Conn, opts *sqlite.VTableConnectOptions) (sqlite.VTable, *sqlite.VTableConfig, error) { + vi := &vectorIndex{conn: conn, cfg: cfg, db: opts.DatabaseName, name: opts.VTableName} + for _, arg := range opts.Args { + arg = strings.TrimSpace(arg) + key, val, hasVal := strings.Cut(arg, "=") + switch { + case arg == "" && len(opts.Args) == 1, arg == "float32": + case arg == "int8": + vi.quantized = true + case hasVal && strings.TrimSpace(key) == "chunk_size": + n, err := strconv.Atoi(strings.TrimSpace(val)) + if err != nil || n < 1 { + return nil, nil, fmt.Errorf("vector_index: chunk_size must be a positive integer, got %q", strings.TrimSpace(val)) + } + vi.capacity = n + default: + return nil, nil, fmt.Errorf("vector_index: unknown argument %q, expected float32, int8, or chunk_size=N", arg) + } + } + if vi.quantized && !cfg.quantEnabled { + return nil, nil, fmt.Errorf("vector_index: int8 storage requires Register with WithQuantRange") + } + var err error + if create { + if vi.capacity == 0 { + vi.capacity = defaultChunkCapacity(vi.vectorSize()) + } + err = vi.create() + } else { + err = vi.check() + } + if err != nil { + return nil, nil, err + } + return vi, &sqlite.VTableConfig{ + Declaration: "CREATE TABLE x(embedding BLOB, distance REAL HIDDEN, k INTEGER HIDDEN)", + }, nil + } +} + +func quoteIdent(s string) string { + return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` +} + +func (vi *vectorIndex) shadow(suffix string) string { + return quoteIdent(vi.db) + "." + quoteIdent(vi.name+"_"+suffix) +} + +func (vi *vectorIndex) storageType() string { + if vi.quantized { + return "int8" + } + return "float32" +} + +// vectorSize is the number of bytes one stored vector occupies in a chunk. +// int8 chunks store raw quantized values without the quantized blob header. +func (vi *vectorIndex) vectorSize() int { + if vi.quantized { + return vi.cfg.dim + } + return vi.cfg.dim * 4 +} + +func (vi *vectorIndex) exec(query string, args ...any) error { + return sqlitex.Execute(vi.conn, query, &sqlitex.ExecOptions{Args: args}) +} + +func (vi *vectorIndex) create() error { + var qmin, qmax any + if vi.quantized { + qmin, qmax = float64(vi.cfg.quantMin), float64(vi.cfg.quantMax) + } + stmts := []struct { + query string + args []any + }{ + {"CREATE TABLE " + vi.shadow("info") + " (dim INTEGER NOT NULL, type TEXT NOT NULL, quant_min REAL, quant_max REAL, chunk_size INTEGER NOT NULL)", nil}, + {"INSERT INTO " + vi.shadow("info") + " VALUES (?, ?, ?, ?, ?)", []any{vi.cfg.dim, vi.storageType(), qmin, qmax, vi.capacity}}, + // Chunk sizes live apart from chunk data because updating any column + // rewrites the whole row, blobs included. UNIQUE (size, id) indexes + // size for finding a chunk with free slots; SQLite's automatic index + // follows the table through renames. + {"CREATE TABLE " + vi.shadow("chunks") + " (id INTEGER PRIMARY KEY, size INTEGER NOT NULL, UNIQUE (size, id))", nil}, + {"CREATE TABLE " + vi.shadow("data") + " (id INTEGER PRIMARY KEY, ids BLOB NOT NULL, vectors BLOB NOT NULL)", nil}, + {"CREATE TABLE " + vi.shadow("rowids") + " (rowid INTEGER PRIMARY KEY, chunk INTEGER NOT NULL, slot INTEGER NOT NULL)", nil}, + } + for _, s := range stmts { + if err := vi.exec(s.query, s.args...); err != nil { + return fmt.Errorf("vector_index: create %s: %w", vi.name, err) + } + } + return nil +} + +func (vi *vectorIndex) check() error { + var found bool + var dim int + var typ string + var qmin, qmax float32 + err := sqlitex.Execute(vi.conn, "SELECT dim, type, quant_min, quant_max, chunk_size FROM "+vi.shadow("info"), &sqlitex.ExecOptions{ + ResultFunc: func(stmt *sqlite.Stmt) error { + found = true + dim = stmt.ColumnInt(0) + typ = stmt.ColumnText(1) + qmin = float32(stmt.ColumnFloat(2)) + qmax = float32(stmt.ColumnFloat(3)) + vi.capacity = stmt.ColumnInt(4) + return nil + }, + }) + if err != nil { + return fmt.Errorf("vector_index: connect %s: %w", vi.name, err) + } + switch { + case !found: + return fmt.Errorf("vector_index: connect %s: missing metadata", vi.name) + case dim != vi.cfg.dim: + return fmt.Errorf("vector_index: table %s has dimension %d, Register was called with %d", vi.name, dim, vi.cfg.dim) + case typ != vi.storageType(): + return fmt.Errorf("vector_index: table %s stores %s, declaration says %s", vi.name, typ, vi.storageType()) + case vi.quantized && (qmin != vi.cfg.quantMin || qmax != vi.cfg.quantMax): + return fmt.Errorf("vector_index: table %s was created with quantization range [%g, %g], Register was called with [%g, %g]", + vi.name, qmin, qmax, vi.cfg.quantMin, vi.cfg.quantMax) + } + return nil +} + +func (vi *vectorIndex) BestIndex(in *sqlite.IndexInputs) (*sqlite.IndexOutputs, error) { + out := &sqlite.IndexOutputs{ + ConstraintUsage: make([]sqlite.IndexConstraintUsage, len(in.Constraints)), + EstimatedCost: 1e9, + EstimatedRows: 1e6, + } + match, k, limit, offset, rowid := -1, -1, -1, -1, -1 + for i, c := range in.Constraints { + if !c.Usable { + continue + } + switch { + case c.Op == sqlite.IndexConstraintMatch && c.Column == indexColEmbedding: + match = i + case c.Op == sqlite.IndexConstraintEq && c.Column == indexColK: + k = i + case c.Op == sqlite.IndexConstraintLimit: + limit = i + case c.Op == sqlite.IndexConstraintOffset: + offset = i + case c.Op == sqlite.IndexConstraintEq && c.Column == -1: + rowid = i + } + } + switch { + case match >= 0: + out.ID.Num = indexPlanKNN + out.ConstraintUsage[match] = sqlite.IndexConstraintUsage{ArgvIndex: 1, Omit: true} + // LIMIT and OFFSET are not omitted: SQLite still applies them to the + // rows the search returns. + switch { + case k >= 0: + out.ID.String = "k" + out.ConstraintUsage[k] = sqlite.IndexConstraintUsage{ArgvIndex: 2, Omit: true} + case limit >= 0 && offset >= 0: + out.ID.String = "limit,offset" + out.ConstraintUsage[limit] = sqlite.IndexConstraintUsage{ArgvIndex: 2} + out.ConstraintUsage[offset] = sqlite.IndexConstraintUsage{ArgvIndex: 3} + case limit >= 0: + out.ID.String = "limit" + out.ConstraintUsage[limit] = sqlite.IndexConstraintUsage{ArgvIndex: 2} + } + out.OrderByConsumed = len(in.OrderBy) == 1 && in.OrderBy[0].Column == indexColDistance && !in.OrderBy[0].Desc + out.EstimatedCost = 1e3 + out.EstimatedRows = 100 + case rowid >= 0: + out.ID.Num = indexPlanRowID + out.ConstraintUsage[rowid] = sqlite.IndexConstraintUsage{ArgvIndex: 1, Omit: true} + out.EstimatedCost = 1 + out.EstimatedRows = 1 + out.IndexFlags = sqlite.IndexScanUnique + } + return out, nil +} + +func (vi *vectorIndex) Open() (sqlite.VTableCursor, error) { + return &indexCursor{vi: vi}, nil +} + +func (vi *vectorIndex) Disconnect() error { return nil } + +func (vi *vectorIndex) Destroy() error { + for _, suffix := range []string{"info", "chunks", "data", "rowids"} { + if err := vi.exec("DROP TABLE " + vi.shadow(suffix)); err != nil { + return fmt.Errorf("vector_index: destroy %s: %w", vi.name, err) + } + } + return nil +} + +func (vi *vectorIndex) Rename(newName string) error { + for _, suffix := range []string{"info", "chunks", "data", "rowids"} { + if err := vi.exec("ALTER TABLE " + vi.shadow(suffix) + " RENAME TO " + quoteIdent(newName+"_"+suffix)); err != nil { + return fmt.Errorf("vector_index: rename %s: %w", vi.name, err) + } + } + vi.name = newName + return nil +} + +// encodeInput validates a float32 embedding blob and converts it to the +// stored representation. +func (vi *vectorIndex) encodeInput(v sqlite.Value, what string) ([]byte, error) { + if v.Type() != sqlite.TypeBlob { + return nil, fmt.Errorf("vector_index: %s must be a float32 blob, got %v", what, v.Type()) + } + b := v.Blob() + if len(b) != vi.cfg.dim*4 { + return nil, fmt.Errorf("vector_index: %s: expected %d bytes (dim=%d), got %d", what, vi.cfg.dim*4, vi.cfg.dim, len(b)) + } + if !vi.quantized { + return slices.Clone(b), nil + } + floats, _ := BlobToFloat32(b) + return quantize(floats, vi.cfg.quantMin, vi.cfg.quantMax)[2:], nil +} + +func (vi *vectorIndex) locate(rowid int64) (chunk int64, slot int, ok bool, err error) { + err = sqlitex.Execute(vi.conn, "SELECT chunk, slot FROM "+vi.shadow("rowids")+" WHERE rowid = ?", &sqlitex.ExecOptions{ + Args: []any{rowid}, + ResultFunc: func(stmt *sqlite.Stmt) error { + chunk, slot, ok = stmt.ColumnInt64(0), stmt.ColumnInt(1), true + return nil + }, + }) + return chunk, slot, ok, err +} + +func (vi *vectorIndex) readAt(column string, chunk int64, offset int, buf []byte) error { + blob, err := vi.conn.OpenBlob(vi.db, vi.name+"_data", column, chunk, false) + if err != nil { + return err + } + defer blob.Close() + if _, err := blob.Seek(int64(offset), io.SeekStart); err != nil { + return err + } + _, err = io.ReadFull(blob, buf) + return err +} + +func (vi *vectorIndex) writeAt(column string, chunk int64, offset int, data []byte) error { + blob, err := vi.conn.OpenBlob(vi.db, vi.name+"_data", column, chunk, true) + if err != nil { + return err + } + if _, err := blob.Seek(int64(offset), io.SeekStart); err != nil { + blob.Close() + return err + } + if _, err := blob.Write(data); err != nil { + blob.Close() + return err + } + return blob.Close() +} + +func (vi *vectorIndex) readVector(chunk int64, slot int) ([]byte, error) { + buf := make([]byte, vi.vectorSize()) + if err := vi.readAt("vectors", chunk, slot*len(buf), buf); err != nil { + return nil, err + } + if vi.quantized { + return append([]byte{0x00, 0x01}, buf...), nil + } + return buf, nil +} + +func (vi *vectorIndex) Update(params sqlite.VTableUpdateParams) (int64, error) { + emb := params.Columns[indexColEmbedding] + if params.IsInsert() { + vec, err := vi.encodeInput(emb, "embedding") + if err != nil { + return 0, err + } + return vi.insert(params.NewRowID, vec) + } + + oldID, newID := params.OldRowID.Int64(), params.NewRowID.Int64() + chunk, slot, _, err := vi.locate(oldID) + if err != nil { + return 0, err + } + var vec []byte + if emb.NoChange() { + if newID == oldID { + return oldID, nil + } + if vec, err = vi.readVector(chunk, slot); err != nil { + return 0, err + } + if vi.quantized { + vec = vec[2:] + } + } else if vec, err = vi.encodeInput(emb, "embedding"); err != nil { + return 0, err + } + if newID == oldID { + return oldID, vi.writeAt("vectors", chunk, slot*vi.vectorSize(), vec) + } + if err := vi.checkRowIDFree(newID); err != nil { + return 0, err + } + if err := vi.DeleteRow(params.OldRowID); err != nil { + return 0, err + } + return vi.insert(params.NewRowID, vec) +} + +func (vi *vectorIndex) checkRowIDFree(rowid int64) error { + _, _, exists, err := vi.locate(rowid) + if err != nil { + return err + } + if exists { + return fmt.Errorf("vector_index: UNIQUE constraint failed: %s.rowid (%d): %w", vi.name, rowid, sqlite.ResultConstraintPrimaryKey.ToError()) + } + return nil +} + +func (vi *vectorIndex) insert(rowid sqlite.Value, vec []byte) (int64, error) { + if rowid.Type() != sqlite.TypeNull { + if err := vi.checkRowIDFree(rowid.Int64()); err != nil { + return 0, err + } + } + + var chunk int64 + var size int + var found bool + err := sqlitex.Execute(vi.conn, "SELECT id, size FROM "+vi.shadow("chunks")+" WHERE size < ? LIMIT 1", &sqlitex.ExecOptions{ + Args: []any{vi.capacity}, + ResultFunc: func(stmt *sqlite.Stmt) error { + chunk, size, found = stmt.ColumnInt64(0), stmt.ColumnInt(1), true + return nil + }, + }) + if err != nil { + return 0, err + } + if !found { + if err := vi.exec("INSERT INTO " + vi.shadow("chunks") + " (size) VALUES (0)"); err != nil { + return 0, err + } + chunk = vi.conn.LastInsertRowID() + err := vi.exec("INSERT INTO "+vi.shadow("data")+" (id, ids, vectors) VALUES (?, zeroblob(?), zeroblob(?))", + chunk, vi.capacity*8, vi.capacity*vi.vectorSize()) + if err != nil { + return 0, err + } + } + + if rowid.Type() == sqlite.TypeNull { + err = vi.exec("INSERT INTO "+vi.shadow("rowids")+" (chunk, slot) VALUES (?, ?)", chunk, size) + } else { + err = vi.exec("INSERT INTO "+vi.shadow("rowids")+" (rowid, chunk, slot) VALUES (?, ?, ?)", rowid.Int64(), chunk, size) + } + if err != nil { + return 0, err + } + id := vi.conn.LastInsertRowID() + + if err := vi.writeAt("ids", chunk, size*8, binary.LittleEndian.AppendUint64(nil, uint64(id))); err != nil { + return 0, err + } + if err := vi.writeAt("vectors", chunk, size*vi.vectorSize(), vec); err != nil { + return 0, err + } + return id, vi.exec("UPDATE "+vi.shadow("chunks")+" SET size = size + 1 WHERE id = ?", chunk) +} + +// DeleteRow removes a vector by moving the chunk's last vector into its slot, +// which keeps every chunk's occupied slots contiguous for scanning. +func (vi *vectorIndex) DeleteRow(rowid sqlite.Value) error { + id := rowid.Int64() + chunk, slot, ok, err := vi.locate(id) + if err != nil || !ok { + return err + } + var size int + err = sqlitex.Execute(vi.conn, "SELECT size FROM "+vi.shadow("chunks")+" WHERE id = ?", &sqlitex.ExecOptions{ + Args: []any{chunk}, + ResultFunc: func(stmt *sqlite.Stmt) error { + size = stmt.ColumnInt(0) + return nil + }, + }) + if err != nil { + return err + } + + if last := size - 1; slot != last { + vs := vi.vectorSize() + movedID := make([]byte, 8) + movedVec := make([]byte, vs) + if err := vi.readAt("ids", chunk, last*8, movedID); err != nil { + return err + } + if err := vi.readAt("vectors", chunk, last*vs, movedVec); err != nil { + return err + } + if err := vi.writeAt("ids", chunk, slot*8, movedID); err != nil { + return err + } + if err := vi.writeAt("vectors", chunk, slot*vs, movedVec); err != nil { + return err + } + err := vi.exec("UPDATE "+vi.shadow("rowids")+" SET slot = ? WHERE rowid = ?", slot, int64(binary.LittleEndian.Uint64(movedID))) + if err != nil { + return err + } + } + + if size == 1 { + if err = vi.exec("DELETE FROM "+vi.shadow("chunks")+" WHERE id = ?", chunk); err == nil { + err = vi.exec("DELETE FROM "+vi.shadow("data")+" WHERE id = ?", chunk) + } + } else { + err = vi.exec("UPDATE "+vi.shadow("chunks")+" SET size = size - 1 WHERE id = ?", chunk) + } + if err != nil { + return err + } + return vi.exec("DELETE FROM "+vi.shadow("rowids")+" WHERE rowid = ?", id) +} + +type indexResult struct { + dist float64 + id int64 +} + +// resultHeap is a max-heap on dist that keeps the k nearest results seen. +type resultHeap []indexResult + +func (h *resultHeap) push(r indexResult, k int) { + s := *h + if len(s) < k { + s = append(s, r) + for i := len(s) - 1; i > 0; { + p := (i - 1) / 2 + if s[p].dist >= s[i].dist { + break + } + s[p], s[i] = s[i], s[p] + i = p + } + *h = s + return + } + if r.dist >= s[0].dist { + return + } + s[0] = r + for i := 0; ; { + l, big := 2*i+1, i + if l < len(s) && s[l].dist > s[big].dist { + big = l + } + if l+1 < len(s) && s[l+1].dist > s[big].dist { + big = l + 1 + } + if big == i { + break + } + s[i], s[big] = s[big], s[i] + i = big + } +} + +type indexChunk struct { + ids []byte + vectors []byte +} + +type indexCursor struct { + vi *vectorIndex + + // k-NN and rowid plans materialize results; the scan plan steps stmt. + knn bool + k int64 + results []indexResult + pos int + + stmt *sqlite.Stmt + eof bool +} + +func (cur *indexCursor) Filter(id sqlite.IndexID, argv []sqlite.Value) error { + if err := cur.Close(); err != nil { + return err + } + cur.knn, cur.results, cur.pos = false, nil, 0 + switch id.Num { + case indexPlanKNN: + cur.knn = true + return cur.search(id.String, argv) + case indexPlanRowID: + rowid := argv[0].Int64() + if _, _, ok, err := cur.vi.locate(rowid); err != nil { + return err + } else if ok { + cur.results = []indexResult{{id: rowid}} + } + return nil + default: + stmt, _, err := cur.vi.conn.PrepareTransient("SELECT rowid, chunk, slot FROM " + cur.vi.shadow("rowids") + " ORDER BY rowid") + if err != nil { + return err + } + cur.stmt = stmt + return cur.Next() + } +} + +func (cur *indexCursor) search(mode string, argv []sqlite.Value) error { + vi := cur.vi + switch mode { + case "k": + cur.k = argv[1].Int64() + if cur.k < 0 { + return fmt.Errorf("vector_index: k must be non-negative, got %d", cur.k) + } + case "limit", "limit,offset": + cur.k = argv[1].Int64() + if cur.k < 0 { + return fmt.Errorf("vector_index: search requires k = ? or a non-negative LIMIT") + } + if mode == "limit,offset" { + cur.k += max(argv[2].Int64(), 0) + } + default: + return fmt.Errorf("vector_index: search requires k = ? or LIMIT") + } + query, err := vi.encodeInput(argv[0], "MATCH query") + if err != nil { + return err + } + if cur.k == 0 { + return nil + } + + dist := l2Squared + if vi.quantized { + scale := float64(vi.cfg.quantMax-vi.cfg.quantMin) / 255 + dist = func(a, b []byte) float64 { return float64(l2SquaredInt8(a, b)) * scale * scale } + } + vs := vi.vectorSize() + k := int(cur.k) + + workers := runtime.GOMAXPROCS(0) + heaps := make([]resultHeap, workers) + jobs := make(chan *indexChunk) + pool := sync.Pool{New: func() any { return new(indexChunk) }} + var wg sync.WaitGroup + for w := range heaps { + wg.Add(1) + go func(h *resultHeap) { + defer wg.Done() + for c := range jobs { + for i := 0; i < len(c.ids)/8; i++ { + d := dist(c.vectors[i*vs:(i+1)*vs], query) + h.push(indexResult{d, int64(binary.LittleEndian.Uint64(c.ids[i*8:]))}, k) + } + pool.Put(c) + } + }(&heaps[w]) + } + err = sqlitex.Execute(vi.conn, "SELECT c.size, d.ids, d.vectors FROM "+vi.shadow("chunks")+" c JOIN "+vi.shadow("data")+" d ON d.id = c.id", &sqlitex.ExecOptions{ + ResultFunc: func(stmt *sqlite.Stmt) error { + n := stmt.ColumnInt(0) + c := pool.Get().(*indexChunk) + c.ids = slices.Grow(c.ids[:0], n*8)[:n*8] + c.vectors = slices.Grow(c.vectors[:0], n*vs)[:n*vs] + stmt.ColumnBytes(1, c.ids) + stmt.ColumnBytes(2, c.vectors) + jobs <- c + return nil + }, + }) + close(jobs) + wg.Wait() + if err != nil { + return err + } + + var all resultHeap + for _, h := range heaps { + for _, r := range h { + all.push(r, k) + } + } + slices.SortFunc(all, func(a, b indexResult) int { + switch { + case a.dist < b.dist: + return -1 + case a.dist > b.dist: + return 1 + case a.id < b.id: + return -1 + case a.id > b.id: + return 1 + } + return 0 + }) + cur.results = all + return nil +} + +func (cur *indexCursor) Next() error { + if cur.stmt == nil { + cur.pos++ + return nil + } + row, err := cur.stmt.Step() + cur.eof = !row + return err +} + +func (cur *indexCursor) EOF() bool { + if cur.stmt == nil { + return cur.pos >= len(cur.results) + } + return cur.eof +} + +func (cur *indexCursor) RowID() (int64, error) { + if cur.stmt == nil { + return cur.results[cur.pos].id, nil + } + return cur.stmt.ColumnInt64(0), nil +} + +func (cur *indexCursor) Column(i int, noChange bool) (sqlite.Value, error) { + switch i { + case indexColEmbedding: + if noChange { + return sqlite.Unchanged(), nil + } + var chunk int64 + var slot int + if cur.stmt != nil { + chunk, slot = cur.stmt.ColumnInt64(1), cur.stmt.ColumnInt(2) + } else { + var err error + if chunk, slot, _, err = cur.vi.locate(cur.results[cur.pos].id); err != nil { + return sqlite.Value{}, err + } + } + vec, err := cur.vi.readVector(chunk, slot) + if err != nil { + return sqlite.Value{}, err + } + return sqlite.BlobValue(vec), nil + case indexColDistance: + if cur.knn { + return sqlite.FloatValue(cur.results[cur.pos].dist), nil + } + case indexColK: + if cur.knn { + return sqlite.IntegerValue(cur.k), nil + } + } + return sqlite.Value{}, nil +} + +func (cur *indexCursor) Close() error { + if cur.stmt == nil { + return nil + } + err := cur.stmt.Finalize() + cur.stmt = nil + return err +} diff --git a/vector_index_test.go b/vector_index_test.go new file mode 100644 index 0000000..0885c35 --- /dev/null +++ b/vector_index_test.go @@ -0,0 +1,460 @@ +package vector + +import ( + "bytes" + "fmt" + "path/filepath" + "strings" + "testing" + + "zombiezen.com/go/sqlite" + "zombiezen.com/go/sqlite/sqlitex" +) + +// indexTestRows spans more than two chunks of indexTestChunkSize so searches, +// deletes, and slot moves cross chunk boundaries. +const ( + indexTestChunkSize = 1024 + indexTestRows = 2*indexTestChunkSize + 500 +) + +func indexTestVector(i int) []float32 { + return []float32{float32(i % 97), float32(i % 89), float32(i % 83), float32(i % 79)} +} + +type indexRow struct { + id int64 + dist float64 +} + +func queryRows(t *testing.T, conn *sqlite.Conn, query string, args ...any) []indexRow { + t.Helper() + var rows []indexRow + err := sqlitex.Execute(conn, query, &sqlitex.ExecOptions{ + Args: args, + ResultFunc: func(stmt *sqlite.Stmt) error { + rows = append(rows, indexRow{stmt.ColumnInt64(0), stmt.ColumnFloat(1)}) + return nil + }, + }) + if err != nil { + t.Fatalf("%s: %v", query, err) + } + return rows +} + +func execArgs(t *testing.T, conn *sqlite.Conn, query string, args ...any) { + t.Helper() + if err := sqlitex.Execute(conn, query, &sqlitex.ExecOptions{Args: args}); err != nil { + t.Fatalf("%s: %v", query, err) + } +} + +// openIndexTestConn creates a vector_index table "idx" of the given storage +// type and a plain reference table "ref" holding the same vectors. +func openIndexTestConn(t *testing.T, storage string) *sqlite.Conn { + t.Helper() + conn := openTestConn(t) + if err := Register(conn, 4, WithQuantRange(0, 100)); err != nil { + t.Fatal(err) + } + execArgs(t, conn, fmt.Sprintf("CREATE VIRTUAL TABLE idx USING vector_index(%s, chunk_size=%d)", storage, indexTestChunkSize)) + execArgs(t, conn, "CREATE TABLE ref (id INTEGER PRIMARY KEY, e BLOB, q BLOB)") + execArgs(t, conn, "BEGIN") + for i := range indexTestRows { + b := Float32ToBlob(indexTestVector(i)) + execArgs(t, conn, "INSERT INTO idx (rowid, embedding) VALUES (?, ?)", i, b) + execArgs(t, conn, "INSERT INTO ref (id, e, q) VALUES (?, ?, vector_quantize(?))", i, b, b) + } + execArgs(t, conn, "COMMIT") + return conn +} + +func refDistance(storage string) string { + if storage == "int8" { + return "vector_distance_q(q, vector_quantize(?1))" + } + return "vector_distance(e, ?1)" +} + +func assertRowsEqual(t *testing.T, got, want []indexRow) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("got %d rows, want %d", len(got), len(want)) + } + for i := range got { + if got[i] != want[i] { + t.Fatalf("row %d = %+v, want %+v", i, got[i], want[i]) + } + } +} + +func TestVectorIndexSearch(t *testing.T) { + query := Float32ToBlob([]float32{50, 40, 30, 20}) + tests := []struct { + name string + query string + ref string + }{ + { + name: "k constraint", + query: "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 AND k = 10", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id LIMIT 10", + }, + { + name: "LIMIT", + query: "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 ORDER BY distance LIMIT 25", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id LIMIT 25", + }, + { + name: "LIMIT with OFFSET", + query: "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 ORDER BY distance LIMIT 10 OFFSET 5", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id LIMIT 10 OFFSET 5", + }, + { + name: "k larger than table", + query: "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 AND k = 100000", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id", + }, + { + name: "k zero", + query: "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 AND k = 0", + ref: "SELECT id, %DIST% AS d FROM ref WHERE 0", + }, + { + name: "join with k", + query: "SELECT r.id, v.distance FROM idx v JOIN ref r ON r.id = v.rowid WHERE v.embedding MATCH ?1 AND v.k = 10 ORDER BY v.distance", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id LIMIT 10", + }, + { + name: "join with LIMIT in subquery", + query: "SELECT r.id, v.distance FROM (SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 ORDER BY distance LIMIT 10) v JOIN ref r ON r.id = v.rowid ORDER BY v.distance", + ref: "SELECT id, %DIST% AS d FROM ref ORDER BY d, id LIMIT 10", + }, + } + for _, storage := range []string{"float32", "int8"} { + conn := openIndexTestConn(t, storage) + for _, tt := range tests { + t.Run(storage+"/"+tt.name, func(t *testing.T) { + got := queryRows(t, conn, tt.query, query) + want := queryRows(t, conn, strings.ReplaceAll(tt.ref, "%DIST%", refDistance(storage)), query) + assertRowsEqual(t, got, want) + }) + } + } +} + +func TestVectorIndexWrites(t *testing.T) { + tests := []struct { + name string + idx []string + ref []string + }{ + { + name: "delete from middle of chunks", + idx: []string{"DELETE FROM idx WHERE rowid % 7 = 3"}, + ref: []string{"DELETE FROM ref WHERE id % 7 = 3"}, + }, + { + name: "delete entire first chunk", + idx: []string{"DELETE FROM idx WHERE rowid < 1100"}, + ref: []string{"DELETE FROM ref WHERE id < 1100"}, + }, + { + name: "delete then insert reuses space", + idx: []string{"DELETE FROM idx WHERE rowid < 300", "INSERT INTO idx (rowid, embedding) VALUES (100000, X'0000c8420000c8420000c8420000c842')"}, + ref: []string{"DELETE FROM ref WHERE id < 300", "INSERT INTO ref (id, e, q) VALUES (100000, X'0000c8420000c8420000c8420000c842', vector_quantize(X'0000c8420000c8420000c8420000c842'))"}, + }, + { + name: "update embedding", + idx: []string{"UPDATE idx SET embedding = X'00004842000020420000f0410000a041' WHERE rowid IN (5, 1500, 2400)"}, + ref: []string{"UPDATE ref SET e = X'00004842000020420000f0410000a041', q = vector_quantize(X'00004842000020420000f0410000a041') WHERE id IN (5, 1500, 2400)"}, + }, + { + name: "change rowid", + idx: []string{"UPDATE idx SET rowid = rowid + 1000000 WHERE rowid IN (7, 2000)"}, + ref: []string{"UPDATE ref SET id = id + 1000000 WHERE id IN (7, 2000)"}, + }, + { + name: "rolled back insert is not visible", + idx: []string{"BEGIN", "INSERT INTO idx (rowid, embedding) VALUES (200000, X'0000c8420000c8420000c8420000c842')", "ROLLBACK"}, + }, + { + name: "insert without rowid assigns one", + idx: []string{"INSERT INTO idx (embedding) VALUES (X'000080400000804000008040000080C0')"}, + ref: []string{"INSERT INTO ref (e, q) VALUES (X'000080400000804000008040000080C0', vector_quantize(X'000080400000804000008040000080C0'))"}, + }, + } + query := Float32ToBlob([]float32{50, 50, 50, 50}) + for _, storage := range []string{"float32", "int8"} { + for _, tt := range tests { + t.Run(storage+"/"+tt.name, func(t *testing.T) { + conn := openIndexTestConn(t, storage) + for _, q := range tt.idx { + execArgs(t, conn, q) + } + for _, q := range tt.ref { + execArgs(t, conn, q) + } + got := queryRows(t, conn, "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 AND k = 1000000", query) + want := queryRows(t, conn, "SELECT id, "+refDistance(storage)+" AS d FROM ref ORDER BY d, id", query) + assertRowsEqual(t, got, want) + + col := "e" + if storage == "int8" { + col = "q" + } + var mismatches, rows int + err := sqlitex.Execute(conn, "SELECT i.embedding, r."+col+" FROM idx i JOIN ref r ON r.id = i.rowid", &sqlitex.ExecOptions{ + ResultFunc: func(stmt *sqlite.Stmt) error { + rows++ + a := make([]byte, stmt.ColumnLen(0)) + b := make([]byte, stmt.ColumnLen(1)) + stmt.ColumnBytes(0, a) + stmt.ColumnBytes(1, b) + if !bytes.Equal(a, b) { + mismatches++ + } + return nil + }, + }) + if err != nil { + t.Fatal(err) + } + if rows != len(want) || mismatches != 0 { + t.Fatalf("stored embeddings: %d rows (want %d), %d mismatches", rows, len(want), mismatches) + } + }) + } + } +} + +func TestVectorIndexErrors(t *testing.T) { + tests := []struct { + name string + setup string + query string + wantErr string + }{ + { + name: "unknown argument", + query: "CREATE VIRTUAL TABLE bad USING vector_index(float64)", + wantErr: `unknown argument "float64"`, + }, + { + name: "zero chunk_size", + query: "CREATE VIRTUAL TABLE bad USING vector_index(chunk_size=0)", + wantErr: `chunk_size must be a positive integer, got "0"`, + }, + { + name: "non-integer chunk_size", + query: "CREATE VIRTUAL TABLE bad USING vector_index(chunk_size=big)", + wantErr: `chunk_size must be a positive integer, got "big"`, + }, + { + name: "int8 without quantization range", + setup: "noquant", + query: "CREATE VIRTUAL TABLE bad USING vector_index(int8)", + wantErr: "int8 storage requires Register with WithQuantRange", + }, + { + name: "insert wrong length", + query: "INSERT INTO idx (rowid, embedding) VALUES (1, X'00000000')", + wantErr: "embedding: expected 16 bytes (dim=4), got 4", + }, + { + name: "insert NULL embedding", + query: "INSERT INTO idx (rowid, embedding) VALUES (1, NULL)", + wantErr: "embedding must be a float32 blob", + }, + { + name: "duplicate rowid", + query: "INSERT INTO idx (rowid, embedding) VALUES (0, X'00000000000000000000000000000000')", + wantErr: "UNIQUE constraint failed", + }, + { + name: "update to existing rowid keeps row", + query: "UPDATE idx SET rowid = 0 WHERE rowid = 1", + wantErr: "UNIQUE constraint failed", + }, + { + name: "MATCH without k or LIMIT", + query: "SELECT rowid FROM idx WHERE embedding MATCH X'00000000000000000000000000000000'", + wantErr: "search requires k = ? or LIMIT", + }, + { + name: "LIMIT through a join", + query: "SELECT r.id FROM idx v JOIN ref r ON r.id = v.rowid WHERE v.embedding MATCH X'00000000000000000000000000000000' LIMIT 5", + wantErr: "search requires k = ? or LIMIT", + }, + { + name: "MATCH wrong length", + query: "SELECT rowid FROM idx WHERE embedding MATCH X'0000' AND k = 5", + wantErr: "MATCH query: expected 16 bytes (dim=4), got 2", + }, + { + name: "negative k", + query: "SELECT rowid FROM idx WHERE embedding MATCH X'00000000000000000000000000000000' AND k = -1", + wantErr: "k must be non-negative", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + conn := openTestConn(t) + if tt.setup == "noquant" { + if err := Register(conn, 4); err != nil { + t.Fatal(err) + } + } else { + if err := Register(conn, 4, WithQuantRange(0, 100)); err != nil { + t.Fatal(err) + } + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index()") + execArgs(t, conn, "CREATE TABLE ref (id INTEGER PRIMARY KEY)") + execArgs(t, conn, "INSERT INTO idx (rowid, embedding) VALUES (0, X'00000000000000000000000000000000'), (1, X'0000803f0000803f0000803f0000803f')") + execArgs(t, conn, "INSERT INTO ref (id) VALUES (0)") + } + err := sqlitex.ExecuteTransient(conn, tt.query, nil) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %v, want containing %q", err, tt.wantErr) + } + if tt.setup != "noquant" { + rows := queryRows(t, conn, "SELECT rowid, 0 FROM idx ORDER BY rowid") + assertRowsEqual(t, rows, []indexRow{{0, 0}, {1, 0}}) + } + }) + } +} + +func TestVectorIndexChunkSize(t *testing.T) { + tests := []struct { + name string + dim int + args string + want int + }{ + {name: "float32 default at dim 4", dim: 4, args: "", want: 65536}, + {name: "int8 default at dim 4", dim: 4, args: "int8", want: 262144}, + {name: "float32 default at dim 1536", dim: 1536, args: "float32", want: 170}, + {name: "int8 default at dim 1536", dim: 1536, args: "int8", want: 682}, + {name: "default when vector exceeds target", dim: 300000, args: "float32", want: 1}, + {name: "explicit chunk_size", dim: 4, args: "chunk_size = 7", want: 7}, + {name: "explicit chunk_size with storage type", dim: 4, args: "int8, chunk_size=7", want: 7}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + conn := openTestConn(t) + if err := Register(conn, tt.dim, WithQuantRange(0, 100)); err != nil { + t.Fatal(err) + } + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index("+tt.args+")") + rows := queryRows(t, conn, "SELECT chunk_size, 0 FROM idx_info") + if len(rows) != 1 || rows[0].id != int64(tt.want) { + t.Fatalf("chunk_size = %+v, want %d", rows, tt.want) + } + if tt.want != 7 { + return + } + v := Float32ToBlob([]float32{1, 2, 3, 4}) + for i := range 20 { + execArgs(t, conn, "INSERT INTO idx (rowid, embedding) VALUES (?, ?)", i, v) + } + rows = queryRows(t, conn, "SELECT count(*), max(size) FROM idx_chunks") + if len(rows) != 1 || rows[0] != (indexRow{3, 7}) { + t.Fatalf("chunks (count, max size) = %+v, want (3, 7)", rows) + } + }) + } +} + +func TestVectorIndexSchema(t *testing.T) { + shadowCount := func(t *testing.T, conn *sqlite.Conn, prefix string) int { + t.Helper() + var n int + err := sqlitex.Execute(conn, "SELECT count(*) FROM sqlite_schema WHERE name LIKE ?", &sqlitex.ExecOptions{ + Args: []any{prefix + "\\_%"}, + ResultFunc: func(stmt *sqlite.Stmt) error { n = stmt.ColumnInt(0); return nil }, + }) + if err != nil { + t.Fatal(err) + } + return n + } + open := func(t *testing.T, path string, dim int, opts ...Option) *sqlite.Conn { + t.Helper() + conn, err := sqlite.OpenConn(path) + if err != nil { + t.Fatal(err) + } + if err := Register(conn, dim, opts...); err != nil { + t.Fatal(err) + } + return conn + } + query := Float32ToBlob([]float32{1, 1, 1, 1}) + + t.Run("reopen uses stored vectors", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "idx.db") + conn := open(t, path, 4, WithQuantRange(0, 100)) + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index(int8)") + execArgs(t, conn, "INSERT INTO idx (rowid, embedding) VALUES (1, ?), (2, ?)", Float32ToBlob([]float32{1, 1, 1, 1}), Float32ToBlob([]float32{9, 9, 9, 9})) + conn.Close() + + conn = open(t, path, 4, WithQuantRange(0, 100)) + defer conn.Close() + rows := queryRows(t, conn, "SELECT rowid, distance FROM idx WHERE embedding MATCH ?1 AND k = 1", query) + if len(rows) != 1 || rows[0].id != 1 { + t.Fatalf("rows = %+v, want rowid 1", rows) + } + }) + + for _, tt := range []struct { + name string + dim int + opts []Option + wantErr string + }{ + {"reopen with different dimension", 8, []Option{WithQuantRange(0, 100)}, "has dimension 4, Register was called with 8"}, + {"reopen with different quantization range", 4, []Option{WithQuantRange(-1, 1)}, "was created with quantization range [0, 100]"}, + } { + t.Run(tt.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "idx.db") + conn := open(t, path, 4, WithQuantRange(0, 100)) + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index(int8)") + conn.Close() + + conn = open(t, path, tt.dim, tt.opts...) + defer conn.Close() + err := sqlitex.ExecuteTransient(conn, "SELECT rowid FROM idx", nil) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %v, want containing %q", err, tt.wantErr) + } + }) + } + + t.Run("rename moves shadow tables", func(t *testing.T) { + conn := open(t, ":memory:", 4) + defer conn.Close() + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index()") + execArgs(t, conn, "INSERT INTO idx (rowid, embedding) VALUES (1, ?)", query) + execArgs(t, conn, "ALTER TABLE idx RENAME TO docs") + if n := shadowCount(t, conn, "idx"); n != 0 { + t.Fatalf("%d idx_ shadow objects remain", n) + } + rows := queryRows(t, conn, "SELECT rowid, distance FROM docs WHERE embedding MATCH ?1 AND k = 1", query) + if len(rows) != 1 || rows[0].id != 1 { + t.Fatalf("rows = %+v, want rowid 1", rows) + } + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index()") + }) + + t.Run("drop removes shadow tables", func(t *testing.T) { + conn := open(t, ":memory:", 4) + defer conn.Close() + execArgs(t, conn, "CREATE VIRTUAL TABLE idx USING vector_index()") + execArgs(t, conn, "DROP TABLE idx") + if n := shadowCount(t, conn, "idx"); n != 0 { + t.Fatalf("%d idx_ shadow objects remain", n) + } + }) +} diff --git a/vector_test.go b/vector_test.go index 1a21105..9b96703 100644 --- a/vector_test.go +++ b/vector_test.go @@ -171,18 +171,6 @@ func TestRegister(t *testing.T) { } }) - t.Run("stub functions return error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable, preventing SQL error propagation") - conn := openTestConn(t) - if err := Register(conn, 3); err != nil { - t.Fatal(err) - } - err := sqlitex.ExecuteTransient(conn, "SELECT vector_encode('[1,2,3]')", nil) - if err == nil { - t.Fatal("expected stub error from vector_encode, got nil") - } - }) - t.Run("re-register overwrites without error", func(t *testing.T) { conn := openTestConn(t) if err := Register(conn, 3); err != nil { @@ -231,7 +219,6 @@ func TestVectorEncode(t *testing.T) { }) t.Run("dimension mismatch", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -243,7 +230,6 @@ func TestVectorEncode(t *testing.T) { }) t.Run("invalid JSON", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -255,7 +241,6 @@ func TestVectorEncode(t *testing.T) { }) t.Run("JSON object not array", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -459,7 +444,6 @@ func TestVectorDistance(t *testing.T) { }) t.Run("wrong dimension blob error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -476,7 +460,6 @@ func TestVectorDistance(t *testing.T) { }) t.Run("quantized blob input error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -641,7 +624,6 @@ func TestVectorQuantize(t *testing.T) { }) t.Run("without WithQuantRange returns error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -654,7 +636,6 @@ func TestVectorQuantize(t *testing.T) { }) t.Run("wrong dimension input error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3, WithQuantRange(-1, 1)); err != nil { t.Fatal(err) @@ -812,7 +793,6 @@ func TestVectorDistanceQ(t *testing.T) { }) t.Run("non-quantized blob input error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3, WithQuantRange(-1, 1)); err != nil { t.Fatal(err) @@ -825,7 +805,6 @@ func TestVectorDistanceQ(t *testing.T) { }) t.Run("without WithQuantRange returns error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -1101,7 +1080,6 @@ func TestVectorEmbed(t *testing.T) { }) t.Run("without WithEmbedder returns error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) if err := Register(conn, 3); err != nil { t.Fatal(err) @@ -1113,7 +1091,6 @@ func TestVectorEmbed(t *testing.T) { }) t.Run("embedder returns wrong dimension", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) emb := &mockEmbedder{vec: []float32{0.1, 0.2}} // dim=2, registered dim=3 if err := Register(conn, 3, WithEmbedder(emb)); err != nil { @@ -1126,7 +1103,6 @@ func TestVectorEmbed(t *testing.T) { }) t.Run("embedder returns error", func(t *testing.T) { - t.Skip("blocked on zombiezen/go/sqlite fix: resultError shadows err variable") conn := openTestConn(t) emb := &mockEmbedder{err: errors.New("embedding service unavailable")} if err := Register(conn, 3, WithEmbedder(emb)); err != nil {