diff --git a/doc/CONFIG_REFERENCE.md b/doc/CONFIG_REFERENCE.md index 4bbb1c3..8cd09ca 100644 --- a/doc/CONFIG_REFERENCE.md +++ b/doc/CONFIG_REFERENCE.md @@ -37,6 +37,7 @@ the DB. | `CIX_CHROMA_PERSIST_DIR` | `/data/chroma` | Legacy chromem-go store. Read on startup for the one-time import into the SQLite vector store, then left untouched as the rollback path. See [VECTORSTORE.md](VECTORSTORE.md). | | `CIX_VECTORS_DIR` | sibling of `CIX_CHROMA_PERSIST_DIR` (`/data/vectors`) | Vector store directory: one SQLite database per embedding namespace. | | `CIX_VECTOR_MMAP_SIZE` | `0` (off) | `PRAGMA mmap_size` for the vector store, in bytes. Roughly 40% lower search latency in exchange for resident memory — mapped database pages count in RSS. | +| `CIX_VECTOR_SCAN_QUANT` | `true` | Scan a compact int8 copy of each vector instead of the float32 original, rescoring the shortlist against the originals. Every score returned is the exact cosine; which documents reach the shortlist is an approximation, measured at recall 1.000 against exact search (see `doc/VECTORSTORE.md`). 3.4x fewer bytes read per query at 2048 dimensions, in exchange for roughly a quarter more disk. Set `false` if the volume cannot take it; existing copies are then ignored, and writes remove the copies of rows they touch so re-enabling rebuilds instead of trusting stale data. | | `CIX_GGUF_CACHE_DIR` | `/data/models` | Where downloaded GGUF files live. | | `CIX_PUBLIC_URL` | — | Externally-reachable URL used to build GitHub webhook delivery URLs. Empty disables webhook URL display. | diff --git a/doc/VECTORSTORE.md b/doc/VECTORSTORE.md index bc95c45..a37a9d9 100644 --- a/doc/VECTORSTORE.md +++ b/doc/VECTORSTORE.md @@ -150,21 +150,96 @@ on disk, deliberately: it keeps the package self-contained and live in `vectors`. A multi-kilobyte `TEXT` column pushes a row past SQLite's local-payload limit, and SQLite then keeps only ~1 kB of the row in the table page and spills the rest — *including the embedding* — into an overflow chain, -roughly doubling the pages a scan touches. Kept apart, a `vectors` row is -~3.2 kB and two of them share an 8 KiB page. Content is read only for the K +roughly doubling the pages a scan touches. Kept apart, a 768-dim `vectors` row +is ~3.2 kB and two of them share an 8 KiB page. Content is read only for the K winners of a search: one extra lookup per result. +**Why the scan reads a second copy of every vector.** The paragraph above stops +being true once the model is bigger than 1024 dimensions. A 2048-dim float32 +embedding is 8192 bytes on its own, past the 8157-byte local-payload limit, so +every `vectors` row spills into an overflow page and the scan is back to the +layout splitting out the content was meant to avoid. Measured with `dbstat` +over 400 rows, bytes a full scan must read per vector: + +| dimensions | representation | leaf | overflow | bytes/vector | +|---|---|---|---|---| +| 768 | float32 | 200 | 0 | 4096 | +| 1024 | float32 | 400 | 0 | 8192 | +| 2048 | float32 | 50 | 400 | 9216 | +| 768 | int8 | 40 | 0 | 819 | +| 1024 | int8 | 58 | 0 | 1188 | +| 2048 | int8 | 134 | 0 | 2744 | + +The pathological line is 1024, not 2048: nothing overflows there and the scan +still reads 8192 bytes to obtain 4096, because two 4.1 kB rows cannot share an +8 KiB page. Halving `output_dimension` to save time bought half the vector +quality for 89% of the I/O. + +`vectors_q8` removes the whole step function by scanning one byte per component +instead of four. `TestScanPackingEfficiency` and `TestScanBytesPerVectorBudget` +assert those numbers — as pages, not milliseconds, so they mean the same thing +in CI, on a laptop, and on the production box. + ## Search ``` -SELECT rowid, embedding FROM vectors INDEXED BY idx_vec_coll - WHERE collection_id = ? [AND ] +SELECT doc_id, scale, embedding FROM vectors_q8 INDEXED BY idx_q8_coll + WHERE collection_id = ? [AND language = ?] ``` -Rows stream past a dot product (embeddings are stored L2-normalised, so cosine -similarity *is* the dot product) into a top-K min-heap that rejects a losing -row with one comparison. Metadata and chunk text are fetched afterwards, for -the winners only. +Rows stream past an integer dot product into a top-K min-heap that rejects a +losing row with one comparison. The heap is wider than the caller's limit — the +int8 ranking chooses a shortlist, it does not produce the answer. The shortlist +is then rescored against the exact float32 vectors in `vectors`, and metadata +and chunk text are fetched for the winners only. + +**What the approximation costs.** Measured on 60k vectors of the load-test +fixture's largest collection (`ziglang/zig`, voyage-code-3 @2048) against 50 +real query-side embeddings, recall of the exact float32 top-K: + +| shortlist | k=10 | k=20 | +|---|---|---| +| 20 | 0.998 | 0.994 | +| 40 | 0.998 | 0.999 | +| **60** | **1.000** | **1.000** | +| 200 | 1.000 | 1.000 | + +Without rescoring at all the int8 ranking alone gives 0.994 at both k — the +quantisation misorders near-ties, it does not lose the documents, which is why +re-reading a few dozen exact vectors recovered every one of them here. +`q8Shortlist` therefore uses a floor of 64 and 4x the limit above it. Scan CPU +in the same run: 127 ms per query float32, 42 ms int8 (3.0x). + +Two different guarantees, worth keeping apart. **Scores are exact by +construction**: every number a caller sees is the cosine against the float32 +vector, computed by the rescore. **The result set is an approximation** whose +error was measured at zero on this corpus and is not proved at zero in general +— the shortlist is a fixed size and `topK` rejects boundary ties strictly, so a +collection holding more than `shortlist` documents inside one quantisation step +of each other (a file vendored a hundred times, say) can truncate a tie in scan +order, and the rescore cannot recover a document that was never shortlisted. +Widening the shortlist to swallow boundary ties would close that; it has not +been needed on any corpus measured so far. + +Scores returned to callers are always the exact cosine, never the int8 +estimate. That is load-bearing beyond cosmetics: `min_score` thresholds on it, +the workspace fan-out normalises across projects with it, and hybrid search +blends it with BM25 — an approximate score would move results between projects +in a way no single-project test would catch. `TestSearchScoresAreExact` pins it. + +**Building and rebuilding the copy.** Writes maintain `vectors_q8` in the same +transaction as `vectors`, so a collection created by this code is complete by +construction, and `q8_state` records that at creation — the readiness check is +a primary-key lookup, never a `COUNT`. A store written before the table existed +is converted by a background pass at open, largest collection first, in 2000-row +transactions at a 50% duty cycle; until a collection is covered its searches +take the float32 scan, which is correct and simply slower. Nothing is ever +marked complete before it is: the flag is written in the same transaction as +the batch that proves it. Set `CIX_VECTOR_SCAN_QUANT=false` to opt out — the +copy is roughly a quarter of the float32 bytes on top of an already large +store, so an operator short of disk needs a way to say no. Turning it off also +withdraws the completion flag from anything written while it is off, so turning +it back on rebuilds rather than trusting a stale copy. `INDEXED BY` is not an optimisation hint, it is a guarantee, and *which* index matters. Measured on the real index, scanning its largest (74k-row) collection: @@ -188,6 +263,10 @@ The metadata filter (`where`) mirrors chromem's semantics exactly, including the two odd cases: an unknown key with a non-empty value matches nothing, and an unknown key with an empty value matches everything. +`TestQ8ScanUsesCollectionIndex` pins the same guarantee for the compact table: +`idx_q8_coll`'s keys are `(collection_id, rowid)` for the same reason, and the +language filter must not change the driving index. + **Concurrency.** One scan per query, and a process-wide semaphore caps concurrent scans at `NumCPU`. Splitting a single query across workers was measured to buy nothing in the low-memory configuration (109 ms at 1 worker vs diff --git a/server/cmd/cix-server/main.go b/server/cmd/cix-server/main.go index 54c23cc..8943e47 100644 --- a/server/cmd/cix-server/main.go +++ b/server/cmd/cix-server/main.go @@ -411,6 +411,7 @@ func run() (restart bool, err error) { Dir: cfg.VectorDirFor(comps), LegacyChromaDir: cfg.ChromaDirFor(comps), MMapBytes: cfg.VectorMMapSize, + ScanQuant: cfg.VectorScanQuantEnabled, Logger: logger, }) } diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 2a346c5..c66f630 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -55,7 +55,19 @@ type Config struct { // 0 (the default) leaves it off. Env: CIX_VECTOR_MMAP_SIZE. It buys // roughly 40% lower search latency and costs resident memory: every // connection maps the database file and mapped pages count in RSS. - VectorMMapSize int64 + VectorMMapSize int64 + // VectorScanQuantEnabled controls the compact int8 copy the vector store + // scans instead of the float32 originals. Env: CIX_VECTOR_SCAN_QUANT, + // default true. + // + // It exists because turning it on costs disk before it saves time: the + // copy is about a quarter of the float32 bytes, added to a store that may + // already be the largest thing on the volume, and it is built by a + // background pass over every existing vector. An operator who is short of + // disk, or who wants to isolate a search-quality question from the + // approximation, needs a way to say no. Off means every collection keeps + // using the exact float32 scan — correct, and as slow as it was before. + VectorScanQuantEnabled bool SQLitePath string MaxFileSize int ExcludedDirs []string @@ -310,6 +322,12 @@ func Load() (*Config, error) { } c.VectorMMapSize = int64(vecMMap) + scanQuant, err := getenvBool("CIX_VECTOR_SCAN_QUANT", true) + if err != nil { + return nil, err + } + c.VectorScanQuantEnabled = scanQuant + authOff, err := getenvBool("CIX_AUTH_DISABLED", false) if err != nil { return nil, err diff --git a/server/internal/httpapi/workspacesearch.go b/server/internal/httpapi/workspacesearch.go index d065d83..dad38f2 100644 --- a/server/internal/httpapi/workspacesearch.go +++ b/server/internal/httpapi/workspacesearch.go @@ -509,26 +509,34 @@ func workspaceSearchResponse( // reindex before BM25 can contribute. A best-effort detector: if any // SQL probe errors out we log + return nil rather than fail the // request, since the warning is informational, not load-bearing. +// +// EXISTS, not COUNT(*), and the difference is not cosmetic: the question is +// "are there any rows", and COUNT walks every matching index entry to answer +// it. Measured on the 45-repo load-test index (1.95M rows in chunks_meta), +// this loop cost 53 ms per workspace search as COUNT and 0.2 ms as EXISTS — +// and it runs BEFORE the fan-out, so every query pays it serially. The LIMIT 1 +// that used to be on these statements did nothing: it bounds the result rows +// of an aggregate that always returns exactly one. func (s *Server) detectStaleFTSRepos(ctx context.Context, projectPaths []string) []workspaceSearchStaleFTSRepoPayload { out := make([]workspaceSearchStaleFTSRepoPayload, 0) for _, pp := range projectPaths { - var nMeta, nFiles int + var hasMeta, hasFiles bool if err := s.Deps.DB.QueryRowContext(ctx, - `SELECT COUNT(*) FROM chunks_meta WHERE project_path = ? LIMIT 1`, pp).Scan(&nMeta); err != nil { + `SELECT EXISTS(SELECT 1 FROM chunks_meta WHERE project_path = ?)`, pp).Scan(&hasMeta); err != nil { s.Deps.Logger.Warn("workspaces search: stale-fts probe (chunks_meta)", "project_path", pp, "err", err) return nil } - if nMeta > 0 { + if hasMeta { continue } if err := s.Deps.DB.QueryRowContext(ctx, - `SELECT COUNT(*) FROM file_hashes WHERE project_path = ? LIMIT 1`, pp).Scan(&nFiles); err != nil { + `SELECT EXISTS(SELECT 1 FROM file_hashes WHERE project_path = ?)`, pp).Scan(&hasFiles); err != nil { s.Deps.Logger.Warn("workspaces search: stale-fts probe (file_hashes)", "project_path", pp, "err", err) return nil } - if nFiles > 0 { + if hasFiles { out = append(out, workspaceSearchStaleFTSRepoPayload{ProjectPath: pp}) } } diff --git a/server/internal/vectorstore/chromemimport.go b/server/internal/vectorstore/chromemimport.go index dceaa12..64bd5fb 100644 --- a/server/internal/vectorstore/chromemimport.go +++ b/server/internal/vectorstore/chromemimport.go @@ -313,6 +313,25 @@ func (s *Store) importCollection(ctx context.Context, dir, name string) (int, er if err := tx.QueryRowContext(ctx, `SELECT id FROM collections WHERE name = ?`, name).Scan(&collID); err != nil { return 0, err } + // No compact scan copy is written here, deliberately: the import is already + // the slowest thing a boot can do, and the background backfill that runs + // right after it (see startQ8Backfill) converts the result at a duty cycle + // that leaves the server usable. Until then those collections search the + // float32 way, which is what they did before that table existed. + // + // But "the import creates the collection, so nothing marked it complete" + // is only true when the collection is NEW. INSERT OR IGNORE also succeeds + // against a collection this binary created and flagged earlier — an + // operator pointing CIX_CHROMA_PERSIST_DIR at a legacy tree after indexing + // the same project live reaches exactly that, because migration_state is + // keyed on the legacy collection name and has never seen it. Imported docs + // would then have no compact rows inside a collection whose flag says it + // is complete, and the backfill skips flagged collections: permanently + // invisible to search. So the flag comes off here, unconditionally. + if err := clearQ8Ready(ctx, tx, collID); err != nil { + return 0, err + } + s.forgetQ8(collID) vecStmt, err := tx.PrepareContext(ctx, upsertVectorSQL) if err != nil { return 0, err diff --git a/server/internal/vectorstore/layout_test.go b/server/internal/vectorstore/layout_test.go new file mode 100644 index 0000000..89eb09a --- /dev/null +++ b/server/internal/vectorstore/layout_test.go @@ -0,0 +1,170 @@ +package vectorstore + +import ( + "context" + "fmt" + "math/rand" + "strings" + "testing" +) + +// --------------------------------------------------------------------------- +// Page layout — what a scan's cost is actually made of. +// +// Search latency on a developer's laptop is not a portable number: it is a +// statement about that machine's disk, page cache and core count. What IS +// portable is how many database pages a full-collection scan is obliged to +// touch, because SQLite's layout rules are the same everywhere. Multiply pages +// by the target machine's read throughput and you have its latency; assert on +// pages and the assertion means the same thing in CI, on a Mac with NVMe, and +// on the 2-vCPU production box whose page cache is smaller than its index. +// --------------------------------------------------------------------------- + +// scanPages reports the pages a full scan must read, from whichever table the +// scan actually walks. +// +// dbstat walks the b-tree page by page, which is the only way to see overflow +// chains: they show up in no COUNT, in no per-table file size, and a row that +// spills is indistinguishable from one that does not until you look at its +// pages. +func scanPages(t *testing.T, s *Store, table string) (leaf, overflow, rows int64) { + t.Helper() + err := s.db.QueryRow(`SELECT COALESCE(SUM(pagetype='leaf'), 0), + COALESCE(SUM(pagetype='overflow'), 0) + FROM dbstat WHERE name = ?`, table).Scan(&leaf, &overflow) + if err != nil { + t.Fatalf("dbstat(%s): %v", table, err) + } + if err := s.db.QueryRow(`SELECT COUNT(*) FROM ` + table).Scan(&rows); err != nil { + t.Fatalf("count(%s): %v", table, err) + } + return leaf, overflow, rows +} + +// scanTable is the table a search walks, and TestLayoutMeasuresTheScannedTable +// checks that scanQ8SQL still names it — otherwise a change of scan source +// would leave every measurement below pointed at a table nobody reads, and the +// numbers would keep passing while meaning nothing. +const scanTable = "vectors_q8" + +// fillDim writes n rows of dimension dim into one collection. +func fillDim(t *testing.T, s *Store, project string, n, dim int) { + t.Helper() + r := rand.New(rand.NewSource(7)) + chunks := make([]Chunk, n) + embs := make([][]float32, n) + for i := range chunks { + chunks[i] = Chunk{ + Content: "package main\n\nfunc main() {}\n", + FilePath: fmt.Sprintf("pkg/mod%03d/file%03d.go", i%16, i), + StartLine: i*10 + 1, + EndLine: i*10 + 9, + ChunkType: "function", + SymbolName: fmt.Sprintf("Handler%03d", i), + Language: "go", + } + embs[i] = randNorm(r, dim) + } + if err := s.UpsertChunks(context.Background(), project, chunks, embs); err != nil { + t.Fatalf("upsert: %v", err) + } +} + +// TestScanPackingEfficiency pins how much of what a scan reads is the data it +// came for. +// +// SQLite keeps a row inside its leaf page only while the payload fits +// usable-35 bytes (8157 on our 8 KiB pages). Past that it keeps about a +// kilobyte local and puts the rest in a chain of overflow pages. Neither the +// row size nor the crossing of that line is visible anywhere in the schema — +// and the row size is set by an operator choosing output_dimension in a config +// file. Measured on this schema, three dimensions behave in three different +// ways: +// +// 768 3.1 kB row two rows share a leaf page 4096 B/vector 1.33x +// 1024 4.1 kB row one row per leaf page 8192 B/vector 2.00x +// 2048 8.2 kB row leaf slice + one overflow page 9216 B/vector 1.12x +// +// The interesting line is 1024, not 2048. Overflow sounds like the pathology +// and is not: at 2048 the overflow page is nearly full, so the scan reads +// 9216 bytes to obtain 8192 useful ones. At 1024 nothing overflows and the +// scan still reads 8192 bytes to obtain 4096, because two 4.1 kB rows cannot +// share an 8 KiB page and the second half of every page is air. An operator +// who halves output_dimension to save disk and time gets half the vector +// quality for 89% of the I/O. +// +// The assertion is the ratio, so it survives a change of dimension, of page +// size, or of which columns live in this table. +func TestScanPackingEfficiency(t *testing.T) { + // A scan that reads more than 1.4x the bytes it needs is spending more on + // structure than any layout choice should cost. 1.33x — two rows to a + // page with the page header and cell pointers on top — is what a healthy + // packing looks like, and is what both 768-dim float32 and 2048-dim int8 + // achieve. + const maxRatio = 1.4 + + for _, dim := range []int{768, 1024, 2048} { + t.Run(fmt.Sprintf("dim%d", dim), func(t *testing.T) { + s := openStore(t) + fillDim(t, s, "/layout", 400, dim) + + leaf, overflow, rows := scanPages(t, s, scanTable) + perVec := float64((leaf+overflow)*pageSize) / float64(rows) + // One byte per component: what the scan reads, against what it + // reads it for. + payload := float64(dim) + ratio := perVec / payload + + t.Logf("dim=%d rows=%d leaf=%d overflow=%d %.0f B/vector %.2fx payload", + dim, rows, leaf, overflow, perVec, ratio) + + if ratio > maxRatio { + t.Errorf("scan reads %.0f B per %d-dim vector to obtain %.0f B of embedding "+ + "(%.2fx, limit %.2fx): leaf=%d overflow=%d over %d rows", + perVec, dim, payload, ratio, maxRatio, leaf, overflow, rows) + } + }) + } +} + +// TestScanBytesPerVectorBudget is the absolute number, and the one the search +// work exists to move. +// +// Packing efficiency says the scan wastes little; it says nothing about the +// scan being affordable. At 2048 dimensions a well-packed float32 scan still +// reads 9 kB per vector, and a workspace query scans every collection: on the +// 45-repo fixture that is 1.9M vectors, about 17 GB of reads for one search. +// No page cache on an 8 GB box holds that, so production pays it at disk +// speed, every query, per repo. +// +// The budget below is what the scan costs when it reads a compact +// representation instead of the float32 original — the float32 blob stays on +// disk for anything that needs exact scores, but the scan stops reading it. +func TestScanBytesPerVectorBudget(t *testing.T) { + const dim = 2048 + // 2048 int8 components + row overhead, three rows to a page. + const maxBytesPerVector = 3072 + + s := openStore(t) + fillDim(t, s, "/budget", 400, dim) + + leaf, overflow, rows := scanPages(t, s, scanTable) + perVec := float64((leaf+overflow)*pageSize) / float64(rows) + t.Logf("dim=%d rows=%d leaf=%d overflow=%d %.0f B/vector", dim, rows, leaf, overflow, perVec) + + if perVec > maxBytesPerVector { + t.Errorf("scan reads %.0f B per vector at %d dims, budget %d B: "+ + "a 1.9M-vector workspace query moves %.1f GB instead of %.1f GB", + perVec, dim, maxBytesPerVector, + perVec*1.9e6/1e9, float64(maxBytesPerVector)*1.9e6/1e9) + } +} + +// TestLayoutMeasuresTheScannedTable keeps the two constants honest. dbstat is +// asked about a table by name, and a name is exactly the kind of thing that +// survives a refactor that moved the data somewhere else. +func TestLayoutMeasuresTheScannedTable(t *testing.T) { + if !strings.Contains(scanQ8SQL, " "+scanTable+" ") { + t.Fatalf("the scan reads a table other than %q:\n%s", scanTable, scanQ8SQL) + } +} diff --git a/server/internal/vectorstore/maintenance.go b/server/internal/vectorstore/maintenance.go index 9503bea..3c79891 100644 --- a/server/internal/vectorstore/maintenance.go +++ b/server/internal/vectorstore/maintenance.go @@ -88,6 +88,26 @@ SELECT c.name, COALESCE(SUM(LENGTH(vc.content) + LENGTH(vc.doc_id)), 0) LEFT JOIN vector_contents vc ON vc.collection_id = c.id GROUP BY c.id` +// q8SizeSQL accounts for the compact scan copy. It is a third of the float32 +// bytes and it is real disk, so leaving it out would make the Resources screen +// under-report the store by ~25% — the same kind of quiet mismatch the WAL +// high-water mark used to cause. +// sizeExprQ8 is the compact copy's per-row logical byte count, in the same +// shape as sizeExprVectors and pasted in no more places than it is. +// +// Reading it walks the compact table's leaf pages — about a fifth of what the +// same question costs over `vectors`, and behind the maintenance service's TTL +// cache either way, so it is answered rarely. If that stops being true, the +// cheap replacement is recording the total at backfill completion rather than +// making the aggregate faster. +const sizeExprQ8 = `LENGTH(q.embedding) + LENGTH(q.doc_id) + LENGTH(q.language) + 16` + +const q8SizeSQL = ` +SELECT c.name, COALESCE(SUM(` + sizeExprQ8 + `), 0) + FROM collections c + LEFT JOIN vectors_q8 q ON q.collection_id = c.id + GROUP BY c.id` + // ListCollections implements Maintainer. Results are sorted by name so callers // (and their tests) see a stable order. func (s *Store) ListCollections() []CollectionInfo { @@ -118,29 +138,45 @@ func (s *Store) ListCollections() []CollectionInfo { return nil } - // Chunk text lives in its own table (see schemaSQL); fold it into the - // reported size with a second aggregate rather than a join that would - // multiply the row counts. - crows, err := s.db.QueryContext(ctx, contentSizeSQL) + // Chunk text and the compact scan copy live in their own tables (see + // schemaSQL); fold each into the reported size with its own aggregate + // rather than joins that would multiply the row counts. + for _, q := range []struct { + what string + sql string + }{ + {"contents", contentSizeSQL}, + {"scan copy", q8SizeSQL}, + } { + sizes, err := s.sizesByCollection(ctx, q.sql) + if err != nil { + s.logger.Error("vectorstore: list collection "+q.what, "err", err) + return out + } + for i := range out { + out[i].SizeBytes += sizes[out[i].Name] + } + } + return out +} + +// sizesByCollection runs one name -> bytes aggregate. +func (s *Store) sizesByCollection(ctx context.Context, query string) (map[string]int64, error) { + rows, err := s.db.QueryContext(ctx, query) if err != nil { - s.logger.Error("vectorstore: list collection contents", "err", err) - return out + return nil, err } - defer crows.Close() - sizes := make(map[string]int64, len(out)) - for crows.Next() { + defer rows.Close() + sizes := map[string]int64{} + for rows.Next() { var name string var n int64 - if err := crows.Scan(&name, &n); err != nil { - s.logger.Error("vectorstore: list collection contents", "err", err) - return out + if err := rows.Scan(&name, &n); err != nil { + return nil, err } sizes[name] = n } - for i := range out { - out[i].SizeBytes += sizes[out[i].Name] - } - return out + return sizes, rows.Err() } // CollectionSizeBytes implements Maintainer. @@ -155,7 +191,7 @@ func (s *Store) CollectionSizeBytes(projectPath string) (int64, bool) { if err != nil || !ok { return 0, false } - var vecBytes, contentBytes int64 + var vecBytes, contentBytes, q8Bytes int64 if err := s.db.QueryRowContext(ctx, ` SELECT COALESCE(SUM(`+sizeExprVectors+`), 0) FROM vectors v WHERE v.collection_id = ?`, collID).Scan(&vecBytes); err != nil { @@ -166,7 +202,12 @@ func (s *Store) CollectionSizeBytes(projectPath string) (int64, bool) { FROM vector_contents WHERE collection_id = ?`, collID).Scan(&contentBytes); err != nil { return 0, false } - return vecBytes + contentBytes, true + if err := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(`+sizeExprQ8+`), 0) + FROM vectors_q8 q WHERE q.collection_id = ?`, collID).Scan(&q8Bytes); err != nil { + return 0, false + } + return vecBytes + contentBytes + q8Bytes, true } // DeleteCollectionByName implements Maintainer. @@ -196,6 +237,8 @@ func (s *Store) DeleteCollectionByName(name string) error { for _, stmt := range []string{ `DELETE FROM vector_contents WHERE collection_id = ?`, + `DELETE FROM vectors_q8 WHERE collection_id = ?`, + `DELETE FROM q8_state WHERE collection_id = ?`, `DELETE FROM vectors WHERE collection_id = ?`, `DELETE FROM collections WHERE id = ?`, } { @@ -207,6 +250,7 @@ func (s *Store) DeleteCollectionByName(name string) error { return fmt.Errorf("vectorstore delete collection %q: %w", name, err) } s.forgetCollection(name) + s.forgetQ8(collID) return nil } diff --git a/server/internal/vectorstore/parity_test.go b/server/internal/vectorstore/parity_test.go index e140083..915970d 100644 --- a/server/internal/vectorstore/parity_test.go +++ b/server/internal/vectorstore/parity_test.go @@ -207,6 +207,19 @@ func TestSearchWhereFilterMirrorsChromemSemantics(t *testing.T) { t.Errorf("start_line filter returned %+v, want the single chunk starting at 11", got) } + // A KNOWN key with an empty value is a filter, not the absence of one: + // chromem compared metadata["language"] to "", which matches only rows + // whose language is empty. The compact scan supports exactly this one + // column, so it is the one place the two scan paths could disagree about + // what an empty value means. + got, err = s.Search(ctx, project, embs[0], 5, map[string]string{"language": ""}) + if err != nil { + t.Fatalf("empty language filter: %v", err) + } + if len(got) != 0 { + t.Errorf("language=\"\" returned %d results, want 0 — every chunk here is Go", len(got)) + } + // Several keys must all match. got, err = s.Search(ctx, project, embs[0], 5, map[string]string{"language": "go", "file_path": "b.go"}) diff --git a/server/internal/vectorstore/plan_test.go b/server/internal/vectorstore/plan_test.go index e2e5069..9b9b8da 100644 --- a/server/internal/vectorstore/plan_test.go +++ b/server/internal/vectorstore/plan_test.go @@ -74,3 +74,33 @@ func TestScanUsesCollectionIndex(t *testing.T) { t.Errorf("delete-by-file plan does not use idx_vec_coll_file:\n%s", plan) } } + +// TestQ8ScanUsesCollectionIndex is TestScanUsesCollectionIndex for the table a +// search actually walks. Same guarantee for the same reason — idx_q8_coll's +// keys are (collection_id, rowid), so the walk stays proportional to the +// collection and yields its rows in table order — and it matters more here, +// because this is the plan every search takes. +func TestQ8ScanUsesCollectionIndex(t *testing.T) { + s := openStore(t) + ctx := context.Background() + chunks, embs := makeChunks(20, "a.go", "go") + if err := s.UpsertChunks(ctx, "/q8plan", chunks, embs); err != nil { + t.Fatalf("upsert: %v", err) + } + + plan := queryPlan(t, s, scanQ8SQL, 1) + if !strings.Contains(plan, "idx_q8_coll") { + t.Errorf("compact scan plan does not use idx_q8_coll:\n%s", plan) + } + if strings.Contains(plan, "SCAN vectors_q8\n") || strings.HasSuffix(plan, "SCAN vectors_q8") { + t.Errorf("compact scan plan falls back to a full table scan:\n%s", plan) + } + + // The language filter must not change the driving index: it is an extra + // test on rows the index already visits, not a reason to pick a different + // one. + plan = queryPlan(t, s, scanQ8SQL+" AND language = ?", 1, "go") + if !strings.Contains(plan, "idx_q8_coll") { + t.Errorf("filtered compact scan plan does not use idx_q8_coll:\n%s", plan) + } +} diff --git a/server/internal/vectorstore/q8.go b/server/internal/vectorstore/q8.go new file mode 100644 index 0000000..93b4ca0 --- /dev/null +++ b/server/internal/vectorstore/q8.go @@ -0,0 +1,448 @@ +package vectorstore + +import ( + "context" + "database/sql" + "errors" + "fmt" + "time" +) + +// --------------------------------------------------------------------------- +// The compact scan copy. +// +// vectors_q8 holds every vector at one byte per component. A search scans it +// instead of the float32 table (3.4x fewer bytes, measured — see schemaSQL), +// takes a shortlist, and rescores that shortlist against the float32 +// originals, which is what keeps the answer exact. +// +// Everything here exists to answer one question cheaply and correctly: does +// this collection have a q8 row for every vector it has? The answer must not +// cost a COUNT per query, and it must never be "yes" when it is not — a +// half-built q8 table would silently hide documents from search, and a store +// that quietly returns fewer results is worse than a slow one. +// +// The invariant is maintained from both ends: +// +// - A collection created by this code is born complete: it has no rows, so +// the empty q8 side matches it, and ensureCollection records that. Every +// upsert afterwards writes both tables in one transaction, so the property +// is preserved by construction and never has to be re-checked. +// - A collection that predates this table has rows and no q8_state row. It +// stays on the float32 scan — correct, just slower — until the backfill +// has quantised all of it and records completion in the same transaction +// as the last batch. +// +// Nothing sets the flag optimistically, and nothing reads q8 without it. +// --------------------------------------------------------------------------- + +// backfillQ8SQL writes one converted vector, and is deliberately weaker than +// the upsert the write path uses. +// +// The backfill reads a batch, quantises it, and writes it in a separate +// transaction. Anything can commit in that gap — the file watcher reindexing a +// saved file, an admin deleting a project — so the write has to be harmless +// against both, without taking a lock that would make the watcher wait on a +// background job. Two clauses do that: +// +// - WHERE EXISTS: a doc deleted in the gap is not resurrected. Without it the +// backfill would reinsert compact rows whose float32 originals are gone — +// rows that DeleteByFile can never clean up afterwards, because its +// subquery finds doc_ids through `vectors`. Every later scan would +// shortlist them and every rescore would silently drop them, which reads +// as a search returning fewer results for no stated reason. +// - DO NOTHING: a doc re-embedded in the gap keeps the compact row upsertBatch +// wrote from its NEW embedding, instead of being overwritten by this batch's +// quantisation of the old one. The backfill only ever fills gaps; it is +// never the more recent writer. +const backfillQ8SQL = `INSERT INTO vectors_q8 (collection_id, doc_id, language, scale, embedding) + SELECT ?,?,?,?,? WHERE EXISTS (SELECT 1 FROM vectors WHERE collection_id = ? AND doc_id = ?) + ON CONFLICT(collection_id, doc_id) DO NOTHING` + +// q8BackfillBatch is how many vectors one backfill transaction converts. +// +// At 2048 dimensions this reads ~16 MB and writes ~4 MB per batch. Small +// enough that a writer (the indexer, the file watcher) never waits long for +// the write lock, large enough that the per-transaction overhead is noise. +const q8BackfillBatch = 2000 + +// q8BackfillDuty is the fraction of wall-clock the backfill is allowed to +// spend working. It sleeps for the rest. +// +// The backfill competes with live searches for the same disk and the same two +// vCPUs on the production box, and it is never urgent: until it finishes, the +// affected collections simply search the way they did before. Yielding half +// the time turns "the server is unusable for ten minutes after an upgrade" +// into "it is a bit slower for twenty". +const q8BackfillDuty = 0.5 + +// markQ8Ready records that a collection's q8 rows are complete. +// +// INSERT OR IGNORE, so calling it for a collection already marked is free and +// keeps the original timestamp — which is the one that says when the data was +// actually built. +func markQ8Ready(ctx context.Context, tx *sql.Tx, collID int64) error { + _, err := tx.ExecContext(ctx, + `INSERT OR IGNORE INTO q8_state(collection_id, built_at) VALUES(?, ?)`, + collID, time.Now().UTC().Format(time.RFC3339)) + if err != nil { + return fmt.Errorf("vectorstore: mark q8 ready for collection %d: %w", collID, err) + } + return nil +} + +// markCollectionQ8Ready records completion outside a caller-owned +// transaction, and updates the in-memory cache so the very next search on a +// freshly created collection takes the fast path. +func (s *Store) markCollectionQ8Ready(ctx context.Context, collID int64) error { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck // no-op after a successful Commit + if err := markQ8Ready(ctx, tx, collID); err != nil { + return err + } + if err := tx.Commit(); err != nil { + return err + } + s.q8Mu.Lock() + s.q8State[collID] = true + s.q8Mu.Unlock() + return nil +} + +// clearQ8Ready withdraws a collection's completion flag. +func clearQ8Ready(ctx context.Context, tx *sql.Tx, collID int64) error { + if _, err := tx.ExecContext(ctx, `DELETE FROM q8_state WHERE collection_id = ?`, collID); err != nil { + return fmt.Errorf("vectorstore: clear q8 state for collection %d: %w", collID, err) + } + return nil +} + +// clearCollectionQ8Ready withdraws a collection's completion flag and forgets +// the cached answer, so the very next search falls back to the exact scan. +func (s *Store) clearCollectionQ8Ready(ctx context.Context, collID int64) error { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() //nolint:errcheck // no-op after a successful Commit + if err := clearQ8Ready(ctx, tx, collID); err != nil { + return err + } + if err := tx.Commit(); err != nil { + return err + } + s.forgetQ8(collID) + return nil +} + +// q8Ready reports whether the scan may read vectors_q8 for this collection. +// +// Cached in memory because it is consulted on every search and the answer only +// ever changes in one direction (not ready -> ready, once, when the backfill +// finishes). A false answer is therefore worth re-checking; a true one is not. +func (s *Store) q8Ready(ctx context.Context, collID int64) bool { + if !s.scanQuant { + return false + } + s.q8Mu.Lock() + _, ready := s.q8State[collID] + s.q8Mu.Unlock() + if ready { + return true + } + + // Only positives are cached, and presence IS the answer — there is no + // stored false to accidentally honour. A collection that is not ready yet + // may become ready at any moment (the backfill is running), so a negative + // has to be re-asked; a positive never reverts without going through + // forgetQ8. + var one int + err := s.db.QueryRowContext(ctx, + `SELECT 1 FROM q8_state WHERE collection_id = ?`, collID).Scan(&one) + switch { + case err == nil: + case errors.Is(err, sql.ErrNoRows): + return false + default: + // A probe that errors must not upgrade the scan: falling back to the + // float32 path answers the query correctly. + s.logger.Warn("vectorstore: q8 readiness probe failed", "collection_id", collID, "err", err) + return false + } + s.q8Mu.Lock() + s.q8State[collID] = true + s.q8Mu.Unlock() + return true +} + +// forgetQ8 drops a collection's cached readiness (after a delete). The next +// collection to be handed this id — which AUTOINCREMENT guarantees is never +// this one — must not inherit its answer. +func (s *Store) forgetQ8(collID int64) { + s.q8Mu.Lock() + delete(s.q8State, collID) + s.q8Mu.Unlock() +} + +// startQ8Backfill converts collections written before vectors_q8 existed. +// +// Runs in the background and returns immediately: the store is fully usable +// while it works, because an unconverted collection is not broken, only slow. +// That is the whole reason this is a background job and not part of Open — the +// alternative is a server that answers nothing for the minutes it takes to +// rewrite a multi-gigabyte store, which is exactly the failure mode the schema +// rebuild already has and which took three false "the server is down" reports +// to diagnose. +func (s *Store) startQ8Backfill(ctx context.Context) { + go func() { + if err := s.backfillQ8(ctx); err != nil && ctx.Err() == nil { + // Warn, not fatal: every collection it failed to convert keeps + // searching the float32 way. + s.logger.Warn("vectorstore: building the compact scan index stopped early", "err", err) + } + }() +} + +// pendingQ8Collections lists collections that have vectors but no completed q8 +// copy, largest first — so the collection whose searches hurt most is the +// first one to get faster. +func (s *Store) pendingQ8Collections(ctx context.Context) ([]int64, error) { + rows, err := s.db.QueryContext(ctx, ` + SELECT v.collection_id, COUNT(*) n + FROM vectors v + WHERE v.collection_id NOT IN (SELECT collection_id FROM q8_state) + GROUP BY v.collection_id + ORDER BY n DESC`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []int64 + for rows.Next() { + var id int64 + var n int64 + if err := rows.Scan(&id, &n); err != nil { + return nil, err + } + out = append(out, id) + } + return out, rows.Err() +} + +// backfillQ8 quantises every pending collection. +func (s *Store) backfillQ8(ctx context.Context) error { + if !s.acquire() { + return nil + } + pending, err := s.pendingQ8Collections(ctx) + s.release() + if err != nil { + return fmt.Errorf("list collections needing a scan copy: %w", err) + } + if len(pending) == 0 { + return nil + } + + // The scan copy is a quarter of the float32 bytes it is derived from. + // Refusing up front beats discovering it a gigabyte in: a failed backfill + // leaves the store fully working, but it also leaves the disk fuller than + // it needs to be and the failure buried in a log line. + if need, err := s.pendingQ8Bytes(ctx); err == nil { + if err := checkFreeSpace(s.dir, need); err != nil { + s.logger.Warn("vectorstore: skipping the compact scan index — not enough free disk", + "db", s.dbPath, "need_mb", need/(1<<20), "err", err) + return nil + } + } + + // WARN for the same reason the schema rebuild and the legacy import log at + // warn: production runs at warn level, and unexplained background I/O on a + // box that was just restarted is indistinguishable from a problem. + started := time.Now() + s.logger.Warn("vectorstore: building the compact scan index in the background", + "db", s.dbPath, "collections", len(pending)) + + var converted int64 + var failed int + for _, collID := range pending { + n, err := s.backfillCollection(ctx, collID) + converted += n + if err != nil { + if ctx.Err() != nil { + return nil + } + // One collection's failure must not cost the other forty-two. + // The realistic cause is a collection deleted out from under the + // walk — an admin removing a project, or the orphan sweep — which + // makes the next insert fail its foreign key. Aborting there would + // leave every collection after it on the float32 scan until + // somebody restarted the server, and the only trace would be one + // warn line. + failed++ + s.logger.Warn("vectorstore: could not build the compact scan index for a collection", + "db", s.dbPath, "collection_id", collID, "err", err) + continue + } + } + s.logger.Warn("vectorstore: compact scan index built", + "db", s.dbPath, "collections", len(pending)-failed, "failed", failed, + "vectors", converted, "took", time.Since(started).Round(time.Second)) + return nil +} + +// pendingQ8Bytes estimates the disk the backfill will add: one byte per +// component of every vector it has to convert, plus the row overhead the size +// accounting already uses for the compact table. +func (s *Store) pendingQ8Bytes(ctx context.Context) (int64, error) { + if !s.acquire() { + return 0, ErrClosed + } + defer s.release() + var n int64 + // LENGTH(embedding)/4 is the compact row's blob length derived from the + // float32 one — same arithmetic as sizeExprQ8, from the only table that + // has the rows yet. + err := s.db.QueryRowContext(ctx, ` + SELECT COALESCE(SUM(LENGTH(embedding)/4 + LENGTH(doc_id) + LENGTH(language) + 16), 0) + FROM vectors + WHERE collection_id NOT IN (SELECT collection_id FROM q8_state)`).Scan(&n) + return n, err +} + +// backfillCollection walks one collection in doc_id order, quantising as it +// goes, and marks the collection ready in the same transaction as its last +// batch — so a kill at any point leaves a collection that is unmarked and +// therefore still searchable the float32 way, never one that is marked and +// incomplete. +func (s *Store) backfillCollection(ctx context.Context, collID int64) (int64, error) { + var ( + after int64 + converted int64 + ) + for { + if ctx.Err() != nil { + return converted, ctx.Err() + } + batchStart := time.Now() + n, last, err := s.backfillBatch(ctx, collID, after) + if err != nil { + return converted, err + } + converted += n + if n == 0 { + return converted, nil + } + after = last + + // Yield. Sleeping proportionally to the work just done keeps the duty + // cycle honest whether a batch took 40 ms on an NVMe laptop or four + // seconds on a network disk. + select { + case <-ctx.Done(): + return converted, ctx.Err() + case <-time.After(time.Duration(float64(time.Since(batchStart)) * (1 - q8BackfillDuty) / q8BackfillDuty)): + } + } +} + +// backfillBatch converts up to q8BackfillBatch vectors whose rowid is above +// `after`, and returns how many it converted and the last rowid it saw. +// +// Keyset pagination rather than OFFSET, so resuming does not re-read +// everything before the cursor. On ROWID and idx_vec_coll rather than on +// doc_id: `vectors` is a rowid table whose composite primary key is a separate +// index, so paging in doc_id order would look up ~9 kB rows scattered across +// the collection's whole rowid span — the same 1.8x that made scanSQL pick +// idx_vec_coll over idx_vec_coll_file (see sqlite.go). Rows inserted ahead of +// the cursor while the walk runs need no special handling: upsertBatch writes +// their compact copy itself. +// +// The batch does NOT hold a lock across its read and its write, and does not +// need to. Two statement-level rules make the gap between them harmless: +// the insert is conditional on the vectors row still existing, so a delete +// that commits in the gap cannot be undone; and it never overwrites an +// existing compact row, so an upsert that commits in the gap keeps its own +// fresher quantisation instead of being clobbered by this one's stale copy. +func (s *Store) backfillBatch(ctx context.Context, collID int64, after int64) (int64, int64, error) { + if !s.acquire() { + return 0, 0, ErrClosed + } + defer s.release() + + type q8Row struct { + docID string + language string + scale float32 + blob []byte + } + batch := make([]q8Row, 0, q8BackfillBatch) + last := after + + // Read and quantise first, write second. Holding a write transaction open + // across a streaming read would keep the write lock for the whole batch, + // and the thing most likely to want that lock is the file watcher + // reindexing a file someone just saved. + rows, err := s.db.QueryContext(ctx, ` + SELECT rowid, doc_id, language, embedding FROM vectors INDEXED BY idx_vec_coll + WHERE collection_id = ? AND rowid > ? + ORDER BY rowid LIMIT ?`, collID, after, q8BackfillBatch) + if err != nil { + return 0, 0, fmt.Errorf("read vectors: %w", err) + } + var scratch []float32 + for rows.Next() { + var ( + rowID int64 + docID, language string + raw sql.RawBytes + ) + if err := rows.Scan(&rowID, &docID, &language, &raw); err != nil { + rows.Close() + return 0, 0, fmt.Errorf("scan vector: %w", err) + } + var vec []float32 + vec, scratch = blobFloats(raw, scratch) + // quantizeInt8 allocates its own output, so nothing here outlives the + // RawBytes it was derived from. + blob, scale := quantizeInt8(vec) + batch = append(batch, q8Row{docID: docID, language: language, scale: scale, blob: blob}) + last = rowID + } + if err := rows.Err(); err != nil { + rows.Close() + return 0, 0, fmt.Errorf("read vectors: %w", err) + } + rows.Close() + if len(batch) == 0 { + // The cursor ran off the end: every row this collection had when the + // walk started now has a compact copy, and every row written since got + // one from upsertBatch. Completeness does not rest on this statement + // being in the same transaction as anything — it rests on the two + // insert rules above, which hold whatever commits in between. + return 0, last, s.markCollectionQ8Ready(ctx, collID) + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return 0, 0, err + } + defer tx.Rollback() //nolint:errcheck // no-op after a successful Commit + stmt, err := tx.PrepareContext(ctx, backfillQ8SQL) + if err != nil { + return 0, 0, err + } + defer stmt.Close() + for _, r := range batch { + if _, err := stmt.ExecContext(ctx, collID, r.docID, r.language, r.scale, r.blob, + collID, r.docID); err != nil { + return 0, 0, fmt.Errorf("write q8 row: %w", err) + } + } + if err := tx.Commit(); err != nil { + return 0, 0, err + } + return int64(len(batch)), last, nil +} diff --git a/server/internal/vectorstore/q8_test.go b/server/internal/vectorstore/q8_test.go new file mode 100644 index 0000000..7f3983c --- /dev/null +++ b/server/internal/vectorstore/q8_test.go @@ -0,0 +1,782 @@ +package vectorstore + +import ( + "context" + "fmt" + "math" + "math/rand" + "strings" + "sync" + "testing" + "time" +) + +// --------------------------------------------------------------------------- +// The compact scan copy. What has to be true of it: +// +// - the scores it reports are the exact ones, because the shortlist is +// rescored on the float32 originals; +// - it returns the same documents an exact scan would, which is an empirical +// property of the shortlist width and is therefore measured, not assumed; +// - a collection it has not covered yet still answers correctly; +// - deleting data deletes it here too, in both directions. +// --------------------------------------------------------------------------- + +// q8Corpus builds a collection that is hostile to quantisation: half the +// vectors are random, and the other half are near-duplicates of a few cluster +// centres, differing by less than the quantisation step. Random unit vectors +// in 2048 dimensions are almost orthogonal to each other and to any query, so +// a corpus of only those has no near-ties to misorder and would let any +// approximation look perfect. Real code corpora are the opposite: boilerplate, +// generated files and copied blocks produce exactly these clusters. +func q8Corpus(t *testing.T, s *Store, project string, n, dim int) ([]Chunk, [][]float32) { + t.Helper() + r := rand.New(rand.NewSource(11)) + centres := make([][]float32, 8) + for i := range centres { + centres[i] = randNorm(r, dim) + } + + chunks := make([]Chunk, n) + embs := make([][]float32, n) + langs := []string{"go", "python", "rust"} + for i := range chunks { + var v []float32 + if i%2 == 0 { + v = randNorm(r, dim) + } else { + base := centres[i%len(centres)] + v = make([]float32, dim) + for j := range v { + v[j] = base[j] + float32(r.NormFloat64())*1e-4 + } + v = normalizeVector(v) + } + chunks[i] = Chunk{ + Content: fmt.Sprintf("chunk %d", i), + FilePath: fmt.Sprintf("src/pkg%02d/f%04d.go", i%20, i), + StartLine: i*10 + 1, + EndLine: i*10 + 9, + ChunkType: "function", + SymbolName: fmt.Sprintf("Fn%04d", i), + Language: langs[i%len(langs)], + } + embs[i] = v + } + if err := s.UpsertChunks(context.Background(), project, chunks, embs); err != nil { + t.Fatalf("upsert: %v", err) + } + return chunks, embs +} + +// exactTopK is the oracle: the ranking a full float32 scan produces, computed +// in Go from the embeddings the test itself wrote. Deliberately not computed +// by asking the store to scan the other way — an oracle that shares code with +// the thing under test can agree with it about a shared mistake. +func exactTopK(chunks []Chunk, embs [][]float32, q []float32, k int, language string) []string { + type sc struct { + key string + score float32 + } + var all []sc + for i, e := range embs { + if language != "" && chunks[i].Language != language { + continue + } + all = append(all, sc{locKey(chunks[i]), dot(q, e)}) + } + for i := 1; i < len(all); i++ { + for j := i; j > 0 && all[j].score > all[j-1].score; j-- { + all[j], all[j-1] = all[j-1], all[j] + } + } + out := make([]string, 0, k) + for i := 0; i < k && i < len(all); i++ { + out = append(out, all[i].key) + } + return out +} + +// locKey identifies a chunk the way a caller sees it — the doc_id is internal. +func locKey(c Chunk) string { return fmt.Sprintf("%s:%d-%d", c.FilePath, c.StartLine, c.EndLine) } + +func resultKeys(rs []SearchResult) []string { + out := make([]string, len(rs)) + for i, r := range rs { + out[i] = fmt.Sprintf("%s:%d-%d", r.FilePath, r.StartLine, r.EndLine) + } + return out +} + +// TestSearchScoresAreExact is the property that does not depend on the corpus, +// the query or the shortlist width: whatever documents come back, the number +// attached to each is the true cosine against the stored float32 vector, not +// the int8 approximation that selected it. +// +// It matters because the score is not decoration. Callers threshold on it +// (min_score), the workspace fan-out normalises across projects with it, and +// hybrid search blends it with BM25. An approximate score would move results +// between projects in ways no per-project test would catch. +func TestSearchScoresAreExact(t *testing.T) { + const dim = 512 + s := openStore(t) + ctx := context.Background() + chunks, embs := q8Corpus(t, s, "/exact", 1500, dim) + + r := rand.New(rand.NewSource(99)) + for qi := 0; qi < 10; qi++ { + q := randNorm(r, dim) + got, err := s.Search(ctx, "/exact", q, 10, nil) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(got) == 0 { + t.Fatal("no results") + } + byKey := map[string]float32{} + for i := range chunks { + byKey[locKey(chunks[i])] = dot(q, embs[i]) + } + for _, res := range got { + key := fmt.Sprintf("%s:%d-%d", res.FilePath, res.StartLine, res.EndLine) + want := round4(byKey[key]) + if math.Abs(float64(res.Score-want)) > 1e-4 { + t.Errorf("query %d: %s scored %v, exact cosine is %v — "+ + "the reported score came from the int8 approximation, not the rescore", + qi, key, res.Score, want) + } + } + } +} + +// TestSearchMatchesExactRanking measures what the shortlist width buys. +// +// The compact scan is an approximation and could in principle drop a document +// that belongs in the top K. q8Shortlist picks its width from a measurement on +// real data (see its comment); this is the same measurement in miniature, on a +// corpus built to contain the near-ties that quantisation actually confuses, +// and it fails loudly if a change to the width, the quantisation or the +// rescore starts losing documents. +func TestSearchMatchesExactRanking(t *testing.T) { + const ( + dim = 512 + k = 10 + ) + s := openStore(t) + ctx := context.Background() + chunks, embs := q8Corpus(t, s, "/rank", 2000, dim) + + collID, ok, err := s.collectionID(ctx, collectionName("/rank")) + if err != nil || !ok { + t.Fatalf("collection id: %v ok=%v", err, ok) + } + if !s.q8Ready(ctx, collID) { + t.Fatal("collection is not on the compact scan — this test would be measuring the float32 path") + } + + r := rand.New(rand.NewSource(7)) + for qi := 0; qi < 25; qi++ { + q := randNorm(r, dim) + got, err := s.Search(ctx, "/rank", q, k, nil) + if err != nil { + t.Fatalf("search: %v", err) + } + want := exactTopK(chunks, embs, q, k, "") + if strings.Join(resultKeys(got), ",") != strings.Join(want, ",") { + t.Errorf("query %d: compact scan returned a different top-%d than an exact scan\n got: %v\nwant: %v", + qi, k, resultKeys(got), want) + } + } +} + +// TestSearchLanguageFilterOnCompactScan pins the one metadata column the +// compact copy carries. It is duplicated from `vectors`, so it can drift; a +// filter that silently matched nothing would look like "no results for that +// language", which is a plausible answer and therefore an invisible bug. +func TestSearchLanguageFilterOnCompactScan(t *testing.T) { + const dim = 512 + s := openStore(t) + ctx := context.Background() + chunks, embs := q8Corpus(t, s, "/lang", 1200, dim) + + r := rand.New(rand.NewSource(3)) + q := randNorm(r, dim) + got, err := s.Search(ctx, "/lang", q, 10, map[string]string{"language": "rust"}) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(got) == 0 { + t.Fatal("language filter returned nothing") + } + for _, res := range got { + if res.Language != "rust" { + t.Fatalf("language filter leaked a %q result", res.Language) + } + } + if want := exactTopK(chunks, embs, q, 10, "rust"); strings.Join(resultKeys(got), ",") != strings.Join(want, ",") { + t.Errorf("filtered compact scan disagrees with an exact filtered scan\n got: %v\nwant: %v", + resultKeys(got), want) + } +} + +// TestSearchUnsupportedFilterFallsBack covers the other half of q8Filterable: +// a filter the compact copy cannot express must take the float32 scan, which +// has every column, rather than be ignored. +func TestSearchUnsupportedFilterFallsBack(t *testing.T) { + const dim = 256 + s := openStore(t) + ctx := context.Background() + chunks, _ := q8Corpus(t, s, "/filter", 600, dim) + + r := rand.New(rand.NewSource(5)) + q := randNorm(r, dim) + target := chunks[123].FilePath + got, err := s.Search(ctx, "/filter", q, 10, map[string]string{"file_path": target}) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(got) == 0 { + t.Fatal("file_path filter returned nothing — the fallback did not run") + } + for _, res := range got { + if res.FilePath != target { + t.Fatalf("file_path filter leaked %q", res.FilePath) + } + } +} + +// stripQ8 turns a store back into what an older binary would have left behind: +// float32 vectors, no compact copy, no completion flag. +func stripQ8(t *testing.T, s *Store) { + t.Helper() + for _, stmt := range []string{`DELETE FROM vectors_q8`, `DELETE FROM q8_state`} { + if _, err := s.db.Exec(stmt); err != nil { + t.Fatalf("%s: %v", stmt, err) + } + } + s.q8Mu.Lock() + s.q8State = map[int64]bool{} + s.q8Mu.Unlock() +} + +// TestSearchWithoutCompactCopy is the guarantee that makes the backfill safe +// to run in the background: a collection with no compact copy answers the same +// queries, correctly, the slow way. +func TestSearchWithoutCompactCopy(t *testing.T) { + const dim = 512 + s := openStore(t) + ctx := context.Background() + chunks, embs := q8Corpus(t, s, "/nocopy", 800, dim) + + r := rand.New(rand.NewSource(21)) + q := randNorm(r, dim) + before, err := s.Search(ctx, "/nocopy", q, 10, nil) + if err != nil { + t.Fatalf("search: %v", err) + } + + stripQ8(t, s) + after, err := s.Search(ctx, "/nocopy", q, 10, nil) + if err != nil { + t.Fatalf("search after strip: %v", err) + } + if len(after) == 0 { + t.Fatal("no results without the compact copy") + } + if strings.Join(resultKeys(after), ",") != strings.Join(exactTopK(chunks, embs, q, 10, ""), ",") { + t.Errorf("fallback scan disagrees with an exact scan: %v", resultKeys(after)) + } + if strings.Join(resultKeys(before), ",") != strings.Join(resultKeys(after), ",") { + t.Errorf("compact and fallback scans disagree\ncompact: %v\nfallback: %v", + resultKeys(before), resultKeys(after)) + } +} + +// TestBackfillConvertsAnOldStore walks the upgrade path end to end: a store +// with no compact copy is opened, the background pass converts it, and +// searches then take the fast path and still agree with an exact scan. +func TestBackfillConvertsAnOldStore(t *testing.T) { + const dim = 512 + dir := t.TempDir() + ctx := context.Background() + + s, err := Open(dir) + if err != nil { + t.Fatalf("open: %v", err) + } + chunks, embs := q8Corpus(t, s, "/old", 900, dim) + stripQ8(t, s) + if err := s.Close(); err != nil { + t.Fatalf("close: %v", err) + } + + s2, err := Open(dir) + if err != nil { + t.Fatalf("reopen: %v", err) + } + defer s2.Close() + + collID, ok, err := s2.collectionID(ctx, collectionName("/old")) + if err != nil || !ok { + t.Fatalf("collection id: %v ok=%v", err, ok) + } + deadline := time.Now().Add(30 * time.Second) + for !s2.q8Ready(ctx, collID) { + if time.Now().After(deadline) { + t.Fatal("backfill did not finish within 30s") + } + time.Sleep(20 * time.Millisecond) + } + + var nVec, nQ8 int + if err := s2.db.QueryRow(`SELECT COUNT(*) FROM vectors`).Scan(&nVec); err != nil { + t.Fatal(err) + } + if err := s2.db.QueryRow(`SELECT COUNT(*) FROM vectors_q8`).Scan(&nQ8); err != nil { + t.Fatal(err) + } + if nVec != nQ8 { + t.Fatalf("backfill left %d of %d vectors unconverted", nVec-nQ8, nVec) + } + + r := rand.New(rand.NewSource(31)) + q := randNorm(r, dim) + got, err := s2.Search(ctx, "/old", q, 10, nil) + if err != nil { + t.Fatalf("search: %v", err) + } + if want := exactTopK(chunks, embs, q, 10, ""); strings.Join(resultKeys(got), ",") != strings.Join(want, ",") { + t.Errorf("backfilled scan disagrees with an exact scan\n got: %v\nwant: %v", resultKeys(got), want) + } +} + +// TestDeletesReachTheCompactCopy checks the direction that fails silently. +// +// A leftover q8 row is a document the scan keeps shortlisting and the rescore +// can no longer score, so it vanishes from results without anything logging a +// word — and it also keeps occupying disk that the "reclaimed" number says was +// freed. +func TestDeletesReachTheCompactCopy(t *testing.T) { + const dim = 256 + s := openStore(t) + ctx := context.Background() + chunks, _ := q8Corpus(t, s, "/del", 400, dim) + + victim := chunks[7].FilePath + if err := s.DeleteByFile(ctx, "/del", victim); err != nil { + t.Fatalf("delete by file: %v", err) + } + var orphans int + if err := s.db.QueryRow(` + SELECT COUNT(*) FROM vectors_q8 q + WHERE NOT EXISTS (SELECT 1 FROM vectors v + WHERE v.collection_id = q.collection_id AND v.doc_id = q.doc_id)`). + Scan(&orphans); err != nil { + t.Fatal(err) + } + if orphans != 0 { + t.Errorf("delete-by-file left %d orphaned rows in the compact copy", orphans) + } + + if err := s.DeleteCollection("/del"); err != nil { + t.Fatalf("delete collection: %v", err) + } + var left, states int + if err := s.db.QueryRow(`SELECT COUNT(*) FROM vectors_q8`).Scan(&left); err != nil { + t.Fatal(err) + } + if err := s.db.QueryRow(`SELECT COUNT(*) FROM q8_state`).Scan(&states); err != nil { + t.Fatal(err) + } + if left != 0 || states != 0 { + t.Errorf("delete-collection left %d compact rows and %d state rows", left, states) + } +} + +// TestQuantizeRoundTrip pins the encoding itself, away from any database. +func TestQuantizeRoundTrip(t *testing.T) { + r := rand.New(rand.NewSource(4)) + for _, dim := range []int{1, 7, 256, 2048} { + v := randNorm(r, dim) + blob, scale := quantizeInt8(v) + if len(blob) != dim { + t.Fatalf("dim %d: blob is %d bytes, want one per component", dim, len(blob)) + } + var maxErr float64 + for i, b := range blob { + got := float64(int8(b)) * float64(scale) + if e := math.Abs(got - float64(v[i])); e > maxErr { + maxErr = e + } + } + // Half a quantisation step is the theoretical bound for round-to- + // nearest; anything above it means the scale or the rounding is wrong. + if bound := float64(scale)/2 + 1e-9; maxErr > bound { + t.Errorf("dim %d: worst component error %g exceeds half a step (%g)", dim, maxErr, bound) + } + } + + // The zero vector must not produce a NaN scale or a panic: it scores 0 + // against everything, which is what the float32 dot product also gives it. + blob, scale := quantizeInt8(make([]float32, 16)) + if scale != 0 || len(blob) != 16 { + t.Errorf("zero vector: scale=%v len=%d, want 0 and 16", scale, len(blob)) + } +} + +// TestDotInt8MatchesFloat checks the integer dot product against the float one +// on the same quantised values, so a mistake in the unrolled loop cannot hide +// behind quantisation error. +func TestDotInt8MatchesFloat(t *testing.T) { + r := rand.New(rand.NewSource(8)) + for _, dim := range []int{3, 4, 5, 64, 2048} { + a, _ := quantizeInt8(randNorm(r, dim)) + b, _ := quantizeInt8(randNorm(r, dim)) + var want int32 + for i := range a { + want += int32(int8(a[i])) * int32(int8(b[i])) + } + if got := dotInt8(a, b); got != want { + t.Errorf("dim %d: dotInt8 = %d, want %d", dim, got, want) + } + } + if got := dotInt8(make([]byte, 4), make([]byte, 5)); got != 0 { + t.Errorf("length mismatch returned %d, want 0", got) + } +} + +// TestScanQuantOffThenOn is the toggle nobody tests until it corrupts +// something. +// +// Turning the compact copy off has to be more than "stop reading it": a store +// that keeps its completion flag while writing rows the copy never sees is a +// store that, once the knob comes back on, answers searches from a copy +// missing everything written in between. Nothing errors, nothing logs — the +// results are just quietly incomplete, which is the failure mode that survives +// review. +func TestScanQuantOffThenOn(t *testing.T) { + const dim = 256 + dir := t.TempDir() + ctx := context.Background() + + on, err := OpenWith(Options{Dir: dir, ScanQuant: true}) + if err != nil { + t.Fatalf("open: %v", err) + } + firstHalf, embs1 := q8Corpus(t, on, "/toggle", 200, dim) + if err := on.Close(); err != nil { + t.Fatalf("close: %v", err) + } + + // Second half written with the copy disabled. + off, err := OpenWith(Options{Dir: dir, ScanQuant: false}) + if err != nil { + t.Fatalf("reopen off: %v", err) + } + r := rand.New(rand.NewSource(77)) + secondHalf := make([]Chunk, 200) + embs2 := make([][]float32, 200) + for i := range secondHalf { + secondHalf[i] = Chunk{ + Content: fmt.Sprintf("late %d", i), + FilePath: fmt.Sprintf("late/f%04d.go", i), + StartLine: i*10 + 1, + EndLine: i*10 + 9, + ChunkType: "function", + SymbolName: fmt.Sprintf("Late%04d", i), + Language: "go", + } + embs2[i] = randNorm(r, dim) + } + if err := off.UpsertChunks(ctx, "/toggle", secondHalf, embs2); err != nil { + t.Fatalf("upsert while off: %v", err) + } + if err := off.Close(); err != nil { + t.Fatalf("close off: %v", err) + } + + back, err := OpenWith(Options{Dir: dir, ScanQuant: true}) + if err != nil { + t.Fatalf("reopen on: %v", err) + } + defer back.Close() + + collID, ok, err := back.collectionID(ctx, collectionName("/toggle")) + if err != nil || !ok { + t.Fatalf("collection id: %v ok=%v", err, ok) + } + deadline := time.Now().Add(30 * time.Second) + for !back.q8Ready(ctx, collID) { + if time.Now().After(deadline) { + t.Fatal("backfill did not re-cover the collection within 30s") + } + time.Sleep(20 * time.Millisecond) + } + + allChunks := append(append([]Chunk{}, firstHalf...), secondHalf...) + allEmbs := append(append([][]float32{}, embs1...), embs2...) + + // Query at one of the vectors written while the copy was off: if that + // window were lost, this is what would go missing. + q := allEmbs[len(embs1)+5] + got, err := back.Search(ctx, "/toggle", q, 10, nil) + if err != nil { + t.Fatalf("search: %v", err) + } + if want := exactTopK(allChunks, allEmbs, q, 10, ""); strings.Join(resultKeys(got), ",") != strings.Join(want, ",") { + t.Errorf("rows written while the compact copy was off are missing from search\n got: %v\nwant: %v", + resultKeys(got), want) + } +} + +// TestQ8FilterableCoversEveryFilter is a canary, not an invariant. +// +// The compact copy carries one metadata column, so exactly one `where` key can +// be answered from it and every other key silently costs 3.4x more bytes per +// vector. That trade is fine as long as somebody chose it. What this catches is +// nobody choosing: a new filterable column added to whereColumns, wired through +// the HTTP layer, and never considered here — after which large collections +// quietly go back to the float32 scan whenever that filter is used. +// +// If this fails, the fix is a decision, not a rubber stamp: either add the +// column to vectors_q8 and to scanQ8, or add it to the list below with a note +// saying the fallback is acceptable for it. +func TestQ8FilterableCoversEveryFilter(t *testing.T) { + fast := map[string]bool{"language": true} + + for key := range whereColumns { + got := q8Filterable(map[string]string{key: "x"}) + if got != fast[key] { + t.Errorf("q8Filterable(%q) = %v, want %v — a filter column changed and "+ + "nobody decided whether the compact scan should carry it", key, got, fast[key]) + } + } + + // An unknown key with an empty value is dropped by buildWhere (chromem + // parity: "" == ""), so it must not disqualify the fast path either. + if !q8Filterable(map[string]string{"nonsense": ""}) { + t.Error("an unknown key with an empty value should not force the exact scan") + } + // An unknown key with a value matches nothing at all; Search short-circuits + // before either scan, so which path it would have taken is moot — but it + // must not be reported as fast-path-able. + if q8Filterable(map[string]string{"nonsense": "x"}) { + t.Error("an unknown key with a value should not be treated as compact-scannable") + } +} + +// currentQ8Mismatches returns the doc_ids whose compact row disagrees with the +// float32 vector it is supposed to be derived from, plus the count of compact +// rows with no float32 row at all. +// +// Both are invisible in normal operation, which is why they are worth asserting +// directly. An orphan is shortlisted by every scan and dropped by every +// rescore, so it costs a result slot and says nothing. A stale row scores its +// document with a vector the document no longer has — it does not disappear, +// it ranks wrong. +func currentQ8Mismatches(t *testing.T, s *Store) (stale []string, orphans int) { + t.Helper() + rows, err := s.db.Query(` + SELECT q.doc_id, q.scale, q.embedding, v.embedding + FROM vectors_q8 q + LEFT JOIN vectors v ON v.collection_id = q.collection_id AND v.doc_id = q.doc_id`) + if err != nil { + t.Fatalf("join: %v", err) + } + defer rows.Close() + for rows.Next() { + var ( + docID string + scale float64 + q8, f32b []byte + ) + if err := rows.Scan(&docID, &scale, &q8, &f32b); err != nil { + t.Fatalf("scan: %v", err) + } + if f32b == nil { + orphans++ + continue + } + vec, _ := blobFloats(f32b, nil) + wantBlob, wantScale := quantizeInt8(vec) + if float32(scale) != wantScale || string(q8) != string(wantBlob) { + stale = append(stale, docID) + } + } + if err := rows.Err(); err != nil { + t.Fatalf("rows: %v", err) + } + return stale, orphans +} + +// TestBackfillSurvivesConcurrentWrites is the regression for the gap between +// the backfill's read and its write. +// +// The backfill reads a batch of float32 vectors, quantises them in Go, and +// writes the results in a separate transaction. It deliberately holds no lock +// across that gap — the file watcher reindexing a file somebody just saved must +// not queue behind a background job — so anything can commit in between, and on +// the load-test fixture the gap is open for 245 seconds of live server. +// +// Two things go wrong without the WHERE EXISTS and DO NOTHING clauses in +// backfillQ8SQL, and neither surfaces as an error: +// +// - a document deleted in the gap gets its compact row reinserted, and +// nothing can ever remove it again — DeleteByFile finds doc_ids through +// `vectors`, where the row no longer is; +// - a document re-embedded in the gap has its fresh compact row overwritten +// by this batch's quantisation of the embedding it just replaced. +// +// The assertions are invariants rather than an expected interleaving, so this +// test can only fail for a real reason, whatever the scheduler does. +func TestBackfillSurvivesConcurrentWrites(t *testing.T) { + const ( + dim = 256 + files = 40 + per = 50 + ) + ctx := context.Background() + s := openStore(t) + + r := rand.New(rand.NewSource(17)) + writeFile := func(file string, seed int64) { + fr := rand.New(rand.NewSource(seed)) + chunks := make([]Chunk, per) + embs := make([][]float32, per) + for i := range chunks { + chunks[i] = Chunk{ + Content: fmt.Sprintf("%s chunk %d", file, i), FilePath: file, + StartLine: i*10 + 1, EndLine: i*10 + 9, + ChunkType: "function", SymbolName: fmt.Sprintf("S%03d", i), Language: "go", + } + embs[i] = randNorm(fr, dim) + } + if err := s.UpsertChunks(ctx, "/race", chunks, embs); err != nil { + t.Errorf("upsert %s: %v", file, err) + } + } + for i := 0; i < files; i++ { + writeFile(fmt.Sprintf("src/f%02d.go", i), int64(i)) + } + _ = r + + // Back to what an older binary would have left: float32 only. + stripQ8(t, s) + + collID, ok, err := s.collectionID(ctx, collectionName("/race")) + if err != nil || !ok { + t.Fatalf("collection id: %v ok=%v", err, ok) + } + + // The backfill walks the collection while the watcher churns it. + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + if _, err := s.backfillCollection(ctx, collID); err != nil { + t.Errorf("backfill: %v", err) + } + }() + go func() { + defer wg.Done() + for i := 0; i < files; i += 2 { + file := fmt.Sprintf("src/f%02d.go", i) + if i%4 == 0 { + // Deleted for good — the case that produces orphans. + if err := s.DeleteByFile(ctx, "/race", file); err != nil { + t.Errorf("delete %s: %v", file, err) + } + continue + } + // Re-embedded with different vectors — the case that produces + // stale rows. A different seed means every chunk's embedding, and + // therefore its quantisation, changes. + writeFile(file, int64(1000+i)) + } + }() + wg.Wait() + + stale, orphans := currentQ8Mismatches(t, s) + if orphans != 0 { + t.Errorf("%d compact rows survive with no vector behind them; "+ + "nothing can delete these — DeleteByFile looks up doc_ids through `vectors`", orphans) + } + if len(stale) != 0 { + t.Errorf("%d compact rows hold the quantisation of a replaced embedding, e.g. %q; "+ + "these documents are scored with vectors they no longer have", len(stale), stale[0]) + } + + // And the collection must still be searchable, by whichever path it ended + // up on — a race that leaves it correct but permanently unconverted would + // pass the assertions above and still be a bug worth seeing. + q := randNorm(rand.New(rand.NewSource(5)), dim) + if _, err := s.Search(ctx, "/race", q, 10, nil); err != nil { + t.Fatalf("search after the race: %v", err) + } +} + +// TestBackfillNeverResurrectsOrOverwrites pins the two SQL clauses the test +// above depends on, without needing a race to expose them. If backfillQ8SQL is +// ever simplified back to a plain upsert, this fails deterministically and says +// which half went missing. +func TestBackfillNeverResurrectsOrOverwrites(t *testing.T) { + ctx := context.Background() + s := openStore(t) + chunks, embs := q8Corpus(t, s, "/rules", 4, 64) + collID, ok, err := s.collectionID(ctx, collectionName("/rules")) + if err != nil || !ok { + t.Fatalf("collection id: %v ok=%v", err, ok) + } + var docID string + if err := s.db.QueryRow( + `SELECT doc_id FROM vectors WHERE collection_id = ? LIMIT 1`, collID).Scan(&docID); err != nil { + t.Fatal(err) + } + _ = chunks + _ = embs + + exec := func(docID string, scale float64, blob []byte) { + t.Helper() + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback() //nolint:errcheck // no-op after a successful Commit + if _, err := tx.ExecContext(ctx, backfillQ8SQL, + collID, docID, "go", scale, blob, collID, docID); err != nil { + t.Fatalf("backfill insert: %v", err) + } + if err := tx.Commit(); err != nil { + t.Fatal(err) + } + } + + // DO NOTHING: an existing compact row is the fresher one and must survive. + var before float64 + if err := s.db.QueryRow( + `SELECT scale FROM vectors_q8 WHERE collection_id = ? AND doc_id = ?`, + collID, docID).Scan(&before); err != nil { + t.Fatal(err) + } + exec(docID, before+1, []byte{1, 2, 3}) + var after float64 + if err := s.db.QueryRow( + `SELECT scale FROM vectors_q8 WHERE collection_id = ? AND doc_id = ?`, + collID, docID).Scan(&after); err != nil { + t.Fatal(err) + } + if after != before { + t.Errorf("backfill overwrote a compact row written by the upsert path (scale %v -> %v)", before, after) + } + + // WHERE EXISTS: a doc with no vector must not gain one. + exec("ghost-doc-id", 0.5, []byte{9}) + var ghosts int + if err := s.db.QueryRow( + `SELECT COUNT(*) FROM vectors_q8 WHERE doc_id = 'ghost-doc-id'`).Scan(&ghosts); err != nil { + t.Fatal(err) + } + if ghosts != 0 { + t.Error("backfill inserted a compact row for a document that has no vector") + } +} diff --git a/server/internal/vectorstore/search.go b/server/internal/vectorstore/search.go index 38fad26..507ffee 100644 --- a/server/internal/vectorstore/search.go +++ b/server/internal/vectorstore/search.go @@ -19,6 +19,18 @@ import ( // of scanners rather than spawn a hundred threads and a hundred page caches. var scanSlots = make(chan struct{}, max(2, runtime.NumCPU())) +// acquireScanSlot takes one of the process-wide scan slots, returning the +// release function. A cancelled context gives up the wait rather than the +// query: a caller that has already gone away must not hold a scanner. +func acquireScanSlot(ctx context.Context) (func(), error) { + select { + case scanSlots <- struct{}{}: + return func() { <-scanSlots }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + // scanSQL streams one collection. // // INDEXED BY is not an optimisation hint, it is a guarantee, and the choice of @@ -36,16 +48,30 @@ var scanSlots = make(chan struct{}, max(2, runtime.NumCPU())) // index costs 1.8x because its keys are ordered by file_path, scattering the // lookups across the collection's whole rowid span. // TestScanUsesCollectionIndex pins the plan. -const scanSQL = `SELECT rowid, embedding FROM vectors INDEXED BY idx_vec_coll WHERE collection_id = ?` +const scanSQL = `SELECT doc_id, embedding FROM vectors INDEXED BY idx_vec_coll WHERE collection_id = ?` + +// scanQ8SQL is the same walk over the compact copy, and it is the one a search +// normally takes. Same INDEXED BY guarantee, same reason: idx_q8_coll's keys +// are (collection_id, rowid), so it visits only this collection's rows and +// yields them in table order. +// +// The row it reads is ~2.1 kB instead of ~8.2 kB, which is the whole point — +// three rows to a leaf page and no overflow chain, against one leaf slice plus +// a dedicated overflow page each. Measured with dbstat: 2731 vs 9216 bytes +// read per vector at 2048 dimensions. +const scanQ8SQL = `SELECT doc_id, scale, embedding FROM vectors_q8 INDEXED BY idx_q8_coll WHERE collection_id = ?` + +// rescoreSQL reads the exact vectors of the shortlist. +const rescoreSQL = `SELECT doc_id, embedding FROM vectors WHERE collection_id = ? AND doc_id IN (%s)` // hydrateSQL fetches the metadata and chunk text of the winners only. The // LEFT JOIN keeps a result whose content row is somehow missing (which should // be impossible — both are written in one transaction) instead of dropping it. -const hydrateSQL = `SELECT v.rowid, v.file_path, v.start_line, v.end_line, +const hydrateSQL = `SELECT v.doc_id, v.file_path, v.start_line, v.end_line, v.chunk_type, v.symbol_name, v.language, COALESCE(c.content, '') FROM vectors v LEFT JOIN vector_contents c ON c.collection_id = v.collection_id AND c.doc_id = v.doc_id - WHERE v.rowid IN (%s)` + WHERE v.collection_id = ? AND v.doc_id IN (%s)` // whereColumns maps chromem metadata keys to their SQL column. start_line and // end_line are integers in the schema but were strings in chromem's metadata, @@ -122,31 +148,120 @@ func (s *Store) Search(ctx context.Context, projectPath string, queryEmbedding [ q = normalizeVector(q) } + best, err := s.rank(ctx, collID, q, limit, where, clauses, args) + if err != nil { + return nil, fmt.Errorf("vectorstore search: %w", err) + } + if len(best) == 0 { + return nil, nil + } + return s.hydrate(ctx, collID, best) +} + +// q8Shortlist is how many candidates the compact scan hands to the rescorer. +// +// The int8 ranking is not the answer, it is a filter: it puts the right +// documents in the shortlist but misorders near-ties, so the shortlist has to +// be wide enough that everything belonging in the top K is inside it. Measured +// on 60k vectors of the fixture's largest collection (ziglang/zig, +// voyage-code-3 @2048) against 50 real query-side embeddings, recall of the +// exact float32 top-K after rescoring: +// +// shortlist k=10 k=20 +// 20 0.998 0.994 +// 40 0.998 0.999 +// 60 1.000 1.000 +// 200 1.000 1.000 +// +// (Without rescoring at all, the int8 ranking alone gives 0.994 at both k.) +// 60 is where both columns reach 1.000, so the floor is 64 and the multiple is +// 4x for larger k. The cost of a wider shortlist is one float32 row each — +// 9 kB — against a scan that just read thousands of times that, which is why +// the floor is generous rather than tight. +// +// A fixed width cannot be exact in the worst case, and it is worth stating +// which case that is: topK rejects boundary ties strictly, so a collection +// holding more than `shortlist` documents within one quantisation step of each +// other — a file vendored a hundred times, the near-duplicate clusters +// q8Corpus models — truncates the tie in scan order, and the rescore cannot +// recover a document that never reached it. Extending the shortlist to +// swallow ties at the boundary would close it. No corpus measured so far has +// needed that, and the documentation says "measured 1.000", not "exact", +// because of it. +func q8Shortlist(limit int) int { + if n := 4 * limit; n > 64 { + return n + } + return 64 +} + +// q8Filterable reports whether the compact scan can answer this filter. +// +// vectors_q8 carries one metadata column, language, because that is the only +// filter any caller actually produces (fetchVectorResults in the HTTP layer, +// from the `languages` query parameter). Anything else — a file_path or +// symbol_name filter, reachable through the Go API but not through HTTP — +// falls back to the float32 scan, which has every column. Slower and correct +// beats fast and wrong. +func q8Filterable(where map[string]string) bool { + for k, v := range where { + if k == "language" { + continue + } + // An unknown key with an empty value matches everything and is dropped + // by buildWhere, so it does not disqualify the fast path. + if _, known := whereColumns[k]; !known && v == "" { + continue + } + return false + } + return true +} + +// rank produces the final ordered candidates, by whichever route this +// collection supports. +func (s *Store) rank(ctx context.Context, collID int64, q []float32, limit int, + where map[string]string, clauses []string, args []any) ([]candidate, error) { + + if !q8Filterable(where) { + // Correct, and quietly ~3.4x more expensive per vector. Said out loud + // because the way this gets slow is a new filter key appearing in the + // HTTP layer: nothing breaks, nothing errors, large collections just + // go back to reading 9 kB per vector. TestQ8FilterableCoversEveryFilter + // is the compile-time half of the same guard. + s.logger.Debug("vectorstore: filter not supported by the compact scan, using the exact one", + "collection_id", collID, "filter_keys", len(where)) + } + if q8Filterable(where) && s.q8Ready(ctx, collID) { + shortlist, err := s.scanQ8(ctx, collID, q, q8Shortlist(limit), where) + if err != nil { + return nil, err + } + if len(shortlist) == 0 { + return nil, nil + } + return s.rescore(ctx, collID, q, shortlist, limit) + } + query := scanSQL queryArgs := append([]any{collID}, args...) if len(clauses) > 0 { query += " AND " + strings.Join(clauses, " AND ") } - top, err := s.scan(ctx, query, queryArgs, q, limit) if err != nil { - return nil, fmt.Errorf("vectorstore search: %w", err) - } - best := top.sorted() - if len(best) == 0 { - return nil, nil + return nil, err } - return s.hydrate(ctx, best) + return top.sorted(), nil } // scan streams the collection past the dot product, keeping the top K. func (s *Store) scan(ctx context.Context, query string, args []any, q []float32, k int) (*topK, error) { - select { - case scanSlots <- struct{}{}: - defer func() { <-scanSlots }() - case <-ctx.Done(): - return nil, ctx.Err() + release, err := acquireScanSlot(ctx) + if err != nil { + return nil, err } + defer release() rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { @@ -155,18 +270,35 @@ func (s *Store) scan(ctx context.Context, query string, args []any, q []float32, defer rows.Close() top := newTopK(k) + if err := streamExact(rows, q, top); err != nil { + return nil, err + } + return top, nil +} + +// streamExact drives (doc_id, embedding) rows past the exact dot product and +// keeps the top K. Shared by the float32 scan and by the rescore, which are +// the same loop over the same columns with different WHERE clauses — and the +// float32 decoding protocol is exactly the kind of thing that rots when it +// exists in two places. +// +// Both columns are read as RawBytes, which is the whole point: the driver +// hands back a view of its own buffer, valid only until the next Next(). The +// embedding is consumed immediately by the dot product, and the doc_id becomes +// a Go string only for a row that actually enters the heap — K allocations +// over a scan instead of one per row, which at 1.9M rows per workspace query +// is the difference between a few thousand strings and sixty megabytes of +// garbage. +func streamExact(rows *sql.Rows, q []float32, top *topK) error { var ( - rowID int64 + docID sql.RawBytes raw sql.RawBytes scratch []float32 ) dim := len(q) for rows.Next() { - // RawBytes avoids a copy of every 3 kB embedding; it is only valid - // until the next Next(), which is fine because the dot product - // consumes it immediately. - if err := rows.Scan(&rowID, &raw); err != nil { - return nil, err + if err := rows.Scan(&docID, &raw); err != nil { + return err } if len(raw)/4 != dim { // A row from a different embedding model (namespaces are supposed @@ -179,10 +311,130 @@ func (s *Store) scan(ctx context.Context, query string, args []any, q []float32, vec, scratch = blobFloats(raw, scratch) score := dot(q, vec) if top.qualifies(score) { - top.add(candidate{rowID: rowID, score: score}) + top.add(candidate{docID: string(docID), score: score}) } } - return top, rows.Err() + return rows.Err() +} + +// scanQ8 streams the compact copy and returns the shortlist in approximate +// score order. +// +// The scores it produces are NOT returned to anyone: they rank the shortlist +// and are then thrown away by rescore, which recomputes them on the exact +// vectors. That is deliberate — an int8 dot product is a good enough ordering +// to choose 64 documents out of 350,000 and not good enough to be shown as a +// similarity. +func (s *Store) scanQ8(ctx context.Context, collID int64, q []float32, n int, where map[string]string) ([]candidate, error) { + release, err := acquireScanSlot(ctx) + if err != nil { + return nil, err + } + defer release() + + // The query is quantised the same way the stored vectors were, and its + // scale is constant across the scan, so it cancels out of every + // comparison. Only the per-row scale has to be applied. + // + // A zero query (scale 0, every component 0) is NOT short-circuited here. + // It scores every row 0, the heap fills with the first rows it sees, and + // the caller gets `limit` results at score 0 — which is exactly what the + // float32 scan does with the same input. Returning nothing instead would + // be more defensible in isolation and wrong in context: the two paths are + // chosen per collection, so a workspace fan-out would answer the same + // broken query with hits from the collections still on float32 and + // silence from the converted ones. + qq, _ := quantizeInt8(q) + + query := scanQ8SQL + args := []any{collID} + // Presence, not emptiness: chromem compared metadata["language"] against + // the filter value, so {"language": ""} asks for rows whose language is + // empty — a real query, not an absent filter. buildWhere gets this right + // for the float32 path by mapping the key to a column and binding + // whatever value came with it; treating "" as "no filter" here would make + // the two paths disagree on the one filter this one supports. + if language, ok := where["language"]; ok { + query += " AND language = ?" + args = append(args, language) + } + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + top := newTopK(n) + var ( + docID sql.RawBytes + scale float64 + raw sql.RawBytes + ) + dim := len(q) + for rows.Next() { + if err := rows.Scan(&docID, &scale, &raw); err != nil { + return nil, err + } + if len(raw) != dim { + // Same guard as the float32 scan: a row left by a different model. + continue + } + score := float32(scale) * float32(dotInt8(raw, qq)) + if top.qualifies(score) { + // Both columns alias the driver's buffer until the next Next(); + // the string is materialised only for a row that survives. See + // streamExact for why that matters at this row count. + top.add(candidate{docID: string(docID), score: score}) + } + } + if err := rows.Err(); err != nil { + return nil, err + } + return top.sorted(), nil +} + +// rescore recomputes the shortlist's scores on the exact float32 vectors and +// returns the true top `limit`. +// +// This is what makes the compact scan lossless in practice: measured against +// exact search over 50 real queries, the shortlist contained every document of +// the exact top-K, and rescoring restored the order the approximation blurred +// (see q8Shortlist for the table, and for the boundary case a fixed shortlist +// width cannot rule out). +func (s *Store) rescore(ctx context.Context, collID int64, q []float32, shortlist []candidate, limit int) ([]candidate, error) { + top := newTopK(limit) + for start := 0; start < len(shortlist); start += hydrateBatch { + batch := shortlist[start:min(start+hydrateBatch, len(shortlist))] + query, args := docIDInList(rescoreSQL, collID, batch) + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("rescore: %w", err) + } + err = streamExact(rows, q, top) + rows.Close() + if err != nil { + return nil, fmt.Errorf("rescore: %w", err) + } + } + return top.sorted(), nil +} + +// docIDInList fills a %s-templated IN-list with one placeholder per candidate +// and returns the statement with its arguments, collection first. +// +// Shared by the rescore and the hydrate because they ask the same question of +// the same key — "these doc_ids, in this collection" — and both are already +// chunked by hydrateBatch so neither can exceed SQLite's bound-parameter +// ceiling. +func docIDInList(tmpl string, collID int64, batch []candidate) (string, []any) { + placeholders := make([]string, len(batch)) + args := make([]any, 0, len(batch)+1) + args = append(args, collID) + for i, c := range batch { + placeholders[i] = "?" + args = append(args, c.docID) + } + return fmt.Sprintf(tmpl, strings.Join(placeholders, ",")), args } // hydrateBatch bounds the IN-list so a caller asking for an enormous limit @@ -191,18 +443,18 @@ const hydrateBatch = 500 // hydrate fetches metadata and chunk text for the winning rows and returns // them in score order. -func (s *Store) hydrate(ctx context.Context, best []candidate) ([]SearchResult, error) { - byRowID := make(map[int64]SearchResult, len(best)) +func (s *Store) hydrate(ctx context.Context, collID int64, best []candidate) ([]SearchResult, error) { + byDocID := make(map[string]SearchResult, len(best)) for start := 0; start < len(best); start += hydrateBatch { batch := best[start:min(start+hydrateBatch, len(best))] - if err := s.hydrateInto(ctx, batch, byRowID); err != nil { + if err := s.hydrateInto(ctx, collID, batch, byDocID); err != nil { return nil, err } } out := make([]SearchResult, 0, len(best)) for _, c := range best { - r, ok := byRowID[c.rowID] + r, ok := byDocID[c.docID] if !ok { // Deleted between the scan and the hydrate. Dropping it is the // honest answer — the chunk no longer exists. @@ -218,14 +470,9 @@ func (s *Store) hydrate(ctx context.Context, best []candidate) ([]SearchResult, } // hydrateInto reads one batch of winners into dst. -func (s *Store) hydrateInto(ctx context.Context, batch []candidate, dst map[int64]SearchResult) error { - placeholders := make([]string, len(batch)) - args := make([]any, len(batch)) - for i, c := range batch { - placeholders[i] = "?" - args[i] = c.rowID - } - rows, err := s.db.QueryContext(ctx, fmt.Sprintf(hydrateSQL, strings.Join(placeholders, ",")), args...) +func (s *Store) hydrateInto(ctx context.Context, collID int64, batch []candidate, dst map[string]SearchResult) error { + query, args := docIDInList(hydrateSQL, collID, batch) + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { return fmt.Errorf("vectorstore search hydrate: %w", err) } @@ -233,14 +480,14 @@ func (s *Store) hydrateInto(ctx context.Context, batch []candidate, dst map[int6 for rows.Next() { var ( - rowID int64 + docID string r SearchResult ) - if err := rows.Scan(&rowID, &r.FilePath, &r.StartLine, &r.EndLine, + if err := rows.Scan(&docID, &r.FilePath, &r.StartLine, &r.EndLine, &r.ChunkType, &r.SymbolName, &r.Language, &r.Content); err != nil { return fmt.Errorf("vectorstore search hydrate: %w", err) } - dst[rowID] = r + dst[docID] = r } if err := rows.Err(); err != nil { return fmt.Errorf("vectorstore search hydrate: %w", err) diff --git a/server/internal/vectorstore/sqlite.go b/server/internal/vectorstore/sqlite.go index 1521c20..feff5e6 100644 --- a/server/internal/vectorstore/sqlite.go +++ b/server/internal/vectorstore/sqlite.go @@ -84,10 +84,11 @@ const idleConnTimeout = 30 * time.Second // schemaSQL is the whole schema. Two tables, deliberately. // -// `vectors` holds only what a scan reads: the metadata columns the `where` -// filter can constrain and the embedding itself. A row is ~3.2 kB, which fits -// inside an 8 KiB table-leaf cell (the local-payload limit is usable-35), so -// the scan reads two rows per page and never follows an overflow chain. +// `vectors` holds the authoritative float32 embedding and the metadata columns +// the `where` filter can constrain. At 768 dimensions a row is ~3.2 kB and fits +// inside an 8 KiB table-leaf cell (the local-payload limit is usable-35); at +// 2048 it is ~8.2 kB and does not, so every row spills into an overflow page. +// That is why the scan no longer reads this table — see `vectors_q8` below. // // `vector_contents` holds the chunk text. It is stored (duplicating chunks_fts // on disk) so SearchResult.Content behaves identically with no cross-database @@ -98,6 +99,35 @@ const idleConnTimeout = 30 * time.Second // chain. That would roughly double the pages a scan touches. Content is read // only for the K winners, so it costs one extra btree lookup per result. // +// `vectors_q8` is what a search actually scans: the same vectors at one byte +// per component, plus the per-vector scale that undoes the quantisation and +// the one metadata column a search can filter on in practice (language — see +// fetchVectorResults, the only caller that passes a filter). It exists because +// the paragraph above stopped being true once models grew past 1024 +// dimensions: a 2048-dim float32 row is 8.2 kB, which does NOT fit an 8 KiB +// leaf cell, so `vectors` is now exactly the overflow-chain layout that +// splitting out the content was meant to avoid. Measured with dbstat on 400 +// rows, per vector read by a full scan: +// +// 768 float32 4096 B two rows per leaf page +// 1024 float32 8192 B one row per leaf page, half of it air +// 2048 float32 9216 B leaf slice plus a whole overflow page +// 2048 int8 2731 B three rows per leaf page, no overflow +// +// The float32 blob stays in `vectors` and stays authoritative: the scan reads +// q8 to pick a shortlist, then rescores that shortlist against the exact +// vectors, which is what keeps the final ranking identical (see vector.go for +// the recall measurement). q8 rows are therefore derived data — losing them +// costs speed, never answers, which is what lets the backfill run in the +// background while searches fall back to the float32 scan. +// +// `q8_state` records that a collection's q8 rows are complete. Without it, +// "is this collection ready" would be a COUNT(*) over both tables on every +// query — the same mistake that made the stale-FTS probe cost 53 ms per +// workspace search. It carries no dimension: the scan compares each row's blob +// length against the query's own, so a row left by a different model is +// skipped per row rather than gated per collection. +// // Two indexes, and the difference between them matters: // // - idx_vec_coll is (collection_id, rowid) — SQLite appends the rowid to @@ -155,6 +185,19 @@ CREATE TABLE IF NOT EXISTS vectors ( ); CREATE INDEX IF NOT EXISTS idx_vec_coll ON vectors(collection_id); CREATE INDEX IF NOT EXISTS idx_vec_coll_file ON vectors(collection_id, file_path); +CREATE TABLE IF NOT EXISTS vectors_q8 ( + collection_id INTEGER NOT NULL REFERENCES collections(id) ON DELETE CASCADE, + doc_id TEXT NOT NULL, + language TEXT NOT NULL DEFAULT '', + scale REAL NOT NULL, + embedding BLOB NOT NULL, + PRIMARY KEY (collection_id, doc_id) +); +CREATE INDEX IF NOT EXISTS idx_q8_coll ON vectors_q8(collection_id); +CREATE TABLE IF NOT EXISTS q8_state ( + collection_id INTEGER PRIMARY KEY REFERENCES collections(id) ON DELETE CASCADE, + built_at TEXT NOT NULL +); CREATE TABLE IF NOT EXISTS vector_contents ( collection_id INTEGER NOT NULL REFERENCES collections(id) ON DELETE CASCADE, doc_id TEXT NOT NULL, diff --git a/server/internal/vectorstore/store.go b/server/internal/vectorstore/store.go index 3021202..76ec7e5 100644 --- a/server/internal/vectorstore/store.go +++ b/server/internal/vectorstore/store.go @@ -78,6 +78,20 @@ type Options struct { // mapped database pages are clean and reclaimable, but they count in RSS // and every connection maps the file. MMapBytes int64 + // ScanQuant enables the compact int8 copy that searches scan instead of + // the float32 originals (see q8.go). It governs the whole lifecycle, not + // just reading: with it off, the scan takes the float32 path, the backfill + // does not run, and writes DELETE the compact rows of the docs they touch + // and withdraw the collection's completion flag. That last part is what + // makes the switch safe to flip back — leaving rows behind under a live + // flag would mean searching a copy that no longer matches the vectors. + // + // The zero value is false so that a caller constructing Options by hand — + // every test, every tool — gets the plain float32 behaviour unless it asks + // otherwise. Open() is the exception: it is the convenience form and turns + // it on, because a store opened with defaults should behave like the + // server's. + ScanQuant bool // Logger receives migration progress. Defaults to a discarding logger. Logger *slog.Logger } @@ -105,6 +119,19 @@ type Store struct { // only invalidation needed is on delete. collMu sync.Mutex collIDs map[string]int64 + + // q8Mu guards q8State, a cache of collection id -> "the compact scan copy + // is complete". See q8.go; the entry only ever flips one way, so a cached + // true is permanent and a cached false is re-probed. + q8Mu sync.Mutex + q8State map[int64]bool + // scanQuant mirrors Options.ScanQuant. + scanQuant bool + + // stopBG cancels background work (the q8 backfill) on Close. The + // goroutines also go through acquire(), so cancelling is about not doing + // pointless work rather than about safety. + stopBG context.CancelFunc } // ErrClosed is returned by every method once Close has run. @@ -113,7 +140,7 @@ var ErrClosed = errors.New("vectorstore: store is closed") // Open opens (creating if needed) a vector store in the namespace directory // dir, with no legacy import. Kept as the simple form used by tests and tools. func Open(dir string) (*Store, error) { - return OpenWith(Options{Dir: dir}) + return OpenWith(Options{Dir: dir, ScanQuant: true}) } // OpenWith opens a vector store and, when o.LegacyChromaDir holds a chromem-go @@ -139,6 +166,8 @@ func OpenWith(o Options) (*Store, error) { legacyDir: strings.TrimSuffix(filepath.Clean(o.LegacyChromaDir), string(os.PathSeparator)), logger: logger, collIDs: map[string]int64{}, + q8State: map[int64]bool{}, + scanQuant: o.ScanQuant, } if o.LegacyChromaDir == "" { s.legacyDir = "" @@ -147,6 +176,11 @@ func OpenWith(o Options) (*Store, error) { db.Close() return nil, err } + bgCtx, stopBG := context.WithCancel(context.Background()) + s.stopBG = stopBG + if s.scanQuant { + s.startQ8Backfill(bgCtx) + } return s, nil } @@ -156,6 +190,9 @@ func (s *Store) Close() error { if s == nil { return nil } + if s.stopBG != nil { + s.stopBG() + } s.closeMu.Lock() defer s.closeMu.Unlock() if s.closed { @@ -230,7 +267,12 @@ func (s *Store) ensureCollection(ctx context.Context, name string) (int64, error if id, ok, err := s.collectionID(ctx, name); err != nil || ok { return id, err } - if _, err := s.db.ExecContext(ctx, `INSERT OR IGNORE INTO collections(name) VALUES(?)`, name); err != nil { + res, err := s.db.ExecContext(ctx, `INSERT OR IGNORE INTO collections(name) VALUES(?)`, name) + if err != nil { + return 0, fmt.Errorf("vectorstore: create collection %q: %w", name, err) + } + created, err := res.RowsAffected() + if err != nil { return 0, fmt.Errorf("vectorstore: create collection %q: %w", name, err) } id, ok, err := s.collectionID(ctx, name) @@ -240,6 +282,23 @@ func (s *Store) ensureCollection(ctx context.Context, name string) (int64, error if !ok { return 0, fmt.Errorf("vectorstore: collection %q vanished after insert", name) } + if created > 0 && s.scanQuant { + // A collection that has just been created has no vectors, so its + // (empty) q8 side already matches it. Recording that here is what + // makes every collection this binary creates exempt from the backfill: + // upsertBatch writes both tables in one transaction from now on, so + // the property holds by construction. See q8.go. + // + // Logged, not returned: this flag is a performance hint, and its own + // transaction can lose a race for the write lock. Failing the caller + // would abort a whole indexing batch over a row whose absence costs + // nothing but a slower scan — and the absence self-heals, because the + // backfill sets the flag on the next open. + if err := s.markCollectionQ8Ready(ctx, id); err != nil { + s.logger.Warn("vectorstore: could not mark a new collection for the compact scan", + "collection", name, "err", err) + } + } return id, nil } @@ -262,6 +321,19 @@ const upsertContentSQL = `INSERT INTO vector_contents (collection_id, doc_id, co VALUES (?,?,?) ON CONFLICT(collection_id, doc_id) DO UPDATE SET content=excluded.content` +// upsertQ8SQL writes the scan copy in the same transaction as the vector it is +// derived from. Same transaction, not a later pass: a q8 row that disagrees +// with its float32 original would shortlist the wrong documents silently, and +// the only cheap way to guarantee they agree is to make them atomic. +const upsertQ8SQL = `INSERT INTO vectors_q8 (collection_id, doc_id, language, scale, embedding) + VALUES (?,?,?,?,?) + ON CONFLICT(collection_id, doc_id) DO UPDATE SET + language=excluded.language, scale=excluded.scale, embedding=excluded.embedding` + +// deleteQ8DocSQL removes one doc's compact copy. Used by the write path when +// the copy is switched off — see upsertBatch for why not writing is not enough. +const deleteQ8DocSQL = `DELETE FROM vectors_q8 WHERE collection_id = ? AND doc_id = ?` + // ErrCollectionDeleted reports that the collection an upsert was writing into // was deleted while the write was in flight — see UpsertChunks. var ErrCollectionDeleted = errors.New("vectorstore: collection was deleted while the upsert was in flight") @@ -324,6 +396,20 @@ func (s *Store) UpsertChunks(ctx context.Context, projectPath string, chunks []C return err } + // With the compact copy switched off, everything written below leaves it + // stale — so the completion flag comes off BEFORE the first byte lands, + // not after the last. Ordered that way because the failure it prevents is + // a crash mid-write with the flag still set: a collection that says it is + // complete while missing whatever the interrupted run had already written. + // Each batch also deletes the compact rows of the docs it touches, so a + // re-enable rebuilds from the float32 side rather than trusting a copy + // that was left behind. + if !s.scanQuant { + if err := s.clearCollectionQ8Ready(ctx, collID); err != nil { + return err + } + } + for start := 0; start < len(chunks); start += upsertBatchSize { end := min(start+upsertBatchSize, len(chunks)) if err := s.upsertBatch(ctx, collID, chunks[start:end], embeddings[start:end], start); err != nil { @@ -358,6 +444,23 @@ func (s *Store) upsertBatch(ctx context.Context, collID int64, chunks []Chunk, e return err } defer contentStmt.Close() + // One statement per doc either way: write the compact copy, or delete it. + // Deleting matters as much as writing. A doc re-embedded while the copy is + // off would otherwise keep the compact row of its PREVIOUS embedding, and + // the backfill deliberately never overwrites an existing compact row (see + // backfillQ8SQL) — so re-enabling the knob would seal that stale row in + // behind a completion flag, and searches would score the doc with a vector + // it no longer has. + var q8Stmt *sql.Stmt + if s.scanQuant { + q8Stmt, err = tx.PrepareContext(ctx, upsertQ8SQL) + } else { + q8Stmt, err = tx.PrepareContext(ctx, deleteQ8DocSQL) + } + if err != nil { + return err + } + defer q8Stmt.Close() for i, c := range chunks { emb := embeddings[i] @@ -372,6 +475,14 @@ func (s *Store) upsertBatch(ctx context.Context, collID int64, chunks []Chunk, e if _, err := contentStmt.ExecContext(ctx, collID, id, c.Content); err != nil { return err } + if s.scanQuant { + q8, scale := quantizeInt8(emb) + if _, err := q8Stmt.ExecContext(ctx, collID, id, c.Language, scale, q8); err != nil { + return err + } + } else if _, err := q8Stmt.ExecContext(ctx, collID, id); err != nil { + return err + } } return tx.Commit() } @@ -402,6 +513,16 @@ func (s *Store) DeleteByFile(ctx context.Context, projectPath, filePath string) collID, collID, filePath); err != nil { return fmt.Errorf("vectorstore delete contents for %q: %w", filePath, err) } + // Before the vectors themselves: the subquery reads file_path from + // `vectors`, so deleting there first would leave every q8 row of that file + // behind, and an orphan q8 row is a document the scan keeps shortlisting + // and the rescore can no longer score. + if _, err := tx.ExecContext(ctx, `DELETE FROM vectors_q8 + WHERE collection_id = ? AND doc_id IN ( + SELECT doc_id FROM vectors WHERE collection_id = ? AND file_path = ?)`, + collID, collID, filePath); err != nil { + return fmt.Errorf("vectorstore delete q8 for %q: %w", filePath, err) + } if _, err := tx.ExecContext(ctx, `DELETE FROM vectors WHERE collection_id = ? AND file_path = ?`, collID, filePath); err != nil { return fmt.Errorf("vectorstore delete by file %q: %w", filePath, err) diff --git a/server/internal/vectorstore/vector.go b/server/internal/vectorstore/vector.go index 9b82312..b6660cc 100644 --- a/server/internal/vectorstore/vector.go +++ b/server/internal/vectorstore/vector.go @@ -125,11 +125,19 @@ func dot(a, b []float32) float32 { // top-K // --------------------------------------------------------------------------- -// candidate is one scored row. Only the rowid is kept during the scan — the +// candidate is one scored row. Only the doc_id is kept during the scan — the // metadata and the chunk text of the K winners are fetched afterwards, so the // scan never materialises a string it is about to throw away. +// +// doc_id rather than rowid because the identity has to survive a VACUUM. +// `vectors` has a composite PRIMARY KEY, so its rowid is implicit, and SQLite +// only promises to preserve implicit rowids across a VACUUM for tables with an +// INTEGER PRIMARY KEY — everywhere else it may renumber them. That is +// survivable while the rowid never leaves a single query, and fatal once a +// second table (vectors_q8) keys off it: the pairing would silently shift and +// every search would score one document with another one's vector. type candidate struct { - rowID int64 + docID string score float32 } @@ -194,7 +202,7 @@ func (t *topK) down(i int) { } // sorted returns the candidates in descending score order. Ties break on the -// lower rowid so the ordering is deterministic across runs (SQLite may hand +// lower doc_id so the ordering is deterministic across runs (SQLite may hand // equal-scoring rows back in a different physical order after churn). func (t *topK) sorted() []candidate { out := append([]candidate(nil), t.h...) @@ -212,5 +220,106 @@ func less(a, b candidate) bool { if a.score != b.score { return a.score > b.score } - return a.rowID < b.rowID + return a.docID < b.docID +} + +// --------------------------------------------------------------------------- +// int8 quantisation +// +// The scan's cost is the bytes it streams, and at 2048 dimensions a float32 +// embedding is 8 KiB — more than an 8 KiB page can hold beside its own header, +// so every row also drags an overflow page behind it. Measured on the 45-repo +// fixture, one workspace query moves 17.6 GB. +// +// int8 makes that 2 KiB, which packs three rows to a page and reads 2.7 kB per +// vector: 3.4x less I/O and, measured, 3.0x less CPU in the dot product. +// +// The quantisation is per vector, not global. A single outlier component +// anywhere in a corpus would otherwise set the scale for every vector in it +// and crush the resolution of all the ordinary ones. Per-vector costs 8 bytes +// of REAL and makes each vector's own dynamic range the thing being spent. +// +// What this loses, measured against exact float32 ranking over 50 real +// query-side embeddings on the fixture's largest collection (60k vectors of +// ziglang/zig, voyage-code-3 @2048): recall@10 = 0.994 from the int8 ranking +// alone. Rescoring the shortlist on the float32 originals brought it to 1.000 +// at k=10 and k=20 — the approximation misorders near-ties, it does not lose +// the documents, so re-reading a few dozen exact vectors recovered every one. +// That is why the float32 blob stays on disk: it is no longer read by the +// scan, only by the rescore. +// +// Note what that does and does not promise. The SCORES are exact — they come +// from the float32 vectors either way. The SET is an approximation measured at +// zero error here, not proved at zero in general; see q8Shortlist for the +// boundary case it cannot rule out. +// --------------------------------------------------------------------------- + +// int8Max is the quantisation range. -128 is deliberately excluded: keeping +// the range symmetric means scale*q reconstructs -maxAbs and +maxAbs alike, so +// no component's sign carries a different error than its opposite. +const int8Max = 127 + +// quantizeInt8 encodes v as one signed byte per component plus the scale that +// undoes it: v[i] ~= scale * int8(blob[i]). +// +// A zero vector (or one of zero length) yields scale 0, which scores 0 against +// every query — the same answer the float32 dot product gives it. +func quantizeInt8(v []float32) (blob []byte, scale float32) { + if len(v) == 0 { + return nil, 0 + } + var maxAbs float32 + for _, x := range v { + if x < 0 { + x = -x + } + if x > maxAbs { + maxAbs = x + } + } + if maxAbs == 0 { + return make([]byte, len(v)), 0 + } + scale = maxAbs / int8Max + blob = make([]byte, len(v)) + for i, x := range v { + r := float64(x) / float64(scale) + q := math.Round(r) + if q > int8Max { + q = int8Max + } else if q < -int8Max { + q = -int8Max + } + blob[i] = byte(int8(q)) + } + return blob, scale +} + +// dotInt8 is the integer dot product of two quantised vectors. +// +// int32 cannot overflow here: every term is at most 127*127 = 16129, so it +// would take more than 133k dimensions to reach 2^31. Accumulating in int32 +// rather than float32 also means the sum is exact — all the error in this path +// is in the quantisation, none of it in the arithmetic. +// +// Returns 0 for a length mismatch, matching dot(). +func dotInt8(a, b []byte) int32 { + if len(a) != len(b) { + return 0 + } + var s0, s1, s2, s3 int32 + i := 0 + for ; i+4 <= len(a); i += 4 { + x := a[i : i+4 : i+4] + y := b[i : i+4 : i+4] + s0 += int32(int8(x[0])) * int32(int8(y[0])) + s1 += int32(int8(x[1])) * int32(int8(y[1])) + s2 += int32(int8(x[2])) * int32(int8(y[2])) + s3 += int32(int8(x[3])) * int32(int8(y[3])) + } + sum := s0 + s1 + s2 + s3 + for ; i < len(a); i++ { + sum += int32(int8(a[i])) * int32(int8(b[i])) + } + return sum }