diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3018b3a --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.tmp/ diff --git a/CLAUDE.md b/CLAUDE.md index 79cef90..725a425 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,41 +5,14 @@ This file governs how Claude Code works in this repository. These rules are mand ## Reference Documents - **SPEC.md**: the authoritative specification. All implementation decisions come from this document. -- **PLAN.md**: the step-by-step implementation plan. Work proceeds in plan part order. - -## Plan Execution Rules - -1. **Follow PLAN.md top-to-bottom.** Do not skip parts. Do not start a later part before the current part is fully validated. - -2. **Validate before moving on.** Every plan part has a Validation section with specific commands and checks. Run every validation step. Only proceed to the next part when all validation checks pass. - -3. **If validation fails:** - - Do NOT silently fix and move on. - - Record the failure in the part's **Failure Log** section in PLAN.md. Include: - - The exact command that failed - - The full error output - - A brief description of the root cause - - What was done to fix it - - Fix the issue, then re-run all validation steps for the current part. - - Update the Failure Log with the resolution. - -4. **Update PLAN.md continuously.** As work progresses: - - Check off completed items with `[x]` - - Mark in-progress items with `[~]` - - Mark failed items with `[!]` and add details to the Failure Log - - The plan must always reflect the current true state of the implementation. ## Git and Commit Rules 1. **A git repository is already initialized.** All work happens on the default branch. -2. **Every plan part must end with a commit.** Commit even if validation initially failed — the failure and fix should both be in the history. The commit message is specified in each plan part. - -3. **No major changes without a commit first.** Before starting any significant refactor, rewrite, or new part, ensure the current state is committed. This guarantees work can always be restored. - -4. **Commit granularity:** one commit per plan part. Do not squash multiple parts into one commit. Do not split a single part across multiple commits unless a mid-part checkpoint is needed to preserve work before a risky change. +2. **No major changes without a commit first.** Before starting any significant refactor or rewrite, ensure the current state is committed. This guarantees work can always be restored. -5. **Never amend previous commits.** Always create new commits. If something from a previous part needs fixing, fix it in the current part and note it in the Failure Log. +3. **Never amend previous commits.** Always create new commits. ## Code Quality Standards diff --git a/PLAN.md b/PLAN.md deleted file mode 100644 index 22bf27a..0000000 --- a/PLAN.md +++ /dev/null @@ -1,374 +0,0 @@ -# Implementation Plan - -Reference: [SPEC.md](SPEC.md) - -Status legend: `[ ]` pending | `[~]` in progress | `[x]` complete | `[!]` failed - ---- - -## Part 1: Project Scaffolding - -### Work -- [x] Run `go mod init github.com/justintout/go-sqlite-vector` -- [x] Run `go get zombiezen.com/go/sqlite` -- [x] Create `LICENSE` file (BSD-3-Clause) -- [x] Create `vector.go` with package declaration and doc comment -- [x] Create `vector_test.go` with package declaration - -### Validation -- [x] `go build ./...` succeeds -- [x] `go vet ./...` succeeds -- [x] `go test ./...` runs (no tests yet, but no errors) -- [x] `go mod tidy` produces no changes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "scaffold project with go module and license" - ---- - -## Part 2: Blob Encoding (Float32ToBlob, BlobToFloat32) - -### Work -- [x] Implement `Float32ToBlob(v []float32) []byte` in `vector.go` - - Little-endian encoding via `encoding/binary` -- [x] Implement `BlobToFloat32(b []byte) ([]float32, error)` in `vector.go` - - Return error if `len(b) % 4 != 0` - - Little-endian decoding via `encoding/binary` - -### Validation -- [x] Write table-driven tests in `vector_test.go`: - - Round-trip: `BlobToFloat32(Float32ToBlob(v))` == `v` for several vectors - - `BlobToFloat32` with invalid length returns error - - Empty slice round-trips correctly - - Known byte values: `Float32ToBlob([]float32{1.0})` == `[]byte{0x00, 0x00, 0x80, 0x3f}` -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "implement Float32ToBlob and BlobToFloat32" - ---- - -## Part 3: Config and Registration Skeleton - -### Work -- [x] Define unexported `config` struct with fields: - - `dim int` - - `quantMin, quantMax float32` - - `quantEnabled bool` -- [x] Define `type Option func(*config)` -- [x] Implement `WithQuantRange(min, max float32) Option` -- [x] Implement `Register(conn *sqlite.Conn, dim int, opts ...Option) error` - - Validate `dim >= 1` - - Apply options to config - - Register all four SQL function stubs via `conn.CreateFunction` that return errors ("not implemented") - - Each SQL function receives the config via closure - -### Validation -- [x] Write test: `Register` with `dim=0` returns error -- [x] Write test: `Register` with `dim=3` succeeds -- [x] Write test: `Register` with `WithQuantRange(-1, 1)` succeeds and marks quant enabled -- [x] Write test: calling any SQL function stub returns a "not implemented" or placeholder error (skipped: blocked on upstream library bug) -- [x] Write test: re-calling `Register` on the same connection does not error (overwrite behavior) -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -- **Stub error propagation test**: zombiezen.com/go/sqlite v1.4.2 has a bug in `func.go:resultError` where `err` is shadowed by `:=` on line 167, causing `sqlite3_result_error_code(ctx, SQLITE_OK)` to cancel the error. All Scalar function error returns are silently swallowed. Test skipped with `t.Skip()`. Upstream patch is in progress. - -### Commit -- [x] Commit: "add config, options, and Register skeleton" - ---- - -## Part 4: vector_encode SQL Function - -### Work -- [x] Implement `vector_encode` function body: - - Check for NULL input, return NULL - - Read text argument as string - - Parse JSON via `encoding/json` into `[]float64` (JSON numbers are float64) - - Validate array length == `dim`; return SQL error with expected vs actual if mismatch - - Convert `[]float64` to `[]float32` - - Use `Float32ToBlob` to produce the result blob - - Return blob result - -### Validation -- [x] Table-driven integration tests (open real SQLite conn, register, run SQL): - - `SELECT vector_encode('[1.0, 2.0, 3.0]')` with dim=3 returns correct 12-byte blob - - `SELECT vector_encode('[1.0, 2.0]')` with dim=3 returns SQL error (skipped: upstream bug) - - `SELECT vector_encode('not json')` returns SQL error (skipped: upstream bug) - - `SELECT vector_encode('{}')` returns SQL error (skipped: upstream bug) - - `SELECT vector_encode(NULL)` returns NULL - - Round-trip: `BlobToFloat32` on the result of `vector_encode` matches input values -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -- Error propagation tests skipped due to upstream zombiezen/go/sqlite resultError bug (same as Part 3). Error-returning code paths are implemented correctly; tests will pass once the library is patched. - -### Commit -- [x] Commit: "implement vector_encode SQL function" - ---- - -## Part 5: L2 Squared Distance (internal) - -### Work -- [x] Implement unexported `l2Squared(a, b []float32) float64` - - Compute `sum((a[i] - b[i])^2)` using float64 accumulator for precision -- [x] Implement unexported `isQuantizedBlob(b []byte) bool` - - Returns true if `len(b) >= 2 && b[0] == 0x00 && b[1] == 0x01` - -### Validation -- [x] Unit tests for `l2Squared`: - - Identical vectors → 0.0 - - Known values: `l2Squared([1,0,0], [0,1,0])` == 2.0 - - Known values: `l2Squared([1,2,3], [4,5,6])` == 27.0 - - Single dimension: `l2Squared([3], [7])` == 16.0 -- [x] Unit tests for `isQuantizedBlob`: - - `[]byte{0x00, 0x01, ...}` → true - - `[]byte{0x00, 0x00, ...}` → false - - Float32 blob → false - - Empty/short slices → false -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "implement l2Squared distance and blob format detection" - ---- - -## Part 6: vector_distance SQL Function - -### Work -- [x] Implement `vector_distance` function body: - - Check for NULL inputs, return NULL - - Read both arguments as blobs ([]byte) - - Validate neither blob has quantized magic bytes; return SQL error if detected - - Validate both blobs have byte length == `dim * 4`; return SQL error with expected vs actual - - Decode both blobs to `[]float32` via `BlobToFloat32` - - Compute `l2Squared` and return as float64 result - -### Validation -- [x] Table-driven integration tests: - - Distance of identical vectors == 0.0 - - Distance of known vectors matches hand-computed value - - `SELECT vector_distance(vector_encode('[1,2,3]'), vector_encode('[4,5,6]'))` == 27.0 (dim=3) - - Wrong dimension blob → SQL error (skipped: upstream bug) - - Quantized blob input → SQL error mentioning `vector_distance_q` (skipped: upstream bug) - - NULL input → NULL output - - Two NULLs → NULL - - One NULL, one valid → NULL -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -- Error propagation tests skipped due to upstream zombiezen/go/sqlite resultError bug (same as Parts 3-4). - -### Commit -- [x] Commit: "implement vector_distance SQL function" - ---- - -## Part 7: Quantization (internal) - -### Work -- [x] Implement unexported `quantize(v []float32, min, max float32) []byte` - - Allocate `2 + len(v)` bytes - - Write magic bytes `0x00, 0x01` - - For each float: `clamp(round((v - min) / (max - min) * 255 - 128), -128, 127)` - - Write as signed int8 -- [x] Implement unexported `dequantize(b []byte, min, max float32) ([]float32, error)` - - Validate magic bytes; return error if missing - - For each int8: `(q + 128) / 255 * (max - min) + min` - - Return `[]float32` - -### Validation -- [x] Unit tests for `quantize`: - - Known values with range [-1, 1]: value 0.0 → int8 value near 0 (middle of range) - - Boundary: value == min → -128, value == max → 127 - - Out-of-range clamping: value > max → 127, value < min → -128 - - Output blob starts with `0x00, 0x01` - - Output blob length == `2 + dim` -- [x] Unit tests for `dequantize`: - - Round-trip: `dequantize(quantize(v))` is approximately equal to `v` (within int8 precision) - - Missing magic bytes → error - - Correct output length == input length - 2 -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "implement scalar int8 quantization and dequantization" - ---- - -## Part 8: vector_quantize SQL Function - -### Work -- [x] Implement `vector_quantize` function body: - - Check for NULL input, return NULL - - Check `quantEnabled`; return SQL error if false - - Read blob argument - - Validate byte length == `dim * 4` - - Decode to `[]float32` via `BlobToFloat32` - - Call `quantize(floats, config.quantMin, config.quantMax)` - - Return quantized blob - -### Validation -- [x] Table-driven integration tests: - - `SELECT vector_quantize(vector_encode('[0.5, -0.5, 0.0]'))` with dim=3 and range [-1,1] produces correct blob (starts with 0x00, 0x01, length 5) - - Calling without `WithQuantRange` → SQL error (skipped: upstream bug) - - Wrong dimension input → SQL error (skipped: upstream bug) - - NULL input → NULL - - Out-of-range values are clamped (no error) -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -- Error propagation tests skipped due to upstream zombiezen/go/sqlite resultError bug. - -### Commit -- [x] Commit: "implement vector_quantize SQL function" - ---- - -## Part 9: vector_distance_q SQL Function - -### Work -- [x] Implement `vector_distance_q` function body: - - Check for NULL inputs, return NULL - - Check `quantEnabled`; return SQL error if false - - Read both arguments as blobs - - Validate both blobs have quantized magic bytes; return SQL error if missing - - Validate both blobs have byte length == `2 + dim` - - Dequantize both via `dequantize(blob, config.quantMin, config.quantMax)` - - Compute `l2Squared` on the dequantized float32 slices - - Return float64 result - -### Validation -- [x] Table-driven integration tests: - - Quantized distance of identical vectors ≈ 0.0 (not exact due to quantization loss) - - Quantized distance ordering matches float32 distance ordering for a set of test vectors - - Non-quantized blob input → SQL error (skipped: upstream bug) - - Calling without `WithQuantRange` → SQL error (skipped: upstream bug) - - Wrong quantized blob dimension → SQL error (skipped: upstream bug) - - NULL inputs → NULL -- [x] `go test ./...` passes -- [x] `go vet ./...` passes - -### Failure Log -- Error propagation tests skipped due to upstream zombiezen/go/sqlite resultError bug. - -### Commit -- [x] Commit: "implement vector_distance_q SQL function" - ---- - -## Part 10: End-to-End Integration Tests - -### Work -- [x] Write full pipeline integration tests that mirror real usage: - - Create table, insert rows with `vector_encode`, query with `vector_distance` + `ORDER BY` + `LIMIT` - - Verify k-NN ordering is correct for a known dataset - - Create table with quantized embeddings, insert with `vector_quantize(vector_encode(...))`, query with `vector_distance_q` - - Verify quantized k-NN ordering matches float32 k-NN ordering - - Test `Register` overwrite: register dim=3, then re-register dim=4, verify dim=4 functions work -- [x] Write test for Go-level `Float32ToBlob` used with SQL parameter binding (bind blob, then compute distance) - -### Validation -- [x] `go test ./...` passes (all tests including new e2e tests) -- [x] `go vet ./...` passes -- [x] `go test -count=1 ./...` passes (no caching) -- [x] `go test -race ./...` passes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "add end-to-end integration tests" - ---- - -## Part 11: Benchmarks - -### Work -- [x] Create `bench_test.go` -- [x] Implement benchmarks using pre-allocated random vectors: - - `BenchmarkL2Distance/dim=384` - - `BenchmarkL2Distance/dim=768` - - `BenchmarkL2Distance/dim=1536` - - `BenchmarkQuantize/dim=384` - - `BenchmarkQuantize/dim=768` - - `BenchmarkQuantize/dim=1536` - - `BenchmarkDequantize/dim=384` - - `BenchmarkDequantize/dim=768` - - `BenchmarkDequantize/dim=1536` - - `BenchmarkVectorEncode/dim=384` - - `BenchmarkVectorEncode/dim=768` - - `BenchmarkVectorEncode/dim=1536` -- [x] Each benchmark uses `b.ResetTimer()` after setup and `b.ReportAllocs()` - -### Validation -- [x] `go test -bench=. -benchmem ./...` runs without errors -- [x] All benchmarks produce non-zero ns/op values -- [x] `go test ./...` still passes (benchmarks don't break tests) -- [x] `go vet ./...` passes - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "add benchmarks for core operations" - ---- - -## Part 12: Final Review and Cleanup - -### Work -- [x] Run `go mod tidy` to clean dependencies -- [x] Verify all exported symbols have doc comments -- [x] Verify no unexported symbols are accidentally exported -- [x] Verify error message strings match SPEC.md table exactly -- [x] Review all SQL functions for correct NULL handling -- [x] Review all SQL functions for correct error propagation (error returns are correct; propagation blocked by upstream library bug) -- [x] Confirm `Register` overwrite behavior works end-to-end - -### Validation -- [x] `go build ./...` succeeds -- [x] `go vet ./...` passes -- [x] `go test ./...` passes -- [x] `go test -race ./...` passes -- [x] `go test -bench=. -benchmem ./...` runs -- [x] `go mod tidy` produces no changes -- [x] No `TODO`, `FIXME`, or `HACK` comments remain in source - -### Failure Log -_(no failures)_ - -### Commit -- [x] Commit: "final review and cleanup" - ---- - -## Completion Checklist - -- [x] All 12 parts have status `[x]` -- [x] All tests pass: `go test -race -count=1 ./...` -- [x] All benchmarks run: `go test -bench=. -benchmem ./...` -- [x] `go vet` clean -- [x] `go mod tidy` clean -- [x] Every part has a commit -- [x] PLAN.md is fully updated with final statuses diff --git a/README.md b/README.md index cb1203b..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. @@ -177,20 +230,55 @@ Each row returned by `vector_chunk` has a `value` (the chunk text) and a `chunk_ go test -bench=. -benchmem ./... ``` -Results on Apple M3 Max: +Results on Apple M3 Max. Distance benchmarks operate on encoded blobs, as the SQL functions do: + +``` +BenchmarkL2Distance/dim=384 4908330 241.5 ns/op 0 B/op 0 allocs/op +BenchmarkL2Distance/dim=768 2496775 483.3 ns/op 0 B/op 0 allocs/op +BenchmarkL2Distance/dim=1536 1246688 965.9 ns/op 0 B/op 0 allocs/op +BenchmarkQuantize/dim=384 2521759 477.5 ns/op 416 B/op 1 allocs/op +BenchmarkQuantize/dim=768 1269049 944.5 ns/op 896 B/op 1 allocs/op +BenchmarkQuantize/dim=1536 630952 1889 ns/op 1792 B/op 1 allocs/op +BenchmarkL2DistanceQuantized/dim=384 11420806 104.5 ns/op 0 B/op 0 allocs/op +BenchmarkL2DistanceQuantized/dim=768 5961984 201.5 ns/op 0 B/op 0 allocs/op +BenchmarkL2DistanceQuantized/dim=1536 3041773 398.2 ns/op 0 B/op 0 allocs/op +``` + +### 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 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. ``` -BenchmarkL2Distance/dim=384 10429109 115.1 ns/op 0 B/op 0 allocs/op -BenchmarkL2Distance/dim=768 5440725 221.4 ns/op 0 B/op 0 allocs/op -BenchmarkL2Distance/dim=1536 2768809 434.3 ns/op 0 B/op 0 allocs/op -BenchmarkQuantize/dim=384 2386230 504.5 ns/op 416 B/op 1 allocs/op -BenchmarkQuantize/dim=768 1000000 1001 ns/op 896 B/op 1 allocs/op -BenchmarkQuantize/dim=1536 606144 1974 ns/op 1792 B/op 1 allocs/op -BenchmarkDequantize/dim=384 2781045 439.3 ns/op 1536 B/op 1 allocs/op -BenchmarkDequantize/dim=768 1607391 739.0 ns/op 3072 B/op 1 allocs/op -BenchmarkDequantize/dim=1536 788222 1473 ns/op 6144 B/op 1 allocs/op +curl -O ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz +mkdir -p .tmp && tar xzf sift.tar.gz -C .tmp +go test -tags sift -run TestSIFT1M -timeout 0 -v -sift.dir .tmp/sift +uv run bench/sift_compare.py .tmp/sift 100 ``` +Results on Apple M3 Max, macOS 15.7, Go 1.24.12. Build time is the time to insert all vectors (HNSW builds use all cores). + +| Implementation | Search | Build | p50 | QPS | Recall@10 | Recall@100 | +|---|---|---|---|---|---|---| +| 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 | +| FAISS 1.15 HNSW M=16 ef=512 | approximate | 28.7s | 0.91 ms | 1,121 | 0.999 | 0.999 | +| hnswlib 0.8 M=16 ef=512 | approximate | 46.8s | 1.22 ms | 852 | 0.999 | 0.998 | + +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. Set the size to at least the database file size: + +```sql +PRAGMA mmap_size = 4294967296; -- 4 GiB +``` + +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 BSD-3-Clause diff --git a/SPEC.md b/SPEC.md index f896ef6..c1f663c 100644 --- a/SPEC.md +++ b/SPEC.md @@ -100,12 +100,86 @@ Converts a float32 blob to a scalar int8 quantized blob using the globally confi vector_distance_q(a BLOB, b BLOB) -> REAL ``` -Computes squared L2 distance between two quantized int8 blobs. Both blobs are dequantized back to float32 using the configured min/max range before computing the distance. +Computes squared L2 distance between the dequantized forms of two quantized int8 blobs, using the configured min/max range. Because dequantization is affine, the distance is computed directly from the int8 values as `sum((qa - qb)^2) * ((max - min) / 255)^2` without materializing float32 vectors. - **Input**: two quantized int8 blobs. - **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 @@ -249,11 +324,10 @@ Go benchmarks (`testing.B`) for the core operations at dimensions 384, 768, and - `BenchmarkL2Distance_{384,768,1536}` - `BenchmarkQuantize_{384,768,1536}` -- `BenchmarkDequantize_{384,768,1536}` +- `BenchmarkL2DistanceQuantized_{384,768,1536}` - `BenchmarkVectorEncode_{384,768,1536}` ## 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/bench/sift_compare.py b/bench/sift_compare.py new file mode 100644 index 0000000..c3f7744 --- /dev/null +++ b/bench/sift_compare.py @@ -0,0 +1,128 @@ +# /// script +# requires-python = ">=3.12,<3.13" +# dependencies = ["numpy", "faiss-cpu", "hnswlib", "sqlite-vec"] +# /// +"""Run SIFT1M through well-known vector search implementations for comparison +with TestSIFT1M in sift_test.go. Queries run single-threaded; index builds may +use all cores. + + uv run bench/sift_compare.py .tmp/sift 100 +""" + +import os +import sqlite3 +import sys +import tempfile +import time + +import faiss +import hnswlib +import numpy as np +import sqlite_vec + +K = 100 + + +def read_vecs(path, dtype): + raw = np.fromfile(path, dtype=dtype) + dim = raw[:1].view(np.int32)[0] + return raw.reshape(-1, dim + 1)[:, 1:].copy() + + +def report(name, build, search): + lat, hits1, hits10, hits100 = [], 0.0, 0.0, 0.0 + for q, t in zip(queries, truth): + start = time.perf_counter() + ids = search(q) + lat.append(time.perf_counter() - start) + ids = list(ids) + hits1 += len(set(ids[:1]) & set(t[:1])) + hits10 += len(set(ids[:10]) & set(t[:10])) / 10 + hits100 += len(set(ids[:100]) & set(t[:100])) / 100 + lat = np.array(lat) + n = len(queries) + print( + f"{name:<28} {build:>9.1f}s {np.percentile(lat, 50) * 1e3:>9.2f}ms " + f"{np.percentile(lat, 99) * 1e3:>9.2f}ms {n / lat.sum():>10.1f} " + f"{hits1 / n:>7.4f} {hits10 / n:>7.4f} {hits100 / n:>7.4f}", + flush=True, + ) + + +def timed(fn): + start = time.perf_counter() + result = fn() + return result, time.perf_counter() - start + + +sift_dir, nq = sys.argv[1], int(sys.argv[2]) +base = read_vecs(os.path.join(sift_dir, "sift_base.fvecs"), np.float32) +queries = read_vecs(os.path.join(sift_dir, "sift_query.fvecs"), np.float32)[:nq] +truth = read_vecs(os.path.join(sift_dir, "sift_groundtruth.ivecs"), np.int32)[:nq] +dim = base.shape[1] + +print(f"{'method':<28} {'build':>10} {'p50':>11} {'p99':>11} {'QPS':>10} {'R@1':>7} {'R@10':>7} {'R@100':>7}") + +# FAISS exact (BLAS brute force). +faiss.omp_set_num_threads(1) +flat = faiss.IndexFlatL2(dim) +_, build = timed(lambda: flat.add(base)) +report("faiss IndexFlatL2", build, lambda q: flat.search(q[None], K)[1][0]) +del flat + +# FAISS HNSW. +faiss.omp_set_num_threads(os.cpu_count()) +hnsw = faiss.IndexHNSWFlat(dim, 16) +hnsw.hnsw.efConstruction = 200 +_, build = timed(lambda: hnsw.add(base)) +faiss.omp_set_num_threads(1) +for ef in (128, 512): + hnsw.hnsw.efSearch = ef + report(f"faiss HNSW M=16 ef={ef}", build, lambda q: hnsw.search(q[None], K)[1][0]) +del hnsw + +# hnswlib. +hl = hnswlib.Index(space="l2", dim=dim) +hl.init_index(max_elements=len(base), M=16, ef_construction=200) +hl.set_num_threads(os.cpu_count()) +_, build = timed(lambda: hl.add_items(base, np.arange(len(base)))) +hl.set_num_threads(1) +for ef in (128, 512): + hl.set_ef(ef) + report(f"hnswlib M=16 ef={ef}", build, lambda q: hl.knn_query(q, k=K)[0][0]) +del hl + +# sqlite-vec, both as a plain-table scan with a scalar distance function (the +# same query shape as this library) and via its vec0 virtual table. +with tempfile.TemporaryDirectory() as tmp: + db = sqlite3.connect(os.path.join(tmp, "sift.db")) + db.enable_load_extension(True) + sqlite_vec.load(db) + + def load_plain(): + db.execute("CREATE TABLE sift (id INTEGER PRIMARY KEY, embedding BLOB)") + db.executemany("INSERT INTO sift VALUES (?, ?)", ((i, v.tobytes()) for i, v in enumerate(base))) + db.commit() + + _, build = timed(load_plain) + report( + "sqlite-vec scalar scan", + build, + lambda q: [r[0] for r in db.execute( + "SELECT id FROM sift ORDER BY vec_distance_l2(embedding, ?) LIMIT 100", (q.tobytes(),))], + ) + + def load_vec0(): + db.execute(f"CREATE VIRTUAL TABLE sift_vec0 USING vec0(embedding float[{dim}])") + db.executemany( + "INSERT INTO sift_vec0 (rowid, embedding) VALUES (?, ?)", ((i, v.tobytes()) for i, v in enumerate(base))) + db.commit() + + _, build = timed(load_vec0) + report( + "sqlite-vec vec0", + build, + lambda q: [r[0] for r in db.execute( + "SELECT rowid FROM sift_vec0 WHERE embedding MATCH ? AND k = 100", (q.tobytes(),))], + ) + db.close() diff --git a/bench_test.go b/bench_test.go index 7aa1be1..aa71447 100644 --- a/bench_test.go +++ b/bench_test.go @@ -16,8 +16,8 @@ func randomFloat32s(n int) []float32 { func BenchmarkL2Distance(b *testing.B) { for _, dim := range []int{384, 768, 1536} { - a := randomFloat32s(dim) - c := randomFloat32s(dim) + a := Float32ToBlob(randomFloat32s(dim)) + c := Float32ToBlob(randomFloat32s(dim)) b.Run("dim="+itoa(dim), func(b *testing.B) { b.ReportAllocs() b.ResetTimer() @@ -41,15 +41,15 @@ func BenchmarkQuantize(b *testing.B) { } } -func BenchmarkDequantize(b *testing.B) { +func BenchmarkL2DistanceQuantized(b *testing.B) { for _, dim := range []int{384, 768, 1536} { - v := randomFloat32s(dim) - qblob := quantize(v, -1.0, 1.0) + qa := quantize(randomFloat32s(dim), -1.0, 1.0) + qc := quantize(randomFloat32s(dim), -1.0, 1.0) b.Run("dim="+itoa(dim), func(b *testing.B) { b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { - dequantize(qblob, -1.0, 1.0) + l2SquaredQuantized(qa, qc, -1.0, 1.0) } }) } 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 new file mode 100644 index 0000000..317e53a --- /dev/null +++ b/sift_test.go @@ -0,0 +1,199 @@ +//go:build sift + +package vector + +import ( + "encoding/binary" + "flag" + "math" + "os" + "path/filepath" + "slices" + "testing" + "time" + + "zombiezen.com/go/sqlite" + "zombiezen.com/go/sqlite/sqlitex" +) + +var ( + siftDir = flag.String("sift.dir", "", "directory containing sift_base.fvecs, sift_query.fvecs, sift_groundtruth.ivecs") + siftQueries = flag.Int("sift.queries", 100, "number of queries to run from sift_query.fvecs") +) + +// TestSIFT1M runs exact and quantized k-NN search over the TEXMEX SIFT1M +// dataset and reports build time, single-threaded query latency, and recall +// against the published ground truth. +// +// go test -tags sift -run TestSIFT1M -timeout 0 -v -sift.dir .tmp/sift +func TestSIFT1M(t *testing.T) { + if *siftDir == "" { + t.Fatal("-sift.dir is required") + } + const dim, k = 128, 100 + + base := readVecs(t, filepath.Join(*siftDir, "sift_base.fvecs"), dim) + queries := readVecs(t, filepath.Join(*siftDir, "sift_query.fvecs"), dim) + truth := readVecs(t, filepath.Join(*siftDir, "sift_groundtruth.ivecs"), k) + nq := min(*siftQueries, len(queries)) + + lo, hi := float32(math.Inf(1)), float32(math.Inf(-1)) + for _, v := range base { + lo = min(lo, slices.Min(v)) + hi = max(hi, slices.Max(v)) + } + + dbPath := filepath.Join(t.TempDir(), "sift.db") + conn, err := sqlite.OpenConn(dbPath) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if err := Register(conn, dim, WithQuantRange(lo, hi)); err != nil { + t.Fatal(err) + } + exec(t, conn, "CREATE TABLE sift (id INTEGER PRIMARY KEY, embedding BLOB, embedding_q BLOB)") + + start := time.Now() + exec(t, conn, "BEGIN") + stmt := conn.Prep("INSERT INTO sift (id, embedding) VALUES (?1, ?2)") + for i, v := range base { + stmt.BindInt64(1, int64(i)) + stmt.BindBytes(2, Float32ToBlob(v)) + if _, err := stmt.Step(); err != nil { + t.Fatal(err) + } + stmt.Reset() + } + exec(t, conn, "COMMIT") + insertTime := time.Since(start) + + start = time.Now() + 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) + } + + 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("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 + query string + arg func([]float32) []byte + }{ + {"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 + stmt := conn.Prep(m.query) + for qi := range nq { + ids := make([]int32, 0, k) + start := time.Now() + stmt.BindBytes(1, m.arg(queries[qi])) + for { + row, err := stmt.Step() + if err != nil { + t.Fatal(err) + } + if !row { + break + } + ids = append(ids, int32(stmt.ColumnInt64(0))) + } + stmt.Reset() + lat[qi] = time.Since(start) + r1 += recall(ids, truth[qi], 1) + r10 += recall(ids, truth[qi], 10) + r100 += recall(ids, truth[qi], 100) + } + var total time.Duration + for _, d := range lat { + total += d + } + slices.Sort(lat) + n := float64(nq) + 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) + } +} + +// readVecs reads a TEXMEX .fvecs or .ivecs file. Each record is a +// little-endian int32 dimension followed by that many float32 or int32 +// values. ivecs values are converted to float32 so both formats share one +// return type; SIFT1M ids fit exactly in float32. +func readVecs(t *testing.T, path string, dim int) [][]float32 { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + isInt := filepath.Ext(path) == ".ivecs" + rec := 4 + dim*4 + if len(data)%rec != 0 { + t.Fatalf("%s: size %d is not a multiple of record size %d", path, len(data), rec) + } + out := make([][]float32, len(data)/rec) + for i := range out { + r := data[i*rec : (i+1)*rec] + if d := int(binary.LittleEndian.Uint32(r)); d != dim { + t.Fatalf("%s: record %d has dimension %d, want %d", path, i, d, dim) + } + v := make([]float32, dim) + for j := range v { + bits := binary.LittleEndian.Uint32(r[4+j*4:]) + if isInt { + v[j] = float32(int32(bits)) + } else { + v[j] = math.Float32frombits(bits) + } + } + out[i] = v + } + return out +} + +// recall returns the fraction of the true top-k ids present in the first k +// results. +func recall(got []int32, truth []float32, k int) float64 { + want := make(map[int32]bool, k) + for _, id := range truth[:k] { + want[int32(id)] = true + } + hits := 0 + for _, id := range got[:min(k, len(got))] { + if want[id] { + hits++ + } + } + return float64(hits) / float64(k) +} + +func exec(t *testing.T, conn *sqlite.Conn, query string) { + t.Helper() + if err := sqlitex.ExecuteTransient(conn, query, nil); err != nil { + t.Fatalf("%s: %v", query, err) + } +} diff --git a/vector.go b/vector.go index 2a54501..0974bde 100644 --- a/vector.go +++ b/vector.go @@ -115,9 +115,7 @@ func Register(conn *sqlite.Conn, dim int, opts ...Option) error { if len(blobB) != expected { return sqlite.Value{}, fmt.Errorf("vector_distance: expected %d bytes (dim=%d), got %d", expected, cfg.dim, len(blobB)) } - a, _ := BlobToFloat32(blobA) - b, _ := BlobToFloat32(blobB) - return sqlite.FloatValue(l2Squared(a, b)), nil + return sqlite.FloatValue(l2Squared(blobA, blobB)), nil }, }) if err != nil { @@ -172,9 +170,7 @@ func Register(conn *sqlite.Conn, dim int, opts ...Option) error { if len(blobB) != expected { return sqlite.Value{}, fmt.Errorf("vector_distance_q: expected %d bytes (dim=%d), got %d", expected, cfg.dim, len(blobB)) } - a, _ := dequantize(blobA, cfg.quantMin, cfg.quantMax) - b, _ := dequantize(blobB, cfg.quantMin, cfg.quantMax) - return sqlite.FloatValue(l2Squared(a, b)), nil + return sqlite.FloatValue(l2SquaredQuantized(blobA, blobB, cfg.quantMin, cfg.quantMax)), nil }, }) if err != nil { @@ -206,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{ @@ -243,10 +247,46 @@ func BlobToFloat32(b []byte) ([]float32, error) { return v, nil } -func l2Squared(a, b []float32) float64 { +// l2Squared computes squared L2 distance directly on two float32 blobs of equal +// length. Decoding in place avoids allocating two []float32 per row scanned, +// and unrolling the loop reduces bounds checks; both dominate scan time. +func l2Squared(a, b []byte) float64 { + b = b[:len(a)] var sum float64 + i := 0 + for ; i+16 <= len(a); i += 16 { + d0 := f32At(a, i) - f32At(b, i) + d1 := f32At(a, i+4) - f32At(b, i+4) + d2 := f32At(a, i+8) - f32At(b, i+8) + d3 := f32At(a, i+12) - f32At(b, i+12) + sum += d0*d0 + d1*d1 + d2*d2 + d3*d3 + } + for ; i+4 <= len(a); i += 4 { + d := f32At(a, i) - f32At(b, i) + sum += d * d + } + return sum +} + +func f32At(b []byte, i int) float64 { + return float64(math.Float32frombits(binary.LittleEndian.Uint32(b[i:]))) +} + +// l2SquaredQuantized computes the squared L2 distance between the dequantized +// forms of two quantized blobs of equal length without dequantizing. +// Dequantization is affine, so each component difference is +// (qa - qb) * (max - min) / 255. +func l2SquaredQuantized(a, b []byte, min, max float32) float64 { + 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 := float64(a[i]) - float64(b[i]) + d := int64(int8(a[i])) - int64(int8(b[i])) sum += d * d } return sum @@ -274,20 +314,6 @@ func quantize(v []float32, min, max float32) []byte { return b } -func dequantize(b []byte, min, max float32) ([]float32, error) { - if len(b) < 2 || b[0] != 0x00 || b[1] != 0x01 { - return nil, fmt.Errorf("dequantize: missing quantized format magic bytes") - } - data := b[2:] - r := float64(max - min) - v := make([]float32, len(data)) - for i, raw := range data { - q := int8(raw) - v[i] = float32((float64(q)+128)/255*r + float64(min)) - } - return v, nil -} - const chunkColValue = 0 const chunkColIndex = 1 const chunkColText = 2 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 ad96091..9b96703 100644 --- a/vector_test.go +++ b/vector_test.go @@ -5,6 +5,7 @@ import ( "context" "errors" "io" + "math" "testing" "zombiezen.com/go/sqlite" @@ -170,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 { @@ -230,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) @@ -242,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) @@ -254,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) @@ -348,10 +334,16 @@ func TestL2Squared(t *testing.T) { b: []float32{7}, want: 16.0, }, + { + name: "unrolled block plus remainder", + a: []float32{1, 2, 3, 4, 5}, + b: []float32{2, 4, 6, 8, 10}, + want: 55.0, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := l2Squared(tt.a, tt.b) + got := l2Squared(Float32ToBlob(tt.a), Float32ToBlob(tt.b)) if got != tt.want { t.Errorf("l2Squared(%v, %v) = %v, want %v", tt.a, tt.b, got, tt.want) } @@ -452,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) @@ -469,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) @@ -561,46 +551,45 @@ func TestQuantize(t *testing.T) { }) } -func TestDequantize(t *testing.T) { - t.Run("round-trip approximate equality", func(t *testing.T) { - original := []float32{0.5, -0.3, 0.0, 1.0, -1.0} - qblob := quantize(original, -1.0, 1.0) - got, err := dequantize(qblob, -1.0, 1.0) - if err != nil { - t.Fatal(err) - } - if len(got) != len(original) { - t.Fatalf("dequantize length = %d, want %d", len(got), len(original)) - } - for i := range got { - diff := got[i] - original[i] - if diff < 0 { - diff = -diff - } - // int8 precision: max error is (max-min)/255 ≈ 0.0078 for range [-1,1] - if diff > 0.01 { - t.Errorf("round-trip[%d]: got %v, want ~%v (diff=%v)", i, got[i], original[i], diff) +func TestL2SquaredQuantized(t *testing.T) { + // With range [0, 255] each integer value quantizes to itself minus 128, + // so the quantized distance equals the float32 distance exactly. + tests := []struct { + name string + a, b []float32 + min, max float32 + want float64 + }{ + { + name: "identical vectors", + a: []float32{0, 10, 255}, + b: []float32{0, 10, 255}, + min: 0, max: 255, + want: 0.0, + }, + { + name: "integer values in unit-step range", + a: []float32{0, 10, 255}, + b: []float32{3, 14, 255}, + min: 0, max: 255, + want: 25.0, + }, + { + name: "scaled by range", + a: []float32{0, 0}, + b: []float32{510, 510}, + min: 0, max: 510, + want: 2 * 255 * 255 * 4, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := l2SquaredQuantized(quantize(tt.a, tt.min, tt.max), quantize(tt.b, tt.min, tt.max), tt.min, tt.max) + if math.Abs(got-tt.want) > 1e-9*max(1, tt.want) { + t.Errorf("l2SquaredQuantized = %v, want %v", got, tt.want) } - } - }) - - t.Run("missing magic bytes error", func(t *testing.T) { - _, err := dequantize([]byte{0x01, 0x02, 0x03}, -1.0, 1.0) - if err == nil { - t.Fatal("expected error for missing magic bytes") - } - }) - - t.Run("correct output length", func(t *testing.T) { - qblob := quantize([]float32{0.1, 0.2, 0.3}, -1.0, 1.0) - got, err := dequantize(qblob, -1.0, 1.0) - if err != nil { - t.Fatal(err) - } - if len(got) != 3 { - t.Errorf("output length = %d, want 3", len(got)) - } - }) + }) + } } func TestVectorQuantize(t *testing.T) { @@ -635,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) @@ -648,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) @@ -806,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) @@ -819,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) @@ -1095,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) @@ -1107,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 { @@ -1120,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 {