From 5ac41181c9d641160892a186d73f6f69d7311d6b Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Thu, 6 Aug 2026 20:21:54 +0200 Subject: [PATCH 1/9] feat(store): share and compress worktree indexes --- .gitattributes | 4 + CLAUDE.md | 17 +- README.md | 43 +- cmd/clean.go | 101 +- cmd/hook.go | 4 +- cmd/index.go | 66 +- cmd/search.go | 5 +- cmd/stdio.go | 83 +- docs/INDEX_STORAGE.md | 175 + e2e_cli_test.go | 2 +- go.mod | 9 +- go.sum | 2 - internal/config/config.go | 66 +- internal/config/config_test.go | 68 + internal/config/service.go | 16 + internal/config/version.go | 3 +- internal/index/index.go | 80 +- internal/index/migrate.go | 136 + internal/index/migrate_test.go | 73 + internal/index/shared.go | 203 + internal/indexlock/lock.go | 39 + internal/indexlock/lock_test.go | 28 + internal/sqlitevec/LICENSE-APACHE | 201 + internal/sqlitevec/LICENSE-MIT | 21 + internal/sqlitevec/lib.go | 32 + internal/sqlitevec/lib_test.go | 29 + internal/sqlitevec/sqlite-vec.c | 10199 ++++++++++++++++++++++++++++ internal/sqlitevec/sqlite-vec.h | 38 + internal/store/hybrid_cte_test.go | 2 +- internal/store/shared.go | 894 +++ internal/store/shared_test.go | 379 ++ internal/store/store.go | 102 +- skills/doctor/SKILL.md | 5 + skills/reindex/SKILL.md | 13 +- 34 files changed, 13001 insertions(+), 137 deletions(-) create mode 100644 docs/INDEX_STORAGE.md create mode 100644 internal/index/migrate.go create mode 100644 internal/index/migrate_test.go create mode 100644 internal/index/shared.go create mode 100644 internal/sqlitevec/LICENSE-APACHE create mode 100644 internal/sqlitevec/LICENSE-MIT create mode 100644 internal/sqlitevec/lib.go create mode 100644 internal/sqlitevec/lib_test.go create mode 100644 internal/sqlitevec/sqlite-vec.c create mode 100644 internal/sqlitevec/sqlite-vec.h create mode 100644 internal/store/shared.go create mode 100644 internal/store/shared_test.go diff --git a/.gitattributes b/.gitattributes index 2d1fe5e1..0165b268 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,6 +1,10 @@ bench-results linguist-generated docs/plans linguist-generated +# Keep the vendored sqlite-vec release source byte-for-byte identical to +# upstream, including its existing whitespace. +internal/sqlitevec/sqlite-vec.c -whitespace linguist-vendored + # Windows batch files require CRLF line endings — the cmd.exe parser # has 512-byte boundary bugs with bare-LF files (GOTO/CALL label # parsing). Pin .bat/.cmd to CRLF regardless of the user's diff --git a/CLAUDE.md b/CLAUDE.md index ec0aed4c..f3dc5039 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -130,6 +130,7 @@ Codex, Cursor, and OpenCode reuse the same repo-root `skills/`, `hooks/`, and | `OLLAMA_HOST` | `http://localhost:11434` | Ollama server URL | | `LM_STUDIO_HOST` | `http://localhost:1234` | LM Studio server URL | | `LUMEN_MAX_CHUNK_TOKENS` | `512` | Max tokens per chunk before splitting | +| `LUMEN_VECTOR_STORAGE` | `int8` | Vector precision (`int8` or `float32`) | ¹ `ordis/jina-embeddings-v2-base-code` (Ollama), `nomic-ai/nomic-embed-code-GGUF` (LM Studio) @@ -157,7 +158,8 @@ Codex, Cursor, and OpenCode reuse the same repo-root `skills/`, `hooks/`, and ├── internal/ │ ├── config/ # Config loading & paths │ ├── index/ # Orchestration (Merkle + embedding + chunking) -│ ├── store/ # SQLite + sqlite-vec operations +│ ├── store/ # Shared SQLite collection + sqlite-vec operations +│ ├── sqlitevec/ # Vendored sqlite-vec v0.1.9 wrapper and sources │ ├── chunker/ # Go AST parsing → chunks │ ├── embedder/ # Ollama/LM Studio HTTP client │ └── merkle/ # Change detection (SHA-256 tree) @@ -189,8 +191,9 @@ because slog writes to the log file while tui writes to the process stderr. ## Key Design Decisions - **Merkle tree for diffs**: Avoid re-indexing unchanged code -- **Model name + IndexVersion in DB path**: Different models or index versions → - separate indexes (SHA-256 hash of path + model name + `IndexVersion`). +- **Repository collection profile in DB path**: Git common directory, indexed + scope, model, dimensions, vector precision, chunking profile, and + `IndexVersion` select a shared content-addressed collection. `IndexVersion` is a hardcoded constant in `internal/config/version.go` — increment it (and document why in the commit message) whenever a chunker, embedder, or index-format change would make existing indexes incompatible. Do @@ -199,13 +202,19 @@ because slog writes to the log file while tui writes to the process stderr. .gitattributes → extension - **Chunk splitting at line boundaries**: Oversized chunks split at `LUMEN_MAX_CHUNK_TOKENS` (512 default) -- **32-batch embedding**: Balance memory vs. API round-trips +- **256-batch embedding**: Only exact embedding inputs missing from the shared + vector table are sent to the backend - **Cosine distance KNN**: Normalized for semantic similarity +- **Vendored sqlite-vec**: v0.1.9 is compiled behind `internal/sqlitevec` so + collection deletion and vector behavior do not drift with system packages - **Plugin system**: Declarative Claude and Cursor packaging at the repo root, plus Codex/OpenCode install surfaces that reuse the same skills and launcher - **No repo-root `.mcp.json`**: Use `mcp.json` for Cursor and `.codex/INSTALL.md` for Codex so Claude project behavior never changes implicitly +See [docs/INDEX_STORAGE.md](docs/INDEX_STORAGE.md) for collection identity, +deduplication, vector precision, migration, cleanup, and status-field semantics. + ## Claude Integration Notes When planning any work related to claude code plugin, marketplace, hooks, diff --git a/README.md b/README.md index 3f2b8c8a..9c873247 100644 --- a/README.md +++ b/README.md @@ -216,8 +216,9 @@ Files → semantic chunks → vector embeddings → SQLite/sqlite-vec → KNN se When Claude needs to understand code, it calls `semantic_search` instead of reading entire files. The index is stored outside your repo -(`~/.local/share/lumen//index.db`), keyed by project path and model name — -different models never share an index. +(`~/.local/share/lumen//index.db`). Git worktrees from the same repository +use one collection for a compatible model, vector-storage, and chunking profile; +non-Git projects use private collections. Different profiles never collide. ## Benchmarks @@ -284,6 +285,7 @@ All configuration is via environment variables: | `OLLAMA_HOST` | `http://localhost:11434` | Ollama server URL | | `LM_STUDIO_HOST` | `http://localhost:1234` | LM Studio server URL | | `LUMEN_MAX_CHUNK_TOKENS` | `512` | Max tokens per chunk before splitting | +| `LUMEN_VECTOR_STORAGE` | `int8` | Vector precision (`int8` or `float32`) | | `LUMEN_EMBED_DIMS` | — | Override embedding dimensions (required for unlisted models) | | `LUMEN_EMBED_CTX` | `8192` (unlisted models) | Override context window length | @@ -387,10 +389,12 @@ Index databases are stored outside your project: ~/.local/share/lumen//index.db ``` -Where `` is derived from the absolute project path, embedding model name, -and binary version. Different models or Lumen versions automatically get -separate indexes. No files are added to your repo, no `.gitignore` modifications -needed. +Where `` identifies the Git common directory (or the absolute path for a +non-Git project), indexed scope, embedding model and dimensions, vector +precision, chunking profile, and index version. Worktrees in one repository +share content-addressed file revisions and vectors while retaining independent +project memberships. Vectors use int8 storage by default; set +`LUMEN_VECTOR_STORAGE=float32` to opt out. No files are added to your repo. You can safely delete the entire `lumen` directory to clear all indexes, or let Lumen reclaim the space for you: @@ -405,11 +409,15 @@ An index counts as used every time Lumen opens it (search, indexing, status, or session start), so indexes for projects you still work on are never removed. Indexes with an indexer currently running are always kept. -**Git worktrees** are detected automatically. When you create a new worktree -(`git worktree add` or `claude --worktree`), Lumen finds a sibling worktree's -existing index and copies it as a seed. The Merkle tree diff then re-indexes -only the files that actually differ — typically a handful of files instead of -the entire codebase. No configuration needed; it just works. +**Git worktrees** are detected automatically. A new worktree attaches unchanged +path-and-content revisions directly from the repository collection and embeds +only missing chunk inputs. Removing an old worktree drops its memberships; +shared revisions and vectors remain until their final reference disappears. +Legacy per-worktree indexes migrate lazily, reusing unchanged float32 vectors +without contacting the embedding backend. + +For the complete storage key, sharing rules, status metrics, migration process, +and cleanup lifecycle, see [Index storage and lifecycle](docs/INDEX_STORAGE.md). ## CLI Reference @@ -471,6 +479,19 @@ to the model, set **Override Domain Type** → **Text Embedding**. Set `LUMEN_EMBED_MODEL` to a model from the supported table above. Each model gets its own database; the old index is not deleted automatically. +Changing `LUMEN_VECTOR_STORAGE`, `LUMEN_EMBED_DIMS`, or `LUMEN_MAX_CHUNK_TOKENS` +also selects a separate collection. Run `lumen clean` after the old profile is +no longer in use if you want to reclaim its disk space. + +**Understanding index size and deduplication** + +Call `index_status` for the project. It reports project-local file and chunk +counts alongside collection-wide unique vectors, shared references, +deduplication ratio, vector precision, database size, and currently reclaimable +SQLite pages. See +[Index storage and lifecycle](docs/INDEX_STORAGE.md#reading-index-status) for +definitions and examples. + **Slow first indexing** The first run embeds every file. Subsequent runs only process changed files diff --git a/cmd/clean.go b/cmd/clean.go index 01685a1d..444f6ff8 100644 --- a/cmd/clean.go +++ b/cmd/clean.go @@ -24,7 +24,6 @@ import ( "github.com/ory/lumen/internal/config" "github.com/ory/lumen/internal/indexlock" "github.com/ory/lumen/internal/store" - "github.com/ory/lumen/internal/tui" "github.com/spf13/cobra" ) @@ -35,7 +34,12 @@ const ( maxCleanDays = 106751 ) -var removeIndexDir = os.RemoveAll +const dailyCleanupInterval = 24 * time.Hour + +var ( + removeIndexDir = os.RemoveAll + cleanupCollectionAt = store.CleanupCollectionAt +) func init() { addCleanFlags(cleanCmd) @@ -46,24 +50,23 @@ func init() { // definition never drifts from what runClean reads. func addCleanFlags(cmd *cobra.Command) { cmd.Flags().Int("days", defaultCleanDays, - "remove indexes not used in the last N days (0 removes every eligible index except those protected by active locks)") + "remove indexes not used in the last N days (0 removes every index that is not currently being written)") } var cleanCmd = &cobra.Command{ Use: "clean", Short: "Remove unused or orphaned lumen indexes", - Long: fmt.Sprintf(`Deletes unused lumen index databases under ~/.local/share/lumen/. + Long: fmt.Sprintf(`Garbage-collects lumen indexes under ~/.local/share/lumen/. -An index is removed when it has not been opened for --days days (default %d), -or when the project it was built for no longer exists — indexes are keyed by -project path, embedding model, and index version, so renamed projects, deleted -checkouts, and abandoned models leave behind data that is never read again. +Shared collections lose project memberships that have not been opened for +--days days (default %d), or whose worktree no longer exists. Unreferenced file +revisions, chunks, and vectors are then deleted and free pages are reclaimed. +Legacy per-project index directories are removed using the same age policy. Indexes written by older binaries that never recorded an access time fall back to their last indexing time; those without any usable timestamp are removed. -Use "lumen clean --days 0" to drop every eligible cached index except those -protected by active locks, and +Use "lumen clean --days 0" to drop every cached index on this host, and "lumen index --force " to rebuild a single project from scratch. Indexes with an indexer currently running are always kept.`, defaultCleanDays), @@ -92,11 +95,10 @@ func runClean(cmd *cobra.Command, _ []string) error { // reported and the sweep continues; the first such failure is returned once // every directory has been considered. func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.Time) error { - progress := tui.NewProgress(stderr) entries, err := os.ReadDir(dataDir) if err != nil { if os.IsNotExist(err) { - progress.Info("No index data found — nothing to clean.") + _, _ = fmt.Fprintln(stderr, "No index data found — nothing to clean.") return nil } return fmt.Errorf("read data dir: %w", err) @@ -104,6 +106,8 @@ func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.T cutoff := now.Add(-time.Duration(days) * 24 * time.Hour) removed, skipped := 0, 0 + projectsRemoved, vectorsRemoved := 0, 0 + var bytesReclaimed int64 var firstErr error for _, entry := range entries { @@ -113,44 +117,68 @@ func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.T continue } hashDir := filepath.Join(dataDir, entry.Name()) - wasRemoved, err := cleanIndex(progress, entry.Name(), hashDir, days, cutoff) + wasRemoved, sharedStats, cleanErr := cleanIndex(stderr, entry.Name(), hashDir, days, cutoff) + projectsRemoved += sharedStats.ProjectsRemoved + vectorsRemoved += sharedStats.VectorsRemoved + bytesReclaimed += sharedStats.BytesReclaimed if wasRemoved { removed++ } else { skipped++ } - if err != nil { - if firstErr == nil { - firstErr = err - } + if cleanErr != nil && firstErr == nil { + firstErr = cleanErr } } _, _ = fmt.Fprintf(stdout, "Removed %d index director%s, skipped %d.\n", removed, pluralY(removed), skipped) + if projectsRemoved > 0 || vectorsRemoved > 0 || bytesReclaimed > 0 { + _, _ = fmt.Fprintf(stdout, "Shared cleanup: %d projects, %d vectors, %d bytes reclaimed.\n", + projectsRemoved, vectorsRemoved, bytesReclaimed) + } return firstErr } -// cleanIndex evaluates and removes one index while holding its writer lock. -func cleanIndex(progress *tui.Progress, name, hashDir string, days int, cutoff time.Time) (bool, error) { +// cleanIndex cleans one legacy index or shared collection while retaining the +// exclusive collection lock for the entire database cleanup and removal. +func cleanIndex(stderr io.Writer, name, hashDir string, days int, cutoff time.Time) (bool, store.CleanupStats, error) { dbPath := filepath.Join(hashDir, "index.db") - lock, err := indexlock.TryAcquire(indexlock.LockPathForDB(dbPath)) - if err != nil || lock == nil { - progress.Info(fmt.Sprintf("Keeping %s: an indexer is currently running.", name)) - return false, nil + lock, lockErr := indexlock.TryAcquire(indexlock.LockPathForDB(dbPath)) + if lockErr != nil || lock == nil { + _, _ = fmt.Fprintf(stderr, "Keeping %s: an indexer is currently running.\n", name) + return false, store.CleanupStats{}, nil } defer lock.Release() + sharedStats, shared, sharedErr := cleanupCollectionAt(dbPath, cutoff) + if shared { + if sharedErr != nil { + _, _ = fmt.Fprintf(stderr, "Failed to clean shared collection %s: %v\n", name, sharedErr) + return false, store.CleanupStats{}, fmt.Errorf("clean shared collection %s: %w", name, sharedErr) + } + if sharedStats.ProjectsLeft > 0 { + _, _ = fmt.Fprintf(stderr, "Cleaned %s: removed %d projects and %d vectors.\n", name, sharedStats.ProjectsRemoved, sharedStats.VectorsRemoved) + return false, sharedStats, nil + } + // Empty collections have no future owner and can be removed as a + // directory, reclaiming sidecars and metadata in one operation. + if err := removeIndexDir(hashDir); err != nil { + return false, sharedStats, fmt.Errorf("remove empty collection %s: %w", hashDir, err) + } + return true, sharedStats, nil + } + stale, reason := isIndexStale(dbPath, days, cutoff) if !stale { - return false, nil + return false, store.CleanupStats{}, nil } if err := removeIndexDir(hashDir); err != nil { - progress.Info(fmt.Sprintf("Failed to remove %s: %v", hashDir, err)) - return false, fmt.Errorf("remove %s: %w", hashDir, err) + _, _ = fmt.Fprintf(stderr, "Failed to remove %s: %v\n", hashDir, err) + return false, store.CleanupStats{}, fmt.Errorf("remove %s: %w", hashDir, err) } - progress.Info(fmt.Sprintf("Removed %s (%s).", name, reason)) - return true, nil + _, _ = fmt.Fprintf(stderr, "Removed %s (%s).\n", name, reason) + return true, store.CleanupStats{}, nil } // isIndexStale reports whether the index at dbPath is no longer worth keeping, @@ -217,3 +245,20 @@ func pluralY(n int) string { } return "ies" } + +// runDailyCleanup performs the MCP-startup maintenance sweep at most once per +// day. The stamp is deliberately outside collection directories so it is not +// mistaken for an index by cleanIndexes. +func runDailyCleanup(dataDir string, now time.Time) { + stampPath := filepath.Join(dataDir, ".last-cleanup") + if info, err := os.Stat(stampPath); err == nil && now.Sub(info.ModTime()) < dailyCleanupInterval { + return + } + if err := os.MkdirAll(dataDir, 0o755); err != nil { + return + } + if err := cleanIndexes(io.Discard, io.Discard, dataDir, defaultCleanDays, now); err != nil { + return + } + _ = os.WriteFile(stampPath, []byte(now.UTC().Format(time.RFC3339)), 0o600) +} diff --git a/cmd/hook.go b/cmd/hook.go index 94e65420..099e7a69 100644 --- a/cmd/hook.go +++ b/cmd/hook.go @@ -155,7 +155,7 @@ func generateSessionContextInternalWithDirective(directive, cwd string, findDono cwd = ancestor } - dbPath := config.DBPathForProject(cwd, modelName) + dbPath := configuredDBPath(cfg, cwd, modelName) if _, err := os.Stat(dbPath); err != nil { // No index yet — kick off background pre-warming so the first search // in this session doesn't pay the full seed + embed cost synchronously. @@ -166,7 +166,7 @@ func generateSessionContextInternalWithDirective(directive, cwd string, findDono return directive + " No index yet — indexing in background." } - s, err := store.New(dbPath, dims) + s, err := store.NewCollection(dbPath, dims, cfg.VectorStorage(), cwd) if err != nil { return directive } diff --git a/cmd/index.go b/cmd/index.go index ef0e9618..2da800bc 100644 --- a/cmd/index.go +++ b/cmd/index.go @@ -221,11 +221,32 @@ func loadConfigWithFlags(cmd *cobra.Command) (*config.ConfigService, error) { // setupIndexer receives dbPath so it is computed exactly once in runIndex. func setupIndexer(cfg *config.ConfigService, emb *embedder.FailoverEmbedder, dbPath string, logger *slog.Logger) (*index.Indexer, error) { - idx, err := index.NewIndexer(dbPath, emb, cfg.MaxChunkTokens()) + return setupIndexerForProject(cfg, emb, dbPath, "", logger) +} + +func configuredDBPath(cfg *config.ConfigService, projectPath, model string) string { + dimensions, known := config.ModelDimensions(model) + if !known { + servers := cfg.Servers() + if len(servers) > 0 && servers[0].Model == model { + dimensions = cfg.ServerDims(0) + } + } + return config.DBPathForProjectProfile(projectPath, model, dimensions, cfg.VectorStorage(), cfg.MaxChunkTokens()) +} + +func setupIndexerForProject(cfg *config.ConfigService, emb *embedder.FailoverEmbedder, dbPath, projectPath string, logger *slog.Logger) (*index.Indexer, error) { + idx, err := index.NewIndexerForProject(dbPath, emb, cfg.MaxChunkTokens(), cfg.VectorStorage(), projectPath) if err != nil { return nil, fmt.Errorf("create indexer: %w", err) } idx.SetLogger(logger) + if projectPath != "" { + legacyPath := config.LegacyDBPathForProject(projectPath, emb.ModelName()) + if err := idx.PrepareLegacyMigration(projectPath, legacyPath); err != nil && logger != nil { + logger.Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + } + } return idx, nil } @@ -233,47 +254,40 @@ func setupIndexer(cfg *config.ConfigService, emb *embedder.FailoverEmbedder, dbP // returns the stats and elapsed time. skipped is true when another indexer holds // the lock. err is nil if indexing was cancelled by a signal. func runIndexer(cmd *cobra.Command, cfg *config.ConfigService, emb *embedder.FailoverEmbedder, projectPath string, p *tui.Progress, logger *slog.Logger) (stats index.Stats, elapsed time.Duration, skipped bool, err error) { - dbPath := config.DBPathForProject(projectPath, emb.ModelName()) + dbPath := configuredDBPath(cfg, projectPath, emb.ModelName()) if mkErr := os.MkdirAll(filepath.Dir(dbPath), 0o755); mkErr != nil { err = fmt.Errorf("create db directory: %w", mkErr) return } - lockPath := indexlock.LockPathForDB(dbPath) - lock, lockErr := indexlock.TryAcquire(lockPath) + collectionLock, lockErr := indexlock.TryAcquireShared(indexlock.LockPathForDB(dbPath)) if lockErr != nil { - err = fmt.Errorf("acquire index lock: %w", lockErr) + err = fmt.Errorf("acquire collection lock: %w", lockErr) return } - if lock == nil { + if collectionLock == nil { skipped = true return } - defer lock.Release() + defer collectionLock.Release() + projectLock, lockErr := indexlock.TryAcquire(indexlock.LockPathForProject(dbPath, projectPath)) + if lockErr != nil { + err = fmt.Errorf("acquire project index lock: %w", lockErr) + return + } + if projectLock == nil { + skipped = true + return + } + defer projectLock.Release() - // Install signal handling before donor seeding: copying a large sibling - // index can take seconds, and cancellation must be able to close/remove the - // seed temp file before the process exits. + // Cancel context on SIGTERM or SIGINT so the indexer stops cleanly and + // the deferred lock releases run before exit. ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT) defer stop() - // Reuse a sibling worktree's index for a brand-new database instead of - // re-embedding from scratch. The index lock serializes CLI indexers, while - // SeedFromDonor's seed lock also serializes this copy with MCP callers. - // A forced rebuild cannot reuse the copied embeddings, so skip the donor - // copy in that mode. force, _ := cmd.Flags().GetBool("force") - if !force { - seedFromDonorIfNew(ctx, dbPath, projectPath, emb.ModelName(), logger, seedOptions{ - status: p.Info, - }) - } - if ctx.Err() != nil { - logger.Info("indexing cancelled by signal", "project", projectPath) - return - } - - idx, setupErr := setupIndexer(cfg, emb, dbPath, logger) + idx, setupErr := setupIndexerForProject(cfg, emb, dbPath, projectPath, logger) if setupErr != nil { err = setupErr return diff --git a/cmd/search.go b/cmd/search.go index 35ca684c..67003b5c 100644 --- a/cmd/search.go +++ b/cmd/search.go @@ -24,7 +24,6 @@ import ( "slices" "time" - "github.com/ory/lumen/internal/config" "github.com/ory/lumen/internal/embedder" "github.com/ory/lumen/internal/index" "github.com/spf13/cobra" @@ -135,11 +134,11 @@ func runSearch(cmd *cobra.Command, args []string) error { tr.record("path resolution", indexRoot) // Span 2: indexer setup - dbPath := config.DBPathForProject(indexRoot, modelName) + dbPath := configuredDBPath(cfg, indexRoot, modelName) if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil { return fmt.Errorf("create db directory: %w", err) } - idx, err := setupIndexer(cfg, emb, dbPath, nil) + idx, err := setupIndexerForProject(cfg, emb, dbPath, indexRoot, nil) if err != nil { return fmt.Errorf("setup indexer: %w", err) } diff --git a/cmd/stdio.go b/cmd/stdio.go index 35114b95..c4ca1e13 100644 --- a/cmd/stdio.go +++ b/cmd/stdio.go @@ -92,13 +92,19 @@ type IndexStatusInput struct { // IndexStatusOutput is the structured output of the index_status tool. type IndexStatusOutput struct { - ProjectPath string `json:"project_path"` - TotalFiles int `json:"total_files"` - IndexedFiles int `json:"indexed_files"` - TotalChunks int `json:"total_chunks"` - LastIndexedAt string `json:"last_indexed_at"` - EmbeddingModel string `json:"embedding_model"` - Stale bool `json:"stale"` + ProjectPath string `json:"project_path"` + TotalFiles int `json:"total_files"` + IndexedFiles int `json:"indexed_files"` + TotalChunks int `json:"total_chunks"` + UniqueVectors int `json:"unique_vectors"` + SharedReferences int `json:"shared_references"` + DeduplicationRate float64 `json:"deduplication_ratio"` + VectorStorage string `json:"vector_storage"` + DatabaseBytes int64 `json:"database_bytes"` + ReclaimableBytes int64 `json:"reclaimable_bytes"` + LastIndexedAt string `json:"last_indexed_at"` + EmbeddingModel string `json:"embedding_model"` + Stale bool `json:"stale"` } // HealthCheckInput defines the parameters for the health_check tool. @@ -164,6 +170,13 @@ func (ic *indexerCache) currentModel() string { return ic.embedder.ModelName() } +func (ic *indexerCache) dbPath(projectPath, model string) string { + if ic.cfg == nil { + return config.DBPathForProject(projectPath, model) + } + return configuredDBPath(ic.cfg, projectPath, model) +} + func (ic *indexerCache) cacheGet(projectPath, model string) (cacheEntry, bool) { if entry, ok := ic.cache[cacheKey(projectPath, model)]; ok { return entry, true @@ -340,7 +353,7 @@ func (ic *indexerCache) findEffectiveRoot(path string, model ...string) string { if _, ok := ic.cacheGet(candidate, modelName); ok { return candidate } - if _, err := os.Stat(config.DBPathForProject(candidate, modelName)); err == nil { + if _, err := os.Stat(ic.dbPath(candidate, modelName)); err == nil { return candidate } } @@ -390,7 +403,7 @@ func (ic *indexerCache) hasIndex(projectPath string, model ...string) bool { if _, ok := ic.cacheGet(projectPath, modelName); ok { return true } - _, err := os.Stat(config.DBPathForProject(projectPath, modelName)) + _, err := os.Stat(ic.dbPath(projectPath, modelName)) return err == nil } @@ -449,7 +462,7 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo // making every first search prohibitively slow. Once an index exists at // the preferred root, subsequent searches reuse it and benefit from the // shared project-wide index. - if _, err := os.Stat(config.DBPathForProject(clean, modelName)); err == nil { + if _, err := os.Stat(ic.dbPath(clean, modelName)); err == nil { effectiveRoot = clean } else { effectiveRoot = ic.findEffectiveRoot(projectPath, modelName) @@ -478,7 +491,7 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo } } - dbPath := config.DBPathForProject(effectiveRoot, modelName) + dbPath := ic.dbPath(effectiveRoot, modelName) if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil { return nil, "", "", fmt.Errorf("create db directory: %w", err) } @@ -503,11 +516,15 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo seed: ic.seedFunc, }) - idx, err := index.NewIndexer(dbPath, ic.embedder, ic.cfg.MaxChunkTokens()) + idx, err := index.NewIndexerForProject(dbPath, ic.embedder, ic.cfg.MaxChunkTokens(), ic.cfg.VectorStorage(), effectiveRoot) if err != nil { return nil, "", "", fmt.Errorf("create indexer: %w", err) } idx.SetLogger(ic.logger()) + legacyPath := config.LegacyDBPathForProject(effectiveRoot, modelName) + if err := idx.PrepareLegacyMigration(effectiveRoot, legacyPath); err != nil { + ic.logger().Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + } // Pre-populate the freshness TTL if the index was recently stamped by // background pre-warming (SessionStart hook). This avoids a redundant @@ -558,7 +575,7 @@ func (ic *indexerCache) handleSemanticSearch(ctx context.Context, req *mcp.CallT progress := buildProgressFunc(ctx, req) - dbPath := config.DBPathForProject(effectiveRoot, modelName) + dbPath := ic.dbPath(effectiveRoot, modelName) out, err := ic.ensureIndexed(idx, input, effectiveRoot, dbPath, progress) if err != nil { return nil, nil, err @@ -763,7 +780,7 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn // If a background indexer holds the exclusive flock, skip EnsureFresh to // avoid duplicating the in-progress Merkle walk. The TOCTOU race is benign: // worst case is redundant work, not corruption (SQLite WAL mode). - if indexlock.IsHeld(indexlock.LockPathForDB(dbPath)) { + if indexlock.IsHeld(indexlock.LockPathForDB(dbPath)) || indexlock.IsHeld(indexlock.LockPathForProject(dbPath, projectDir)) { ic.logger().Info("skipping reindex: background indexer is running", "project", projectDir) out.StaleWarning = staleIndexWarning return out, nil @@ -807,7 +824,8 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn } bgCtx, bgCancel := context.WithTimeout(bgParent, backgroundReindexMaxDuration) - lockPath := indexlock.LockPathForDB(dbPath) + collectionLockPath := indexlock.LockPathForDB(dbPath) + projectLockPath := indexlock.LockPathForProject(dbPath, projectDir) ic.wg.Go(func() { defer bgCancel() defer func() { @@ -816,19 +834,25 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn ic.mu.Unlock() }() - lk, lockErr := indexlock.TryAcquire(lockPath) + collectionLock, lockErr := indexlock.TryAcquireShared(collectionLockPath) if lockErr != nil { ic.logger().Warn("background reindex: failed to acquire lock", "project", projectDir, "err", lockErr) done <- freshResult{skipped: true} return } - if lk == nil { + if collectionLock == nil { // Another process grabbed the lock between our IsHeld check and now. ic.logger().Debug("background reindex: lock held by another process, skipping", "project", projectDir) done <- freshResult{skipped: true} return } - defer lk.Release() + defer collectionLock.Release() + projectLock, lockErr := indexlock.TryAcquire(projectLockPath) + if lockErr != nil || projectLock == nil { + done <- freshResult{skipped: true} + return + } + defer projectLock.Release() // If a recent external process (e.g. lumen index from SessionStart) // already updated the index within freshnessTTL, trust the DB timestamp @@ -1012,12 +1036,18 @@ func (ic *indexerCache) handleIndexStatus(_ context.Context, _ *mcp.CallToolRequ } out := IndexStatusOutput{ - ProjectPath: info.ProjectPath, - TotalFiles: info.TotalFiles, - IndexedFiles: info.IndexedFiles, - TotalChunks: info.TotalChunks, - LastIndexedAt: info.LastIndexedAt, - EmbeddingModel: info.EmbeddingModel, + ProjectPath: info.ProjectPath, + TotalFiles: info.TotalFiles, + IndexedFiles: info.IndexedFiles, + TotalChunks: info.TotalChunks, + UniqueVectors: info.UniqueVectors, + SharedReferences: info.SharedReferences, + DeduplicationRate: info.DeduplicationRate, + VectorStorage: info.VectorStorage, + DatabaseBytes: info.DatabaseBytes, + ReclaimableBytes: info.ReclaimableBytes, + LastIndexedAt: info.LastIndexedAt, + EmbeddingModel: info.EmbeddingModel, } fresh, err := idx.IsFresh(effectiveRoot) @@ -1377,6 +1407,10 @@ func formatIndexStatus(out IndexStatusOutput) string { var b strings.Builder fmt.Fprintf(&b, "Index: %s\n", out.ProjectPath) fmt.Fprintf(&b, "Files: %d | Indexed: %d | Chunks: %d | Model: %s\n", out.TotalFiles, out.IndexedFiles, out.TotalChunks, out.EmbeddingModel) + if out.VectorStorage != "" { + fmt.Fprintf(&b, "Vectors: %d unique | Shared refs: %d | Dedup: %.1f%% | Storage: %s | DB: %d bytes | Reclaimable: %d bytes\n", + out.UniqueVectors, out.SharedReferences, out.DeduplicationRate*100, out.VectorStorage, out.DatabaseBytes, out.ReclaimableBytes) + } stale := "no" if out.Stale { stale = "yes" @@ -1412,6 +1446,7 @@ func runStdio(_ *cobra.Command, _ []string) error { "backend", cfg.Servers()[0].Backend, "freshness_ttl", cfg.FreshnessTTL().String(), ) + runDailyCleanup(filepath.Join(config.XDGDataDir(), "lumen"), time.Now()) closeCtx, closeFn := context.WithCancel(context.Background()) indexers := &indexerCache{ diff --git a/docs/INDEX_STORAGE.md b/docs/INDEX_STORAGE.md new file mode 100644 index 00000000..586b90b9 --- /dev/null +++ b/docs/INDEX_STORAGE.md @@ -0,0 +1,175 @@ +# Index storage and lifecycle + +Lumen stores indexes outside the source repository in repository-scoped, +content-addressed SQLite collections. This lets Git worktrees reuse unchanged +file revisions and embeddings without copying a complete database per worktree. + +## Collection identity + +The default location is: + +```text +~/.local/share/lumen//index.db +``` + +`XDG_DATA_HOME` replaces `~/.local/share` when it is set. The profile hash is +derived from: + +- the resolved Git common directory, or the absolute project path for a non-Git + project; +- the indexed scope; +- embedding model and vector dimensions; +- vector storage (`int8` or `float32`); +- maximum chunk size; and +- Lumen's index-format version. + +All worktrees whose settings resolve to the same profile use one physical +collection. Each worktree still has an independent project membership, Merkle +state, access time, and set of files. A different model, dimension count, vector +precision, chunk size, repository, or format version selects a different +collection automatically. + +The embedding backend is not part of the profile. If the same model name is +served by both Ollama and LM Studio, use distinct configured model names unless +the services return compatible embeddings. + +## What is shared + +Lumen deduplicates at two levels inside a collection: + +1. A file revision is identified by its relative path and content hash. When + another worktree contains the same revision, Lumen attaches it directly. +2. A vector is identified by the SHA-256 hash of the exact embedding input: + + ```text + // + + ``` + +Including the relative path preserves Lumen's existing filepath-aware search +semantics. Identical text at different paths is therefore not assumed to have +the same vector. + +When a file changes, Lumen chunks the new revision, checks which exact inputs +already exist, and sends only missing inputs to the embedding service. Unique +constraints make concurrent worktree indexing safe; two indexers may compute the +same embedding during a race, but only one stored vector remains. + +Search remains project-local. Lumen scans collection vectors adaptively, then +joins candidates through the requested project's memberships and optional path +filter. Results from another worktree are never returned merely because the +underlying storage is shared. + +## Vector storage + +`int8` is the default: + +```bash +export LUMEN_VECTOR_STORAGE=int8 +``` + +Lumen max-absolute-normalizes each vector before quantizing it to signed bytes. +Cosine KNN then operates directly on sqlite-vec's int8 representation. This +substantially reduces vector storage while retaining the vector's direction. + +Use float32 when you need the unquantized representation: + +```bash +export LUMEN_VECTOR_STORAGE=float32 +``` + +Changing this setting creates a separate profile rather than rewriting the +active collection in place. The same applies to the embedding model, dimensions, +and `LUMEN_MAX_CHUNK_TOKENS`. + +## Reading `index_status` + +The MCP `index_status` tool reports: + +| Field | Scope | Meaning | +| --------------------- | --------------- | ------------------------------------------------------------ | +| `total_files` | Current project | Files found by the last completed indexing walk | +| `indexed_files` | Current project | File memberships stored for this project | +| `total_chunks` | Current project | Chunk references reachable by this project | +| `unique_vectors` | Collection | Physical vectors stored once in the collection | +| `shared_references` | Collection | Chunk references beyond the unique-vector count | +| `deduplication_ratio` | Collection | `shared_references / (unique_vectors + shared_references)` | +| `vector_storage` | Collection | `int8` or `float32` | +| `database_bytes` | Collection | Allocated SQLite database pages | +| `reclaimable_bytes` | Collection | Free SQLite pages that incremental vacuum may reclaim | +| `last_indexed_at` | Current project | Last completed indexing timestamp | +| `stale` | Current project | Whether the source tree differs from the stored Merkle state | + +For example, `10,000` unique vectors and `15,000` shared references means +`25,000` total chunk references are represented by `10,000` physical vectors, +for a deduplication ratio of 60%. + +Collection-wide values can stay unchanged after indexing a second worktree even +though that worktree's project-local counts increase. That is the expected sign +that its revisions and vectors were reused. + +## Legacy migration + +Indexes created before shared collections used one float32 database per +worktree. Migration is lazy and automatic when Lumen first opens a project with +the new index format: + +1. Lumen opens the legacy database read-only. +2. It verifies each candidate file still matches its stored content hash. +3. It re-chunks matching files to reconstruct the exact filepath-aware input + hashes. +4. Matching vectors are reused, and quantized when the destination uses int8; + only missing or changed inputs are embedded. +5. The legacy database is removed only after the new project's file hashes have + been verified. + +If legacy recovery cannot be completed, Lumen leaves the source database in +place and embeds the missing inputs normally. + +## Cleanup and disk reclamation + +Run cleanup manually with: + +```bash +lumen clean # memberships unused for 30 days or missing projects +lumen clean --days 7 # use a seven-day inactivity cutoff +lumen clean --days 0 # remove all cached indexes not actively being written +``` + +For shared collections, cleanup first removes stale project memberships, then +garbage-collects file revisions, chunks, and vectors that no remaining project +references. It incrementally vacuums free SQLite pages. The collection directory +is deleted when no projects remain. Legacy index directories use the same age +policy. + +Opening a project for search, indexing, status, or session startup refreshes its +access time. An active indexer lock prevents cleanup from deleting a collection +being written. + +The MCP server also performs this cleanup on startup, throttled to at most once +every 24 hours. The throttle stamp is stored at +`~/.local/share/lumen/.last-cleanup` (or below `XDG_DATA_HOME`). + +`reclaimable_bytes` is a point-in-time estimate from SQLite's freelist. It may +be nonzero until incremental vacuum can truncate pages, and it does not include +WAL or filesystem allocation details. + +## Reindexing versus wiping + +Use the narrowest operation that matches the problem: + +```bash +lumen index . # refresh changed files +lumen index --force . # reprocess every file in the current project +lumen clean # reclaim stale memberships and unreferenced data +lumen clean --days 0 # wipe every cached index on the host +``` + +`--force` does not wipe other worktrees from a shared collection. It rebuilds +the current project's memberships and chunk definitions while the collection +continues to deduplicate physical vectors. Use the full wipe only when you +intend to rebuild all Lumen indexes on the machine. + +To delete indexes manually, stop active Lumen indexers and remove the Lumen data +directory. No source-tree files are stored there, and no files are added to the +project repository. diff --git a/e2e_cli_test.go b/e2e_cli_test.go index 37c7718b..a31fcc11 100644 --- a/e2e_cli_test.go +++ b/e2e_cli_test.go @@ -24,9 +24,9 @@ import ( "strings" "testing" - sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo" _ "github.com/mattn/go-sqlite3" "github.com/ory/lumen/internal/config" + sqlite_vec "github.com/ory/lumen/internal/sqlitevec" ) // gitSampleProject copies testdata/sample-project into a temp directory and diff --git a/go.mod b/go.mod index 9e2b52a3..6ed2603f 100644 --- a/go.mod +++ b/go.mod @@ -2,13 +2,12 @@ module github.com/ory/lumen go 1.25 -require ( - github.com/asg017/sqlite-vec-go-bindings v0.1.6 - github.com/mattn/go-sqlite3 v1.14.34 -) +require github.com/mattn/go-sqlite3 v1.14.34 require ( github.com/alexaandru/go-sitter-forest/dart v1.9.4 + github.com/alexaandru/go-sitter-forest/svelte v1.9.2 + github.com/alexaandru/go-sitter-forest/swift v1.9.5 github.com/bradleyjkemp/cupaloy/v2 v2.8.0 github.com/fsnotify/fsnotify v1.9.0 github.com/gofrs/flock v0.13.0 @@ -30,8 +29,6 @@ require ( atomicgo.dev/cursor v0.2.0 // indirect atomicgo.dev/keyboard v0.2.9 // indirect atomicgo.dev/schedule v0.1.0 // indirect - github.com/alexaandru/go-sitter-forest/svelte v1.9.2 // indirect - github.com/alexaandru/go-sitter-forest/swift v1.9.5 // indirect github.com/clipperhouse/uax29/v2 v2.7.0 // indirect github.com/containerd/console v1.0.5 // indirect github.com/davecgh/go-spew v1.1.1 // indirect diff --git a/go.sum b/go.sum index b74516ce..87fe1a44 100644 --- a/go.sum +++ b/go.sum @@ -21,8 +21,6 @@ github.com/alexaandru/go-sitter-forest/svelte v1.9.2 h1:ixFAFy5oaPU3QU5XeR9tomoR github.com/alexaandru/go-sitter-forest/svelte v1.9.2/go.mod h1:Mfat0bA+ML4NR4cTz7eTz6ziL4KPSQSo2dX80U7GXpM= github.com/alexaandru/go-sitter-forest/swift v1.9.5 h1:CCfvj4BRjvN7HtznqDbgU7ylHHO9ML34ezsJFbErjV0= github.com/alexaandru/go-sitter-forest/swift v1.9.5/go.mod h1:EzSPcZpETNyJIoAyPdbQgFUxWM+vcO3y5eYh8kmNvNc= -github.com/asg017/sqlite-vec-go-bindings v0.1.6 h1:Nx0jAzyS38XpkKznJ9xQjFXz2X9tI7KqjwVxV8RNoww= -github.com/asg017/sqlite-vec-go-bindings v0.1.6/go.mod h1:A8+cTt/nKFsYCQF6OgzSNpKZrzNo5gQsXBTfsXHXY0Q= github.com/atomicgo/cursor v0.0.1/go.mod h1:cBON2QmmrysudxNBFthvMtN32r3jxVRIvzkUiF/RuIk= github.com/bradleyjkemp/cupaloy/v2 v2.8.0 h1:any4BmKE+jGIaMpnU8YgH/I2LPiLBufr6oMMlVBbn9M= github.com/bradleyjkemp/cupaloy/v2 v2.8.0/go.mod h1:bm7JXdkRd4BHJk9HpwqAI8BoAY1lps46Enkdqw6aRX0= diff --git a/internal/config/config.go b/internal/config/config.go index 84becddd..235b0ca7 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -20,6 +20,10 @@ import ( "fmt" "os" "path/filepath" + "strconv" + + "github.com/ory/lumen/internal/git" + "github.com/ory/lumen/internal/models" ) const ( @@ -29,11 +33,8 @@ const ( BackendLMStudio = "lmstudio" ) -// DBPathForProject returns the SQLite database path for a given project, -// derived from a SHA-256 hash of the project path, embedding model name, and -// IndexVersion. Including the model ensures that switching models creates a -// fresh index automatically. Including IndexVersion ensures that incompatible -// chunker/index format changes never share an index with older data. +// DBPathForProject returns the default int8/512-token collection path. New +// runtime code with non-default settings should use DBPathForProjectProfile. func DBPathForProject(projectPath, model string) string { return DBPathForProjectBase(XDGDataDir(), projectPath, model) } @@ -42,7 +43,60 @@ func DBPathForProject(projectPath, model string) string { // using an explicit data directory instead of reading XDG_DATA_HOME from the // environment. Safe to call from parallel goroutines. func DBPathForProjectBase(dataDir, projectPath, model string) string { - hash := fmt.Sprintf("%x", sha256.Sum256([]byte(projectPath+"\x00"+model+"\x00"+IndexVersion))) + dimensions, _ := ModelDimensions(model) + return DBPathForProjectProfileBase(dataDir, projectPath, model, dimensions, "int8", 512) +} + +// ModelDimensions resolves dimensions for a model in the built-in registry. +func ModelDimensions(model string) (int, bool) { + canonical := model + if resolved, ok := models.ModelAliases[model]; ok { + canonical = resolved + } + spec, ok := models.KnownModels[canonical] + return spec.Dims, ok +} + +// DBPathForProjectProfile returns the repository-scoped collection path for an +// exact embedding and chunking profile. Git worktrees sharing a common Git +// directory resolve to the same collection; non-Git projects remain private. +func DBPathForProjectProfile(projectPath, model string, dimensions int, vectorStorage string, maxChunkTokens int) string { + return DBPathForProjectProfileBase(XDGDataDir(), projectPath, model, dimensions, vectorStorage, maxChunkTokens) +} + +// DBPathForProjectProfileBase is DBPathForProjectProfile with an explicit data +// directory, primarily for tests and callers that must not read process-wide +// environment state. +func DBPathForProjectProfileBase(dataDir, projectPath, model string, dimensions int, vectorStorage string, maxChunkTokens int) string { + identity := projectPath + scope := "." + if abs, err := filepath.Abs(projectPath); err == nil { + identity = filepath.Clean(abs) + } + if commonDir, err := git.CommonDir(projectPath); err == nil { + identity = commonDir + // Indexing currently normalizes Git projects to the worktree root. Keep + // scope explicit in the key so subdirectory collections can be added + // without another on-disk format change. + scope = "." + } + profile := identity + "\x00" + scope + "\x00" + model + "\x00" + + strconv.Itoa(dimensions) + "\x00" + vectorStorage + "\x00" + + strconv.Itoa(maxChunkTokens) + "\x00" + IndexVersion + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(profile))) + return filepath.Join(dataDir, "lumen", hash[:16], "index.db") +} + +// LegacyDBPathForProject returns the final one-database-per-worktree path used +// before shared collections. It exists solely for lazy migration. +func LegacyDBPathForProject(projectPath, model string) string { + return LegacyDBPathForProjectBase(XDGDataDir(), projectPath, model) +} + +// LegacyDBPathForProjectBase is LegacyDBPathForProject with an explicit data +// directory. +func LegacyDBPathForProjectBase(dataDir, projectPath, model string) string { + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(projectPath+"\x00"+model+"\x00"+"3"))) return filepath.Join(dataDir, "lumen", hash[:16], "index.db") } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 5843952a..d5b771fe 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -16,6 +16,7 @@ package config import ( "os" + "os/exec" "path/filepath" "strings" "testing" @@ -68,6 +69,44 @@ func TestDBPathForProject(t *testing.T) { }) } +func TestDBPathForProjectProfileSharesGitWorktrees(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git is unavailable") + } + root := filepath.Join(t.TempDir(), "repo") + worktree := filepath.Join(t.TempDir(), "worktree") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatal(err) + } + run := func(dir string, args ...string) { + t.Helper() + cmd := exec.Command("git", args...) + cmd.Dir = dir + cmd.Env = append(os.Environ(), "GIT_AUTHOR_NAME=Lumen", "GIT_AUTHOR_EMAIL=lumen@example.test", "GIT_COMMITTER_NAME=Lumen", "GIT_COMMITTER_EMAIL=lumen@example.test") + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("git %v: %v: %s", args, err, out) + } + } + run(root, "init") + if err := os.WriteFile(filepath.Join(root, "README"), []byte("test"), 0o644); err != nil { + t.Fatal(err) + } + run(root, "add", "README") + run(root, "commit", "-m", "initial") + run(root, "worktree", "add", worktree, "-b", "worktree-test") + + dataDir := t.TempDir() + mainPath := DBPathForProjectProfileBase(dataDir, root, "model", 768, "int8", 512) + worktreePath := DBPathForProjectProfileBase(dataDir, worktree, "model", 768, "int8", 512) + if mainPath != worktreePath { + t.Fatalf("worktrees should share a collection: %q != %q", mainPath, worktreePath) + } + floatPath := DBPathForProjectProfileBase(dataDir, root, "model", 768, "float32", 512) + if floatPath == mainPath { + t.Fatal("vector storage must be part of the collection profile") + } +} + func TestXDGConfigDir(t *testing.T) { t.Run("uses XDG_CONFIG_HOME when set", func(t *testing.T) { t.Setenv("XDG_CONFIG_HOME", "/custom/config") @@ -84,3 +123,32 @@ func TestXDGConfigDir(t *testing.T) { } }) } + +func TestVectorStorageConfiguration(t *testing.T) { + t.Run("defaults to int8", func(t *testing.T) { + t.Setenv("LUMEN_VECTOR_STORAGE", "") + cfg, err := NewConfigService("") + if err != nil { + t.Fatal(err) + } + if got := cfg.VectorStorage(); got != "int8" { + t.Fatalf("VectorStorage() = %q, want int8", got) + } + }) + t.Run("accepts float32 override", func(t *testing.T) { + t.Setenv("LUMEN_VECTOR_STORAGE", "FLOAT32") + cfg, err := NewConfigService("") + if err != nil { + t.Fatal(err) + } + if got := cfg.VectorStorage(); got != "float32" { + t.Fatalf("VectorStorage() = %q, want float32", got) + } + }) + t.Run("rejects unknown storage", func(t *testing.T) { + t.Setenv("LUMEN_VECTOR_STORAGE", "float16") + if _, err := NewConfigService(""); err == nil { + t.Fatal("expected invalid vector storage to fail validation") + } + }) +} diff --git a/internal/config/service.go b/internal/config/service.go index 9eb3f37e..9b408262 100644 --- a/internal/config/service.go +++ b/internal/config/service.go @@ -62,6 +62,7 @@ func defaultServerForBackend(backend string) ServerConfig { func defaultsMap() map[string]any { return map[string]any{ "max_chunk_tokens": 512, + "vector_storage": "int8", "freshness_ttl": "60s", "reindex_timeout": "0s", "log_level": "info", @@ -220,6 +221,9 @@ func applyEnvOverrides(k *koanf.Koanf) { if v := os.Getenv("LUMEN_LOG_LEVEL"); v != "" { globals["log_level"] = v } + if v := os.Getenv("LUMEN_VECTOR_STORAGE"); v != "" { + globals["vector_storage"] = strings.ToLower(v) + } if len(globals) > 0 { _ = k.Load(confmap.Provider(globals, "."), nil) } @@ -303,6 +307,14 @@ func (s *ConfigService) MaxChunkTokens() int { return s.k.Int("max_chunk_tokens") } +// VectorStorage returns the sqlite-vec element type used for persisted +// embeddings. int8 is the default; float32 is retained as an opt-in override. +func (s *ConfigService) VectorStorage() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.k.String("vector_storage") +} + func (s *ConfigService) FreshnessTTL() time.Duration { s.mu.RLock() defer s.mu.RUnlock() @@ -429,6 +441,10 @@ func (s *ConfigService) ServersForModel(model string) ([]int, error) { // other goroutines, or on a temporary ConfigService (as in reload). Must not // be called while holding s.mu — it acquires RLock via Servers(). func (s *ConfigService) validate() error { + storage := s.VectorStorage() + if storage != "int8" && storage != "float32" { + return fmt.Errorf("config: vector_storage must be int8 or float32, got %q", storage) + } servers := s.Servers() if len(servers) == 0 { return fmt.Errorf("config: servers list is empty") diff --git a/internal/config/version.go b/internal/config/version.go index 5d59d5f0..afb8db2f 100644 --- a/internal/config/version.go +++ b/internal/config/version.go @@ -29,4 +29,5 @@ package config // 2 — leading comments included in tree-sitter chunk content; Ruby methods // now produce class-qualified symbols (e.g. Animal.speak) // 3 — Svelte chunker added; .svelte files now indexed via two-phase TS injection -const IndexVersion = "3" +// 4 — repository-scoped content-addressed collections with int8 vectors +const IndexVersion = "4" diff --git a/internal/index/index.go b/internal/index/index.go index b3436378..715dfe84 100644 --- a/internal/index/index.go +++ b/internal/index/index.go @@ -72,12 +72,18 @@ type Stats struct { // StatusInfo holds information about the current index state for a project. type StatusInfo struct { - ProjectPath string - TotalFiles int - IndexedFiles int - TotalChunks int - LastIndexedAt string - EmbeddingModel string + ProjectPath string + TotalFiles int + IndexedFiles int + TotalChunks int + UniqueVectors int + SharedReferences int + DeduplicationRate float64 + VectorStorage string + DatabaseBytes int64 + ReclaimableBytes int64 + LastIndexedAt string + EmbeddingModel string } // Indexer orchestrates chunking, embedding, and storage for a code index. @@ -89,6 +95,10 @@ type Indexer struct { maxChunkTokens int logger *slog.Logger dsn string // path to the SQLite database file; used for corruption recovery + vectorStorage string + projectPath string + legacyVectors map[[32]byte][]float32 + legacySource string } // SetLogger attaches a logger to the indexer for structured diagnostic output. @@ -100,7 +110,19 @@ func (idx *Indexer) SetLogger(l *slog.Logger) { // using the given embedder for vector generation. maxChunkTokens controls // the maximum estimated token count per chunk before splitting; 0 disables splitting. func NewIndexer(dsn string, emb embedder.Embedder, maxChunkTokens int) (*Indexer, error) { - s, err := store.New(dsn, emb.Dimensions()) + return NewIndexerWithStorage(dsn, emb, maxChunkTokens, "int8") +} + +// NewIndexerWithStorage creates a shared-collection indexer with the selected +// sqlite-vec element type. +func NewIndexerWithStorage(dsn string, emb embedder.Embedder, maxChunkTokens int, vectorStorage string) (*Indexer, error) { + return NewIndexerForProject(dsn, emb, maxChunkTokens, vectorStorage, "") +} + +// NewIndexerForProject opens a shared collection with projectPath selected up +// front, allowing metadata such as last_indexed_at to be read before indexing. +func NewIndexerForProject(dsn string, emb embedder.Embedder, maxChunkTokens int, vectorStorage, projectPath string) (*Indexer, error) { + s, err := store.NewCollection(dsn, emb.Dimensions(), vectorStorage, projectPath) if err != nil { return nil, fmt.Errorf("create store: %w", err) } @@ -110,6 +132,8 @@ func NewIndexer(dsn string, emb embedder.Embedder, maxChunkTokens int) (*Indexer chunker: chunker.NewMultiChunker(chunker.DefaultLanguages(maxChunkTokens)), maxChunkTokens: maxChunkTokens, dsn: dsn, + vectorStorage: vectorStorage, + projectPath: projectPath, }, nil } @@ -123,7 +147,7 @@ func (idx *Indexer) rebuildStore() error { _ = os.Remove(idx.dsn + suffix) } } - s, err := store.New(idx.dsn, idx.emb.Dimensions()) + s, err := store.NewCollection(idx.dsn, idx.emb.Dimensions(), idx.vectorStorage, idx.projectPath) if err != nil { return fmt.Errorf("open fresh store: %w", err) } @@ -161,6 +185,9 @@ func (idx *Indexer) Index(ctx context.Context, projectDir string, force bool, pr idx.mu.Lock() defer idx.mu.Unlock() + if err := idx.selectProject(projectDir); err != nil { + return Stats{}, err + } storedHash, err := idx.store.GetMeta("root_hash") if err != nil && err != sql.ErrNoRows { @@ -226,6 +253,9 @@ func (idx *Indexer) EnsureFresh(ctx context.Context, projectDir string, progress idx.mu.Lock() defer idx.mu.Unlock() + if err := idx.selectProject(projectDir); err != nil { + return false, Stats{}, err + } storedHash, err := idx.store.GetMeta("root_hash") if err != nil && err != sql.ErrNoRows { @@ -277,6 +307,9 @@ func (idx *Indexer) EnsureFresh(ctx context.Context, projectDir string, progress // merkle tree, so callers that already have one (e.g. EnsureFresh) do not need // to build it again. func (idx *Indexer) indexWithTree(ctx context.Context, projectDir, oldRootHash string, force bool, curTree *merkle.Tree, progress ProgressFunc) (Stats, error) { + if idx.store.IsShared() { + return idx.indexSharedWithTree(ctx, projectDir, oldRootHash, force, curTree, progress) + } var stats Stats stats.TotalFiles = len(curTree.Files) @@ -537,6 +570,9 @@ func (idx *Indexer) LastIndexedAt() (time.Time, bool) { // IsFresh does not acquire the indexer mutex; it reads through the store's // read-only connection (SQLite WAL isolation). func (idx *Indexer) IsFresh(projectDir string) (bool, error) { + if err := idx.selectProject(projectDir); err != nil { + return false, err + } curTree, err := merkle.BuildTree(projectDir, makeSkip(projectDir)) if err != nil { return false, fmt.Errorf("build merkle tree: %w", err) @@ -559,7 +595,11 @@ func (idx *Indexer) IsFresh(projectDir string) (bool, error) { // Search uses a dedicated read-only database connection so it can execute // concurrently with write operations (e.g. during indexing). It does not // acquire the indexer mutex, relying on SQLite WAL mode for isolation. -func (idx *Indexer) Search(ctx context.Context, _ string, queryVec []float32, limit int, maxDistance float64, pathPrefix string) ([]store.SearchResult, error) { + +func (idx *Indexer) Search(ctx context.Context, projectDir string, queryVec []float32, limit int, maxDistance float64, pathPrefix string) ([]store.SearchResult, error) { + if err := idx.selectProject(projectDir); err != nil { + return nil, err + } return idx.store.Search(ctx, queryVec, limit, maxDistance, pathPrefix) } @@ -571,6 +611,9 @@ func (idx *Indexer) Search(ctx context.Context, _ string, queryVec []float32, li func (idx *Indexer) Status(projectDir string) (StatusInfo, error) { var info StatusInfo info.ProjectPath = projectDir + if err := idx.selectProject(projectDir); err != nil { + return info, err + } storeStats, err := idx.store.Stats() if err != nil { @@ -578,6 +621,14 @@ func (idx *Indexer) Status(projectDir string) (StatusInfo, error) { } info.IndexedFiles = storeStats.TotalFiles info.TotalChunks = storeStats.TotalChunks + info.UniqueVectors = storeStats.UniqueVectors + info.SharedReferences = storeStats.SharedReferences + if references := info.UniqueVectors + info.SharedReferences; references > 0 { + info.DeduplicationRate = float64(info.SharedReferences) / float64(references) + } + info.VectorStorage = storeStats.VectorStorage + info.DatabaseBytes = storeStats.DatabaseBytes + info.ReclaimableBytes = storeStats.ReclaimableBytes meta, err := idx.store.GetMetaBatch([]string{"embedding_model", "last_indexed_at", "total_files"}) if err != nil { @@ -592,6 +643,17 @@ func (idx *Indexer) Status(projectDir string) (StatusInfo, error) { return info, nil } +func (idx *Indexer) selectProject(projectDir string) error { + if projectDir == "" { + projectDir = idx.projectPath + } + if err := idx.store.UseProject(projectDir); err != nil { + return fmt.Errorf("select project: %w", err) + } + idx.projectPath = projectDir + return nil +} + // isBinaryContent reports whether data appears to be binary by checking // for NUL bytes in the first 512 bytes — the same heuristic used by git. func isBinaryContent(data []byte) bool { diff --git a/internal/index/migrate.go b/internal/index/migrate.go new file mode 100644 index 00000000..6a7b6779 --- /dev/null +++ b/internal/index/migrate.go @@ -0,0 +1,136 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +package index + +import ( + "crypto/sha256" + "database/sql" + "encoding/binary" + "encoding/hex" + "fmt" + "math" + "os" + "path/filepath" + + "github.com/ory/lumen/internal/merkle" + "github.com/ory/lumen/internal/store" +) + +// PrepareLegacyMigration recovers vectors for unchanged chunks from a legacy +// per-worktree database. It is safe to ignore a missing source. +func (idx *Indexer) PrepareLegacyMigration(projectDir, legacyPath string) error { + if legacyPath == "" || legacyPath == idx.dsn { + return nil + } + if _, err := os.Stat(legacyPath); err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + db, err := sql.Open("sqlite3", legacyPath) + if err != nil { + return err + } + defer func() { _ = db.Close() }() + if _, err := db.Exec(`PRAGMA query_only=ON`); err != nil { + return err + } + + legacyByChunk := make(map[string][]float32) + rows, err := db.Query(`SELECT c.id, v.embedding FROM chunks c JOIN vec_chunks v ON v.id = c.id`) + if err != nil { + return fmt.Errorf("read legacy vectors: %w", err) + } + for rows.Next() { + var id string + var blob []byte + if err := rows.Scan(&id, &blob); err != nil { + _ = rows.Close() + return err + } + vector, ok := decodeFloat32Vector(blob, idx.emb.Dimensions()) + if ok { + legacyByChunk[id] = vector + } + } + if err := rows.Close(); err != nil { + return err + } + + fileRows, err := db.Query(`SELECT path, hash FROM files WHERE hash <> ''`) + if err != nil { + return fmt.Errorf("read legacy files: %w", err) + } + recovered := make(map[[32]byte][]float32) + for fileRows.Next() { + var relativePath, storedHash string + if err := fileRows.Scan(&relativePath, &storedHash); err != nil { + _ = fileRows.Close() + return err + } + content, err := os.ReadFile(filepath.Join(projectDir, relativePath)) + if err != nil { + continue + } + contentSum := sha256.Sum256(content) + if hex.EncodeToString(contentSum[:]) != storedHash { + continue + } + chunks, err := idx.chunker.Chunk(relativePath, content) + if err != nil { + continue + } + chunks = splitOversizedChunks(chunks, idx.maxChunkTokens) + chunks = mergeUndersizedChunks(chunks) + chunks = splitOversizedChunks(chunks, idx.maxChunkTokens) + for _, chunk := range chunks { + if vector, ok := legacyByChunk[chunk.ID]; ok { + h := sha256.Sum256([]byte(store.EmbeddingInput(chunk))) + recovered[h] = vector + } + } + } + if err := fileRows.Close(); err != nil { + return err + } + idx.legacyVectors = recovered + idx.legacySource = legacyPath + return nil +} + +func decodeFloat32Vector(blob []byte, dimensions int) ([]float32, bool) { + if len(blob) != dimensions*4 { + return nil, false + } + vector := make([]float32, dimensions) + for i := range vector { + bits := binary.LittleEndian.Uint32(blob[i*4 : i*4+4]) + vector[i] = math.Float32frombits(bits) + } + return vector, true +} + +func (idx *Indexer) finishLegacyMigration(tree *merkle.Tree) { + if idx.legacySource == "" { + return + } + hashes, err := idx.store.GetFileHashes() + if err != nil || len(hashes) != len(tree.Files) { + return + } + for path, hash := range tree.Files { + if hashes[path] != hash { + return + } + } + for _, suffix := range []string{"", "-wal", "-shm"} { + _ = os.Remove(idx.legacySource + suffix) + } + _ = os.Remove(filepath.Dir(idx.legacySource)) + idx.legacySource = "" + idx.legacyVectors = nil +} diff --git a/internal/index/migrate_test.go b/internal/index/migrate_test.go new file mode 100644 index 00000000..50bd1009 --- /dev/null +++ b/internal/index/migrate_test.go @@ -0,0 +1,73 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +package index + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "os" + "path/filepath" + "testing" + + "github.com/ory/lumen/internal/store" +) + +func TestLegacyMigrationReusesUnchangedVectors(t *testing.T) { + projectDir := t.TempDir() + content := []byte("package demo\n\nfunc Hello() {}\n") + if err := os.WriteFile(filepath.Join(projectDir, "main.go"), content, 0o644); err != nil { + t.Fatal(err) + } + emb := &mockEmbedder{dims: 4, model: "test-model"} + newPath := filepath.Join(t.TempDir(), "shared.db") + idx, err := NewIndexerForProject(newPath, emb, 0, "int8", projectDir) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + + chunks, err := idx.chunker.Chunk("main.go", content) + if err != nil { + t.Fatal(err) + } + legacyPath := filepath.Join(t.TempDir(), "index.db") + legacy, err := store.New(legacyPath, 4) + if err != nil { + t.Fatal(err) + } + sum := sha256.Sum256(content) + if err := legacy.UpsertFile("main.go", hex.EncodeToString(sum[:])); err != nil { + t.Fatal(err) + } + vectors := make([][]float32, len(chunks)) + for i := range vectors { + vectors[i] = []float32{1, 0, 0, 0} + } + if err := legacy.InsertChunks(chunks, vectors); err != nil { + t.Fatal(err) + } + if err := legacy.Close(); err != nil { + t.Fatal(err) + } + + if err := idx.PrepareLegacyMigration(projectDir, legacyPath); err != nil { + t.Fatal(err) + } + stats, err := idx.Index(context.Background(), projectDir, false, nil) + if err != nil { + t.Fatal(err) + } + if stats.ChunksCreated == 0 { + t.Fatal("expected migrated chunks") + } + if emb.callCount != 0 { + t.Fatalf("legacy vectors should avoid embedding, got %d calls", emb.callCount) + } + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("legacy database should be removed after verification, stat err=%v", err) + } +} diff --git a/internal/index/shared.go b/internal/index/shared.go new file mode 100644 index 00000000..2e78f561 --- /dev/null +++ b/internal/index/shared.go @@ -0,0 +1,203 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +package index + +import ( + "context" + "crypto/sha256" + "fmt" + "os" + "path/filepath" + "slices" + "strconv" + "time" + + "github.com/ory/lumen/internal/merkle" + "github.com/ory/lumen/internal/store" +) + +// indexSharedWithTree indexes one project membership in a repository-scoped +// content-addressed collection. Existing path+content revisions and exact +// embedding inputs are reused without calling the embedding backend. +func (idx *Indexer) indexSharedWithTree(ctx context.Context, projectDir, _ string, force bool, curTree *merkle.Tree, progress ProgressFunc) (Stats, error) { + stats := Stats{TotalFiles: len(curTree.Files)} + oldHashes, err := idx.store.GetFileHashes() + if err != nil { + return stats, fmt.Errorf("get project file hashes: %w", err) + } + for path := range oldHashes { + if !supportedExts[filepath.Ext(path)] { + if err := idx.store.DeleteFileChunks(path); err != nil { + return stats, fmt.Errorf("purge stale file %s: %w", path, err) + } + delete(oldHashes, path) + } + } + + var filesToIndex, filesToRemove []string + if force { + for path := range curTree.Files { + filesToIndex = append(filesToIndex, path) + } + for path := range oldHashes { + if _, ok := curTree.Files[path]; !ok { + filesToRemove = append(filesToRemove, path) + } + } + stats.FilesAdded = len(filesToIndex) + stats.FilesRemoved = len(filesToRemove) + } else { + added, removed, modified := merkle.Diff(&merkle.Tree{Files: oldHashes}, curTree) + filesToIndex = append(filesToIndex, added...) + filesToIndex = append(filesToIndex, modified...) + filesToRemove = removed + stats.FilesAdded = len(added) + stats.FilesModified = len(modified) + stats.FilesRemoved = len(removed) + } + slices.Sort(filesToIndex) + slices.Sort(filesToRemove) + stats.FilesChanged = len(filesToIndex) + len(filesToRemove) + + for _, path := range filesToRemove { + if err := idx.store.DeleteFileChunks(path); err != nil { + return stats, fmt.Errorf("remove project file %s: %w", path, err) + } + } + + if progress != nil { + progress(0, len(filesToIndex), fmt.Sprintf("Found %d files to index", len(filesToIndex))) + } + + for fileIndex, relativePath := range filesToIndex { + if err := ctx.Err(); err != nil { + return stats, err + } + if progress != nil { + progress(fileIndex, len(filesToIndex), fmt.Sprintf("Processing file %d/%d: %s", fileIndex+1, len(filesToIndex), relativePath)) + } + contentHash := curTree.Files[relativePath] + if !force { + attached, err := idx.store.AttachExistingFileRevision(relativePath, contentHash) + if err != nil { + return stats, fmt.Errorf("reuse file revision %s: %w", relativePath, err) + } + if attached { + stats.IndexedFiles++ + continue + } + } + + content, err := os.ReadFile(filepath.Join(projectDir, relativePath)) + if err != nil { + if os.IsPermission(err) { + stats.FilesSkipped++ + continue + } + return stats, fmt.Errorf("read file %s: %w", relativePath, err) + } + if isBinaryContent(content) { + if err := idx.store.DeleteFileChunks(relativePath); err != nil { + return stats, fmt.Errorf("remove binary file %s: %w", relativePath, err) + } + continue + } + + chunks, err := idx.chunker.Chunk(relativePath, content) + if err != nil { + if idx.logger != nil { + idx.logger.Warn("skipping unchunkable file", "path", relativePath, "error", err) + } + stats.FilesSkipped++ + if _, err := idx.store.StoreFileRevision(relativePath, contentHash, nil, nil); err != nil { + return stats, fmt.Errorf("record skipped file %s: %w", relativePath, err) + } + continue + } + chunks = splitOversizedChunks(chunks, idx.maxChunkTokens) + chunks = mergeUndersizedChunks(chunks) + chunks = splitOversizedChunks(chunks, idx.maxChunkTokens) + + var missing []int + if force { + missing = make([]int, len(chunks)) + for i := range chunks { + missing[i] = i + } + } else { + missing, err = idx.store.MissingChunkInputs(chunks) + if err != nil { + return stats, fmt.Errorf("check shared vectors for %s: %w", relativePath, err) + } + } + vectors := make(map[int][]float32, len(missing)) + const embedBatchSize = 256 + var needsEmbedding []int + for _, position := range missing { + h := sha256.Sum256([]byte(store.EmbeddingInput(chunks[position]))) + if vector, ok := idx.legacyVectors[h]; ok { + vectors[position] = vector + } else { + needsEmbedding = append(needsEmbedding, position) + } + } + for start := 0; start < len(needsEmbedding); start += embedBatchSize { + end := min(start+embedBatchSize, len(needsEmbedding)) + positions := needsEmbedding[start:end] + texts := make([]string, len(positions)) + for i, position := range positions { + texts[i] = store.EmbeddingInput(chunks[position]) + } + embedded, err := idx.emb.Embed(ctx, texts) + if err != nil { + return stats, fmt.Errorf("embed %s: %w", relativePath, err) + } + if len(embedded) != len(positions) { + return stats, fmt.Errorf("embed %s returned %d vectors for %d inputs", relativePath, len(embedded), len(positions)) + } + for i, position := range positions { + vectors[position] = embedded[i] + } + if progress != nil { + progress(fileIndex+1, len(filesToIndex), fmt.Sprintf("Embedded %d chunks for %s", len(positions), relativePath)) + } + } + created, err := idx.store.StoreFileRevision(relativePath, contentHash, chunks, vectors) + if err != nil { + return stats, fmt.Errorf("store file revision %s: %w", relativePath, err) + } + if created || force { + stats.ChunksCreated += len(chunks) + } + stats.IndexedFiles++ + } + + if len(filesToIndex) > 0 { + idx.store.Analyze() + } + if err := idx.store.SetMeta("root_hash", curTree.RootHash); err != nil { + return stats, err + } + for key, value := range map[string]string{ + "embedding_model": idx.emb.ModelName(), + "project_path": projectDir, + "last_indexed_at": time.Now().UTC().Format(time.RFC3339), + "total_files": strconv.Itoa(stats.TotalFiles), + "vector_storage": idx.vectorStorage, + } { + if err := idx.store.SetMeta(key, value); err != nil { + return stats, fmt.Errorf("store %s metadata: %w", key, err) + } + } + if progress != nil && len(filesToIndex) > 0 { + progress(len(filesToIndex), len(filesToIndex), fmt.Sprintf("Indexing complete: %d files, %d new chunks", len(filesToIndex), stats.ChunksCreated)) + } + idx.finishLegacyMigration(curTree) + return stats, nil +} diff --git a/internal/indexlock/lock.go b/internal/indexlock/lock.go index ea887661..021f6984 100644 --- a/internal/indexlock/lock.go +++ b/internal/indexlock/lock.go @@ -7,7 +7,9 @@ package indexlock import ( "context" + "crypto/sha256" "errors" + "fmt" "os" "time" @@ -20,6 +22,13 @@ func LockPathForDB(dbPath string) string { return dbPath + ".lock" } +// LockPathForProject returns a stable per-project lock alongside a shared +// collection. Different worktrees therefore do not serialize one another. +func LockPathForProject(dbPath, projectPath string) string { + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(projectPath))) + return dbPath + ".project-" + hash[:16] + ".lock" +} + // Lock is an exclusive advisory lock held on an index lock file. // Release it when indexing is complete. Safe to call Release on a nil Lock. type Lock struct { @@ -64,6 +73,21 @@ func Acquire(ctx context.Context, lockPath string) (*Lock, error) { return &Lock{fl: fl}, nil } +// TryAcquireShared takes a shared lock. Indexers hold this on the collection +// guard while taking an exclusive per-project lock; cleanup takes an exclusive +// probe and therefore never runs during active indexing. +func TryAcquireShared(lockPath string) (*Lock, error) { + fl := flock.New(lockPath) + locked, err := fl.TryRLock() + if err != nil { + return nil, err + } + if !locked { + return nil, nil + } + return &Lock{fl: fl}, nil +} + // IsHeld reports whether another process currently holds an exclusive lock on // lockPath. Returns true on any error (fail-closed: callers skip work rather // than risk concurrent writes). Does NOT create the lock file — if it doesn't @@ -87,6 +111,21 @@ func IsHeld(lockPath string) bool { return false } +// IsAnyHeld reports whether lockPath has either shared or exclusive holders. +// It is used by destructive collection maintenance. +func IsAnyHeld(lockPath string) bool { + if _, err := os.Stat(lockPath); err != nil { + return false + } + fl := flock.New(lockPath) + locked, err := fl.TryLock() + if err != nil || !locked { + return true + } + _ = fl.Unlock() + return false +} + // Release releases the exclusive lock and closes the underlying file. // Safe to call on a nil *Lock. func (l *Lock) Release() { diff --git a/internal/indexlock/lock_test.go b/internal/indexlock/lock_test.go index 56b488b1..7def2479 100644 --- a/internal/indexlock/lock_test.go +++ b/internal/indexlock/lock_test.go @@ -19,6 +19,34 @@ func TestLockPathForDB(t *testing.T) { } } +func TestLockPathForProject(t *testing.T) { + a := indexlock.LockPathForProject("/data/index.db", "/repo/a") + b := indexlock.LockPathForProject("/data/index.db", "/repo/b") + if a == b { + t.Fatal("different projects must have different lock paths") + } + if a != indexlock.LockPathForProject("/data/index.db", "/repo/a") { + t.Fatal("project lock path must be deterministic") + } +} + +func TestSharedCollectionGuardsCanCoexist(t *testing.T) { + path := filepath.Join(t.TempDir(), "collection.lock") + first, err := indexlock.TryAcquireShared(path) + if err != nil || first == nil { + t.Fatalf("first shared lock: lock=%v err=%v", first, err) + } + defer first.Release() + second, err := indexlock.TryAcquireShared(path) + if err != nil || second == nil { + t.Fatalf("second shared lock: lock=%v err=%v", second, err) + } + defer second.Release() + if !indexlock.IsAnyHeld(path) { + t.Fatal("exclusive maintenance probe should see shared holders") + } +} + // TestTryAcquire_Free verifies acquiring a lock on a fresh path succeeds. func TestTryAcquire_Free(t *testing.T) { lockPath := filepath.Join(t.TempDir(), "index.db.lock") diff --git a/internal/sqlitevec/LICENSE-APACHE b/internal/sqlitevec/LICENSE-APACHE new file mode 100644 index 00000000..abbef71f --- /dev/null +++ b/internal/sqlitevec/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2024 Alex Garcia + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/internal/sqlitevec/LICENSE-MIT b/internal/sqlitevec/LICENSE-MIT new file mode 100644 index 00000000..9c106bc4 --- /dev/null +++ b/internal/sqlitevec/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Alex Garcia + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/internal/sqlitevec/lib.go b/internal/sqlitevec/lib.go new file mode 100644 index 00000000..4550f58f --- /dev/null +++ b/internal/sqlitevec/lib.go @@ -0,0 +1,32 @@ +// Package sqlitevec registers the bundled sqlite-vec v0.1.9 extension and +// contains the small amount of vector serialization needed by the store. +package sqlitevec + +// #cgo CFLAGS: -DSQLITE_CORE +// #cgo linux LDFLAGS: -lm +// #include "sqlite-vec.h" +import "C" + +import ( + "bytes" + "encoding/binary" +) + +// Auto registers sqlite-vec for every SQLite connection opened afterward. +func Auto() { + C.sqlite3_auto_extension((*[0]byte)(C.sqlite3_vec_init)) +} + +// Cancel cancels the automatic sqlite-vec extension registration. +func Cancel() { + C.sqlite3_cancel_auto_extension((*[0]byte)(C.sqlite3_vec_init)) +} + +// SerializeFloat32 encodes a vector as sqlite-vec's little-endian float BLOB. +func SerializeFloat32(vector []float32) ([]byte, error) { + buf := new(bytes.Buffer) + if err := binary.Write(buf, binary.LittleEndian, vector); err != nil { + return nil, err + } + return buf.Bytes(), nil +} diff --git a/internal/sqlitevec/lib_test.go b/internal/sqlitevec/lib_test.go new file mode 100644 index 00000000..61514d91 --- /dev/null +++ b/internal/sqlitevec/lib_test.go @@ -0,0 +1,29 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +package sqlitevec + +import ( + "database/sql" + "testing" + + _ "github.com/mattn/go-sqlite3" +) + +func TestBundledVersion(t *testing.T) { + Auto() + db, err := sql.Open("sqlite3", ":memory:") + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + var version string + if err := db.QueryRow(`SELECT vec_version()`).Scan(&version); err != nil { + t.Fatal(err) + } + if version != "v0.1.9" { + t.Fatalf("vec_version() = %q, want v0.1.9", version) + } +} diff --git a/internal/sqlitevec/sqlite-vec.c b/internal/sqlitevec/sqlite-vec.c new file mode 100644 index 00000000..de3176f9 --- /dev/null +++ b/internal/sqlitevec/sqlite-vec.c @@ -0,0 +1,10199 @@ +#include "sqlite-vec.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef SQLITE_VEC_OMIT_FS +#include +#endif + +#ifndef SQLITE_CORE +#include "sqlite3ext.h" +SQLITE_EXTENSION_INIT1 +#else +#include "sqlite3.h" +#endif + +#ifndef UINT32_TYPE +#ifdef HAVE_UINT32_T +#define UINT32_TYPE uint32_t +#else +#define UINT32_TYPE unsigned int +#endif +#endif +#ifndef UINT16_TYPE +#ifdef HAVE_UINT16_T +#define UINT16_TYPE uint16_t +#else +#define UINT16_TYPE unsigned short int +#endif +#endif +#ifndef INT16_TYPE +#ifdef HAVE_INT16_T +#define INT16_TYPE int16_t +#else +#define INT16_TYPE short int +#endif +#endif +#ifndef UINT8_TYPE +#ifdef HAVE_UINT8_T +#define UINT8_TYPE uint8_t +#else +#define UINT8_TYPE unsigned char +#endif +#endif +#ifndef INT8_TYPE +#ifdef HAVE_INT8_T +#define INT8_TYPE int8_t +#else +#define INT8_TYPE signed char +#endif +#endif +#ifndef LONGDOUBLE_TYPE +#define LONGDOUBLE_TYPE long double +#endif + +#ifndef _WIN32 +#ifndef __EMSCRIPTEN__ +#ifndef __COSMOPOLITAN__ +#ifndef __wasi__ +typedef u_int8_t uint8_t; +typedef u_int16_t uint16_t; +typedef u_int64_t uint64_t; +#endif +#endif +#endif +#endif + +typedef int8_t i8; +typedef uint8_t u8; +typedef int16_t i16; +typedef int32_t i32; +typedef sqlite3_int64 i64; +typedef uint32_t u32; +typedef uint64_t u64; +typedef float f32; +typedef size_t usize; + +#ifndef UNUSED_PARAMETER +#define UNUSED_PARAMETER(X) (void)(X) +#endif + +// sqlite3_vtab_in() was added in SQLite version 3.38 (2022-02-22) +// https://www.sqlite.org/changes.html#version_3_38_0 +#if SQLITE_VERSION_NUMBER >= 3038000 +#define COMPILER_SUPPORTS_VTAB_IN 1 +#endif + +#ifndef SQLITE_SUBTYPE +#define SQLITE_SUBTYPE 0x000100000 +#endif + +#ifndef SQLITE_RESULT_SUBTYPE +#define SQLITE_RESULT_SUBTYPE 0x001000000 +#endif + +#ifndef SQLITE_INDEX_CONSTRAINT_LIMIT +#define SQLITE_INDEX_CONSTRAINT_LIMIT 73 +#endif + +#ifndef SQLITE_INDEX_CONSTRAINT_OFFSET +#define SQLITE_INDEX_CONSTRAINT_OFFSET 74 +#endif + +#define countof(x) (sizeof(x) / sizeof((x)[0])) +#define min(a, b) (((a) <= (b)) ? (a) : (b)) + +enum VectorElementType { + // clang-format off + SQLITE_VEC_ELEMENT_TYPE_FLOAT32 = 223 + 0, + SQLITE_VEC_ELEMENT_TYPE_BIT = 223 + 1, + SQLITE_VEC_ELEMENT_TYPE_INT8 = 223 + 2, + // clang-format on +}; + +#ifdef SQLITE_VEC_ENABLE_AVX +#include +#define PORTABLE_ALIGN32 __attribute__((aligned(32))) +#define PORTABLE_ALIGN64 __attribute__((aligned(64))) + +static f32 l2_sqr_float_avx(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + f32 PORTABLE_ALIGN32 TmpRes[8]; + size_t qty16 = qty >> 4; + + const f32 *pEnd1 = pVect1 + (qty16 << 4); + + __m256 diff, v1, v2; + __m256 sum = _mm256_set1_ps(0); + + while (pVect1 < pEnd1) { + v1 = _mm256_loadu_ps(pVect1); + pVect1 += 8; + v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + diff = _mm256_sub_ps(v1, v2); + sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); + + v1 = _mm256_loadu_ps(pVect1); + pVect1 += 8; + v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + diff = _mm256_sub_ps(v1, v2); + sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); + } + + _mm256_store_ps(TmpRes, sum); + return sqrt(TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + + TmpRes[5] + TmpRes[6] + TmpRes[7]); +} +#endif + +#ifdef SQLITE_VEC_ENABLE_NEON +#include + +#define PORTABLE_ALIGN32 __attribute__((aligned(32))) + +// thx https://github.com/nmslib/hnswlib/pull/299/files +static f32 l2_sqr_float_neon(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + size_t qty16 = qty >> 4; + + const f32 *pEnd1 = pVect1 + (qty16 << 4); + + float32x4_t diff, v1, v2; + float32x4_t sum0 = vdupq_n_f32(0); + float32x4_t sum1 = vdupq_n_f32(0); + float32x4_t sum2 = vdupq_n_f32(0); + float32x4_t sum3 = vdupq_n_f32(0); + + while (pVect1 < pEnd1) { + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + diff = vsubq_f32(v1, v2); + sum0 = vfmaq_f32(sum0, diff, diff); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + diff = vsubq_f32(v1, v2); + sum1 = vfmaq_f32(sum1, diff, diff); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + diff = vsubq_f32(v1, v2); + sum2 = vfmaq_f32(sum2, diff, diff); + + v1 = vld1q_f32(pVect1); + pVect1 += 4; + v2 = vld1q_f32(pVect2); + pVect2 += 4; + diff = vsubq_f32(v1, v2); + sum3 = vfmaq_f32(sum3, diff, diff); + } + + f32 sum_scalar = + vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); + const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4)); + while (pVect1 < pEnd2) { + f32 diff = *pVect1 - *pVect2; + sum_scalar += diff * diff; + pVect1++; + pVect2++; + } + + return sqrt(sum_scalar); +} + +static f32 l2_sqr_int8_neon(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + i8 *pVect1 = (i8 *)pVect1v; + i8 *pVect2 = (i8 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + + const i8 *pEnd1 = pVect1 + qty; + i32 sum_scalar = 0; + + while (pVect1 < pEnd1 - 7) { + // loading 8 at a time + int8x8_t v1 = vld1_s8(pVect1); + int8x8_t v2 = vld1_s8(pVect2); + pVect1 += 8; + pVect2 += 8; + + // widen to protect against overflow + int16x8_t v1_wide = vmovl_s8(v1); + int16x8_t v2_wide = vmovl_s8(v2); + + int16x8_t diff = vsubq_s16(v1_wide, v2_wide); + int16x8_t squared_diff = vmulq_s16(diff, diff); + int32x4_t sum = vpaddlq_s16(squared_diff); + + sum_scalar += vgetq_lane_s32(sum, 0) + vgetq_lane_s32(sum, 1) + + vgetq_lane_s32(sum, 2) + vgetq_lane_s32(sum, 3); + } + + // handle leftovers + while (pVect1 < pEnd1) { + i16 diff = (i16)*pVect1 - (i16)*pVect2; + sum_scalar += diff * diff; + pVect1++; + pVect2++; + } + + return sqrtf(sum_scalar); +} + +static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + i8 *pVect1 = (i8 *)pVect1v; + i8 *pVect2 = (i8 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + + const int8_t *pEnd1 = pVect1 + qty; + + int32x4_t acc1 = vdupq_n_s32(0); + int32x4_t acc2 = vdupq_n_s32(0); + int32x4_t acc3 = vdupq_n_s32(0); + int32x4_t acc4 = vdupq_n_s32(0); + + while (pVect1 < pEnd1 - 63) { + int8x16_t v1 = vld1q_s8(pVect1); + int8x16_t v2 = vld1q_s8(pVect2); + int8x16_t diff1 = vabdq_s8(v1, v2); + acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff1))); + + v1 = vld1q_s8(pVect1 + 16); + v2 = vld1q_s8(pVect2 + 16); + int8x16_t diff2 = vabdq_s8(v1, v2); + acc2 = vaddq_s32(acc2, vpaddlq_u16(vpaddlq_u8(diff2))); + + v1 = vld1q_s8(pVect1 + 32); + v2 = vld1q_s8(pVect2 + 32); + int8x16_t diff3 = vabdq_s8(v1, v2); + acc3 = vaddq_s32(acc3, vpaddlq_u16(vpaddlq_u8(diff3))); + + v1 = vld1q_s8(pVect1 + 48); + v2 = vld1q_s8(pVect2 + 48); + int8x16_t diff4 = vabdq_s8(v1, v2); + acc4 = vaddq_s32(acc4, vpaddlq_u16(vpaddlq_u8(diff4))); + + pVect1 += 64; + pVect2 += 64; + } + + while (pVect1 < pEnd1 - 15) { + int8x16_t v1 = vld1q_s8(pVect1); + int8x16_t v2 = vld1q_s8(pVect2); + int8x16_t diff = vabdq_s8(v1, v2); + acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff))); + pVect1 += 16; + pVect2 += 16; + } + + int32x4_t acc = vaddq_s32(vaddq_s32(acc1, acc2), vaddq_s32(acc3, acc4)); + + int32_t sum = 0; + while (pVect1 < pEnd1) { + int32_t diff = abs((int32_t)*pVect1 - (int32_t)*pVect2); + sum += diff; + pVect1++; + pVect2++; + } + + return vaddvq_s32(acc) + sum; +} + +static double l1_f32_neon(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + + const f32 *pEnd1 = pVect1 + qty; + float64x2_t acc = vdupq_n_f64(0); + + while (pVect1 < pEnd1 - 3) { + float32x4_t v1 = vld1q_f32(pVect1); + float32x4_t v2 = vld1q_f32(pVect2); + pVect1 += 4; + pVect2 += 4; + + // f32x4 -> f64x2 pad for overflow + float64x2_t low_diff = vabdq_f64(vcvt_f64_f32(vget_low_f32(v1)), + vcvt_f64_f32(vget_low_f32(v2))); + float64x2_t high_diff = + vabdq_f64(vcvt_high_f64_f32(v1), vcvt_high_f64_f32(v2)); + + acc = vaddq_f64(acc, vaddq_f64(low_diff, high_diff)); + } + + double sum = 0; + while (pVect1 < pEnd1) { + sum += fabs((double)*pVect1 - (double)*pVect2); + pVect1++; + pVect2++; + } + + return vaddvq_f64(acc) + sum; +} +#endif + +static f32 l2_sqr_float(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + + f32 res = 0; + for (size_t i = 0; i < qty; i++) { + f32 t = *pVect1 - *pVect2; + pVect1++; + pVect2++; + res += t * t; + } + return sqrt(res); +} + +static f32 l2_sqr_int8(const void *pA, const void *pB, const void *pD) { + i8 *a = (i8 *)pA; + i8 *b = (i8 *)pB; + size_t d = *((size_t *)pD); + + f32 res = 0; + for (size_t i = 0; i < d; i++) { + f32 t = *a - *b; + a++; + b++; + res += t * t; + } + return sqrt(res); +} + +static f32 distance_l2_sqr_float(const void *a, const void *b, const void *d) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)d) > 16) { + return l2_sqr_float_neon(a, b, d); + } +#endif +#ifdef SQLITE_VEC_ENABLE_AVX + if (((*(const size_t *)d) % 16 == 0)) { + return l2_sqr_float_avx(a, b, d); + } +#endif + return l2_sqr_float(a, b, d); +} + +static f32 distance_l2_sqr_int8(const void *a, const void *b, const void *d) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)d) > 7) { + return l2_sqr_int8_neon(a, b, d); + } +#endif + return l2_sqr_int8(a, b, d); +} + +static i32 l1_int8(const void *pA, const void *pB, const void *pD) { + i8 *a = (i8 *)pA; + i8 *b = (i8 *)pB; + size_t d = *((size_t *)pD); + + i32 res = 0; + for (size_t i = 0; i < d; i++) { + res += abs(*a - *b); + a++; + b++; + } + + return res; +} + +static i32 distance_l1_int8(const void *a, const void *b, const void *d) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)d) > 15) { + return l1_int8_neon(a, b, d); + } +#endif + return l1_int8(a, b, d); +} + +static double l1_f32(const void *pA, const void *pB, const void *pD) { + f32 *a = (f32 *)pA; + f32 *b = (f32 *)pB; + size_t d = *((size_t *)pD); + + double res = 0; + for (size_t i = 0; i < d; i++) { + res += fabs((double)*a - (double)*b); + a++; + b++; + } + + return res; +} + +static double distance_l1_f32(const void *a, const void *b, const void *d) { +#ifdef SQLITE_VEC_ENABLE_NEON + if ((*(const size_t *)d) > 3) { + return l1_f32_neon(a, b, d); + } +#endif + return l1_f32(a, b, d); +} + +static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v, + const void *qty_ptr) { + f32 *pVect1 = (f32 *)pVect1v; + f32 *pVect2 = (f32 *)pVect2v; + size_t qty = *((size_t *)qty_ptr); + + f32 dot = 0; + f32 aMag = 0; + f32 bMag = 0; + for (size_t i = 0; i < qty; i++) { + dot += *pVect1 * *pVect2; + aMag += *pVect1 * *pVect1; + bMag += *pVect2 * *pVect2; + pVect1++; + pVect2++; + } + return 1 - (dot / (sqrt(aMag) * sqrt(bMag))); +} +static f32 distance_cosine_int8(const void *pA, const void *pB, + const void *pD) { + i8 *a = (i8 *)pA; + i8 *b = (i8 *)pB; + size_t d = *((size_t *)pD); + + f32 dot = 0; + f32 aMag = 0; + f32 bMag = 0; + for (size_t i = 0; i < d; i++) { + dot += *a * *b; + aMag += *a * *a; + bMag += *b * *b; + a++; + b++; + } + return 1 - (dot / (sqrt(aMag) * sqrt(bMag))); +} + +// https://github.com/facebookresearch/faiss/blob/77e2e79cd0a680adc343b9840dd865da724c579e/faiss/utils/hamming_distance/common.h#L34 +static u8 hamdist_table[256] = { + 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4, + 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, + 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4, + 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, + 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, + 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, + 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, + 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, + 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, + 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, + 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8}; + +static f32 distance_hamming_u8(u8 *a, u8 *b, size_t n) { + int same = 0; + for (unsigned long i = 0; i < n; i++) { + same += hamdist_table[a[i] ^ b[i]]; + } + return (f32)same; +} + +#ifdef _MSC_VER +#if !defined(__clang__) && (defined(_M_ARM) || defined(_M_ARM64)) +// From +// https://github.com/ngtcp2/ngtcp2/blob/b64f1e77b5e0d880b93d31f474147fae4a1d17cc/lib/ngtcp2_ringbuf.c, +// line 34-43 +static unsigned int __builtin_popcountl(unsigned int x) { + unsigned int c = 0; + for (; x; ++c) { + x &= x - 1; + } + return c; +} +#else +#include +#define __builtin_popcountl __popcnt64 +#endif +#endif + +static f32 distance_hamming_u64(u64 *a, u64 *b, size_t n) { + int same = 0; + for (unsigned long i = 0; i < n; i++) { + same += __builtin_popcountl(a[i] ^ b[i]); + } + return (f32)same; +} + +/** + * @brief Calculate the hamming distance between two bitvectors. + * + * @param a - first bitvector, MUST have d dimensions + * @param b - second bitvector, MUST have d dimensions + * @param d - pointer to size_t, MUST be divisible by CHAR_BIT + * @return f32 + */ +static f32 distance_hamming(const void *a, const void *b, const void *d) { + size_t dimensions = *((size_t *)d); + + if ((dimensions % 64) == 0) { + return distance_hamming_u64((u64 *)a, (u64 *)b, dimensions / 8 / CHAR_BIT); + } + return distance_hamming_u8((u8 *)a, (u8 *)b, dimensions / CHAR_BIT); +} + +#ifdef SQLITE_VEC_TEST +f32 _test_distance_l2_sqr_float(const f32 *a, const f32 *b, size_t dims) { + return distance_l2_sqr_float(a, b, &dims); +} +f32 _test_distance_cosine_float(const f32 *a, const f32 *b, size_t dims) { + return distance_cosine_float(a, b, &dims); +} +f32 _test_distance_hamming(const u8 *a, const u8 *b, size_t dims) { + return distance_hamming(a, b, &dims); +} +#endif + +// from SQLite source: +// https://github.com/sqlite/sqlite/blob/a509a90958ddb234d1785ed7801880ccb18b497e/src/json.c#L153 +static const char vecJsonIsSpaceX[] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, +}; + +#define vecJsonIsspace(x) (vecJsonIsSpaceX[(unsigned char)x]) + +typedef void (*vector_cleanup)(void *p); + +void vector_cleanup_noop(void *_) { UNUSED_PARAMETER(_); } + +#define JSON_SUBTYPE 74 + +void vtab_set_error(sqlite3_vtab *pVTab, const char *zFormat, ...) { + va_list args; + sqlite3_free(pVTab->zErrMsg); + va_start(args, zFormat); + pVTab->zErrMsg = sqlite3_vmprintf(zFormat, args); + va_end(args); +} +struct Array { + size_t element_size; + size_t length; + size_t capacity; + void *z; +}; + +/** + * @brief Initial an array with the given element size and capacity. + * + * @param array + * @param element_size + * @param init_capacity + * @return SQLITE_OK on success, error code on failure. Only error is + * SQLITE_NOMEM + */ +int array_init(struct Array *array, size_t element_size, size_t init_capacity) { + int sz = element_size * init_capacity; + void *z = sqlite3_malloc(sz); + if (!z) { + return SQLITE_NOMEM; + } + memset(z, 0, sz); + + array->element_size = element_size; + array->length = 0; + array->capacity = init_capacity; + array->z = z; + return SQLITE_OK; +} + +int array_append(struct Array *array, const void *element) { + if (array->length == array->capacity) { + size_t new_capacity = array->capacity * 2 + 100; + void *z = sqlite3_realloc64(array->z, array->element_size * new_capacity); + if (z) { + array->capacity = new_capacity; + array->z = z; + } else { + return SQLITE_NOMEM; + } + } + memcpy(&((unsigned char *)array->z)[array->length * array->element_size], + element, array->element_size); + array->length++; + return SQLITE_OK; +} + +void array_cleanup(struct Array *array) { + if (!array) + return; + array->element_size = 0; + array->length = 0; + array->capacity = 0; + sqlite3_free(array->z); + array->z = NULL; +} + +char *vector_subtype_name(int subtype) { + switch (subtype) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: + return "float32"; + case SQLITE_VEC_ELEMENT_TYPE_INT8: + return "int8"; + case SQLITE_VEC_ELEMENT_TYPE_BIT: + return "bit"; + } + return ""; +} +char *type_name(int type) { + switch (type) { + case SQLITE_INTEGER: + return "INTEGER"; + case SQLITE_BLOB: + return "BLOB"; + case SQLITE_TEXT: + return "TEXT"; + case SQLITE_FLOAT: + return "FLOAT"; + case SQLITE_NULL: + return "NULL"; + } + return ""; +} + +typedef void (*fvec_cleanup)(void *vector); + +void fvec_cleanup_noop(void *_) { UNUSED_PARAMETER(_); } + +static int fvec_from_value(sqlite3_value *value, f32 **vector, + size_t *dimensions, fvec_cleanup *cleanup, + char **pzErr) { + int value_type = sqlite3_value_type(value); + + if (value_type == SQLITE_BLOB) { + const void *blob = sqlite3_value_blob(value); + int bytes = sqlite3_value_bytes(value); + if (bytes == 0) { + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + if ((bytes % sizeof(f32)) != 0) { + *pzErr = sqlite3_mprintf("invalid float32 vector BLOB length. Must be " + "divisible by %d, found %d", + sizeof(f32), bytes); + return SQLITE_ERROR; + } + f32 *buf = sqlite3_malloc(bytes); + if (!buf) { + *pzErr = sqlite3_mprintf("out of memory"); + return SQLITE_NOMEM; + } + memcpy(buf, blob, bytes); + *vector = buf; + *dimensions = bytes / sizeof(f32); + *cleanup = sqlite3_free; + return SQLITE_OK; + } + + if (value_type == SQLITE_TEXT) { + const char *source = (const char *)sqlite3_value_text(value); + int source_len = sqlite3_value_bytes(value); + if (source_len == 0) { + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + int i = 0; + + struct Array x; + int rc = array_init(&x, sizeof(f32), ceil(source_len / 2.0)); + if (rc != SQLITE_OK) { + return rc; + } + + // advance leading whitespace to first '[' + while (i < source_len) { + if (vecJsonIsspace(source[i])) { + i++; + continue; + } + if (source[i] == '[') { + break; + } + + *pzErr = sqlite3_mprintf( + "JSON array parsing error: Input does not start with '['"); + array_cleanup(&x); + return SQLITE_ERROR; + } + if (source[i] != '[') { + *pzErr = sqlite3_mprintf( + "JSON array parsing error: Input does not start with '['"); + array_cleanup(&x); + return SQLITE_ERROR; + } + int offset = i + 1; + + while (offset < source_len) { + char *ptr = (char *)&source[offset]; + char *endptr; + + errno = 0; + double result = strtod(ptr, &endptr); + if ((errno != 0 && result == 0) // some interval error? + || (errno == ERANGE && + (result == HUGE_VAL || result == -HUGE_VAL)) // too big / smalls + ) { + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("JSON parsing error"); + return SQLITE_ERROR; + } + + if (endptr == ptr) { + if (*ptr != ']') { + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("JSON parsing error"); + return SQLITE_ERROR; + } + goto done; + } + + f32 res = (f32)result; + array_append(&x, (const void *)&res); + + offset += (endptr - ptr); + while (offset < source_len) { + if (vecJsonIsspace(source[offset])) { + offset++; + continue; + } + if (source[offset] == ',') { + offset++; + continue; + } + if (source[offset] == ']') + goto done; + break; + } + } + + done: + + if (x.length > 0) { + *vector = (f32 *)x.z; + *dimensions = x.length; + *cleanup = sqlite3_free; + return SQLITE_OK; + } + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + + *pzErr = sqlite3_mprintf( + "Input must have type BLOB (compact format) or TEXT (JSON), found %s", + type_name(value_type)); + return SQLITE_ERROR; +} + +static int bitvec_from_value(sqlite3_value *value, u8 **vector, + size_t *dimensions, vector_cleanup *cleanup, + char **pzErr) { + int value_type = sqlite3_value_type(value); + if (value_type == SQLITE_BLOB) { + const void *blob = sqlite3_value_blob(value); + int bytes = sqlite3_value_bytes(value); + if (bytes == 0) { + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + *vector = (u8 *)blob; + *dimensions = bytes * CHAR_BIT; + *cleanup = vector_cleanup_noop; + return SQLITE_OK; + } + *pzErr = sqlite3_mprintf("Unknown type for bitvector."); + return SQLITE_ERROR; +} + +static int int8_vec_from_value(sqlite3_value *value, i8 **vector, + size_t *dimensions, vector_cleanup *cleanup, + char **pzErr) { + int value_type = sqlite3_value_type(value); + if (value_type == SQLITE_BLOB) { + const void *blob = sqlite3_value_blob(value); + int bytes = sqlite3_value_bytes(value); + if (bytes == 0) { + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + *vector = (i8 *)blob; + *dimensions = bytes; + *cleanup = vector_cleanup_noop; + return SQLITE_OK; + } + + if (value_type == SQLITE_TEXT) { + const char *source = (const char *)sqlite3_value_text(value); + int source_len = sqlite3_value_bytes(value); + int i = 0; + + if (source_len == 0) { + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + + struct Array x; + int rc = array_init(&x, sizeof(i8), ceil(source_len / 2.0)); + if (rc != SQLITE_OK) { + return rc; + } + + // advance leading whitespace to first '[' + while (i < source_len) { + if (vecJsonIsspace(source[i])) { + i++; + continue; + } + if (source[i] == '[') { + break; + } + + *pzErr = sqlite3_mprintf( + "JSON array parsing error: Input does not start with '['"); + array_cleanup(&x); + return SQLITE_ERROR; + } + if (source[i] != '[') { + *pzErr = sqlite3_mprintf( + "JSON array parsing error: Input does not start with '['"); + array_cleanup(&x); + return SQLITE_ERROR; + } + int offset = i + 1; + + while (offset < source_len) { + char *ptr = (char *)&source[offset]; + char *endptr; + + errno = 0; + long result = strtol(ptr, &endptr, 10); + if ((errno != 0 && result == 0) || + (errno == ERANGE && (result == LONG_MAX || result == LONG_MIN))) { + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("JSON parsing error"); + return SQLITE_ERROR; + } + + if (endptr == ptr) { + if (*ptr != ']') { + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("JSON parsing error"); + return SQLITE_ERROR; + } + goto done; + } + + if (result < INT8_MIN || result > INT8_MAX) { + sqlite3_free(x.z); + *pzErr = + sqlite3_mprintf("JSON parsing error: value out of range for int8"); + return SQLITE_ERROR; + } + + i8 res = (i8)result; + array_append(&x, (const void *)&res); + + offset += (endptr - ptr); + while (offset < source_len) { + if (vecJsonIsspace(source[offset])) { + offset++; + continue; + } + if (source[offset] == ',') { + offset++; + continue; + } + if (source[offset] == ']') + goto done; + break; + } + } + + done: + + if (x.length > 0) { + *vector = (i8 *)x.z; + *dimensions = x.length; + *cleanup = (vector_cleanup)sqlite3_free; + return SQLITE_OK; + } + sqlite3_free(x.z); + *pzErr = sqlite3_mprintf("zero-length vectors are not supported."); + return SQLITE_ERROR; + } + + *pzErr = sqlite3_mprintf("Unknown type for int8 vector."); + return SQLITE_ERROR; +} + +/** + * @brief Extract a vector from a sqlite3_value. Can be a float32, int8, or bit + * vector. + * + * @param value: the sqlite3_value to read from. + * @param vector: Output pointer to vector data. + * @param dimensions: Output number of dimensions + * @param dimensions: Output vector element type + * @param cleanup + * @param pzErrorMessage + * @return int SQLITE_OK on success, error code otherwise + */ +int vector_from_value(sqlite3_value *value, void **vector, size_t *dimensions, + enum VectorElementType *element_type, + vector_cleanup *cleanup, char **pzErrorMessage) { + int subtype = sqlite3_value_subtype(value); + if (!subtype || (subtype == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) || + (subtype == JSON_SUBTYPE)) { + int rc = fvec_from_value(value, (f32 **)vector, dimensions, + (fvec_cleanup *)cleanup, pzErrorMessage); + if (rc == SQLITE_OK) { + *element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32; + } + return rc; + } + + if (subtype == SQLITE_VEC_ELEMENT_TYPE_BIT) { + int rc = bitvec_from_value(value, (u8 **)vector, dimensions, cleanup, + pzErrorMessage); + if (rc == SQLITE_OK) { + *element_type = SQLITE_VEC_ELEMENT_TYPE_BIT; + } + return rc; + } + if (subtype == SQLITE_VEC_ELEMENT_TYPE_INT8) { + int rc = int8_vec_from_value(value, (i8 **)vector, dimensions, cleanup, + pzErrorMessage); + if (rc == SQLITE_OK) { + *element_type = SQLITE_VEC_ELEMENT_TYPE_INT8; + } + return rc; + } + *pzErrorMessage = sqlite3_mprintf("Unknown subtype: %d", subtype); + return SQLITE_ERROR; +} + +int ensure_vector_match(sqlite3_value *aValue, sqlite3_value *bValue, void **a, + void **b, enum VectorElementType *element_type, + size_t *dimensions, vector_cleanup *outACleanup, + vector_cleanup *outBCleanup, char **outError) { + int rc; + enum VectorElementType aType, bType; + size_t aDims, bDims; + char *error = NULL; + vector_cleanup aCleanup, bCleanup; + + rc = vector_from_value(aValue, a, &aDims, &aType, &aCleanup, &error); + if (rc != SQLITE_OK) { + *outError = sqlite3_mprintf("Error reading 1st vector: %s", error); + sqlite3_free(error); + return SQLITE_ERROR; + } + + rc = vector_from_value(bValue, b, &bDims, &bType, &bCleanup, &error); + if (rc != SQLITE_OK) { + *outError = sqlite3_mprintf("Error reading 2nd vector: %s", error); + sqlite3_free(error); + aCleanup(*a); + return SQLITE_ERROR; + } + + if (aType != bType) { + *outError = + sqlite3_mprintf("Vector type mistmatch. First vector has type %s, " + "while the second has type %s.", + vector_subtype_name(aType), vector_subtype_name(bType)); + aCleanup(*a); + bCleanup(*b); + return SQLITE_ERROR; + } + if (aDims != bDims) { + *outError = sqlite3_mprintf( + "Vector dimension mistmatch. First vector has %ld dimensions, " + "while the second has %ld dimensions.", + aDims, bDims); + aCleanup(*a); + bCleanup(*b); + return SQLITE_ERROR; + } + *element_type = aType; + *dimensions = aDims; + *outACleanup = aCleanup; + *outBCleanup = bCleanup; + return SQLITE_OK; +} + +int _cmp(const void *a, const void *b) { return (*(i64 *)a - *(i64 *)b); } + +struct VecNpyFile { + char *path; + size_t pathLength; +}; +#define SQLITE_VEC_NPY_FILE_NAME "vec0-npy-file" + +#ifndef SQLITE_VEC_OMIT_FS +static void vec_npy_file(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 1); + char *path = (char *)sqlite3_value_text(argv[0]); + size_t pathLength = sqlite3_value_bytes(argv[0]); + struct VecNpyFile *f; + + f = sqlite3_malloc(sizeof(*f)); + if (!f) { + sqlite3_result_error_nomem(context); + return; + } + memset(f, 0, sizeof(*f)); + + f->path = path; + f->pathLength = pathLength; + sqlite3_result_pointer(context, f, SQLITE_VEC_NPY_FILE_NAME, sqlite3_free); +} +#endif + +#pragma region scalar functions +static void vec_f32(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 1); + int rc; + f32 *vector = NULL; + size_t dimensions; + fvec_cleanup cleanup; + char *errmsg; + rc = fvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, errmsg, -1); + sqlite3_free(errmsg); + return; + } + sqlite3_result_blob(context, vector, dimensions * sizeof(f32), + (void (*)(void *))cleanup); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32); +} + +static void vec_bit(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 1); + int rc; + u8 *vector; + size_t dimensions; + vector_cleanup cleanup; + char *errmsg; + rc = bitvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, errmsg, -1); + sqlite3_free(errmsg); + return; + } + sqlite3_result_blob(context, vector, dimensions / CHAR_BIT, SQLITE_TRANSIENT); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT); + cleanup(vector); +} +static void vec_int8(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 1); + int rc; + i8 *vector; + size_t dimensions; + vector_cleanup cleanup; + char *errmsg; + rc = int8_vec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, errmsg, -1); + sqlite3_free(errmsg); + return; + } + sqlite3_result_blob(context, vector, dimensions, SQLITE_TRANSIENT); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8); + cleanup(vector); +} + +static void vec_length(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 1); + int rc; + void *vector; + size_t dimensions; + vector_cleanup cleanup; + char *errmsg; + enum VectorElementType elementType; + rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, &cleanup, + &errmsg); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, errmsg, -1); + sqlite3_free(errmsg); + return; + } + sqlite3_result_int64(context, dimensions); + cleanup(vector); +} + +static void vec_distance_cosine(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a = NULL, *b = NULL; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error( + context, "Cannot calculate cosine distance between two bitvectors.", + -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + f32 result = distance_cosine_float(a, b, &dimensions); + sqlite3_result_double(context, result); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + f32 result = distance_cosine_int8(a, b, &dimensions); + sqlite3_result_double(context, result); + goto finish; + } + } + +finish: + aCleanup(a); + bCleanup(b); + return; +} + +static void vec_distance_l2(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a = NULL, *b = NULL; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error( + context, "Cannot calculate L2 distance between two bitvectors.", -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + f32 result = distance_l2_sqr_float(a, b, &dimensions); + sqlite3_result_double(context, result); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + f32 result = distance_l2_sqr_int8(a, b, &dimensions); + sqlite3_result_double(context, result); + goto finish; + } + } + +finish: + aCleanup(a); + bCleanup(b); + return; +} + +static void vec_distance_l1(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a, *b; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error( + context, "Cannot calculate L1 distance between two bitvectors.", -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + double result = distance_l1_f32(a, b, &dimensions); + sqlite3_result_double(context, result); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + i64 result = distance_l1_int8(a, b, &dimensions); + sqlite3_result_int(context, result); + goto finish; + } + } + +finish: + aCleanup(a); + bCleanup(b); + return; +} + +static void vec_distance_hamming(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a = NULL, *b = NULL; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_double(context, distance_hamming(a, b, &dimensions)); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + sqlite3_result_error( + context, + "Cannot calculate hamming distance between two float32 vectors.", -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + sqlite3_result_error( + context, "Cannot calculate hamming distance between two int8 vectors.", + -1); + goto finish; + } + } + +finish: + aCleanup(a); + bCleanup(b); + return; +} + +char *vec_type_name(enum VectorElementType elementType) { + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: + return "float32"; + case SQLITE_VEC_ELEMENT_TYPE_INT8: + return "int8"; + case SQLITE_VEC_ELEMENT_TYPE_BIT: + return "bit"; + } + return ""; +} + +static void vec_type(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 1); + void *vector; + size_t dimensions; + vector_cleanup cleanup; + char *pzError; + enum VectorElementType elementType; + int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, + &cleanup, &pzError); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, pzError, -1); + sqlite3_free(pzError); + return; + } + sqlite3_result_text(context, vec_type_name(elementType), -1, SQLITE_STATIC); + cleanup(vector); +} +static void vec_quantize_binary(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 1); + void *vector; + size_t dimensions; + vector_cleanup vectorCleanup; + char *pzError; + enum VectorElementType elementType; + int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, + &vectorCleanup, &pzError); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, pzError, -1); + sqlite3_free(pzError); + return; + } + + if (dimensions <= 0) { + sqlite3_result_error(context, "Zero length vectors are not supported.", -1); + goto cleanup; + return; + } + if ((dimensions % CHAR_BIT) != 0) { + sqlite3_result_error( + context, + "Binary quantization requires vectors with a length divisible by 8", + -1); + goto cleanup; + return; + } + + int sz = dimensions / CHAR_BIT; + u8 *out = sqlite3_malloc(sz); + if (!out) { + sqlite3_result_error_code(context, SQLITE_NOMEM); + goto cleanup; + return; + } + memset(out, 0, sz); + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + + for (size_t i = 0; i < dimensions; i++) { + int res = ((f32 *)vector)[i] > 0.0; + out[i / 8] |= (res << (i % 8)); + } + break; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + for (size_t i = 0; i < dimensions; i++) { + int res = ((i8 *)vector)[i] > 0; + out[i / 8] |= (res << (i % 8)); + } + break; + } + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error(context, + "Can only binary quantize float or int8 vectors", -1); + sqlite3_free(out); + return; + } + } + sqlite3_result_blob(context, out, sz, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT); + +cleanup: + vectorCleanup(vector); +} + +static void vec_quantize_int8(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 2); + f32 *srcVector; + size_t dimensions; + fvec_cleanup srcCleanup; + char *err; + i8 *out = NULL; + int rc = fvec_from_value(argv[0], &srcVector, &dimensions, &srcCleanup, &err); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, err, -1); + sqlite3_free(err); + return; + } + + int sz = dimensions * sizeof(i8); + out = sqlite3_malloc(sz); + if (!out) { + sqlite3_result_error_nomem(context); + goto cleanup; + } + memset(out, 0, sz); + + if ((sqlite3_value_type(argv[1]) != SQLITE_TEXT) || + (sqlite3_value_bytes(argv[1]) != strlen("unit")) || + (sqlite3_stricmp((const char *)sqlite3_value_text(argv[1]), "unit") != + 0)) { + sqlite3_result_error( + context, "2nd argument to vec_quantize_int8() must be 'unit'.", -1); + sqlite3_free(out); + goto cleanup; + } + f32 step = (1.0 - (-1.0)) / 255; + for (size_t i = 0; i < dimensions; i++) { + double val = ((srcVector[i] - (-1.0)) / step) - 128; + if (!(val <= 127.0)) val = 127.0; /* also clamps NaN */ + if (!(val >= -128.0)) val = -128.0; + out[i] = (i8)val; + } + + sqlite3_result_blob(context, out, dimensions * sizeof(i8), sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8); + +cleanup: + srcCleanup(srcVector); +} + +static void vec_add(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a = NULL, *b = NULL; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error(context, "Cannot add two bitvectors together.", -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + size_t outSize = dimensions * sizeof(f32); + f32 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + goto finish; + } + memset(out, 0, outSize); + for (size_t i = 0; i < dimensions; i++) { + out[i] = ((f32 *)a)[i] + ((f32 *)b)[i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + size_t outSize = dimensions * sizeof(i8); + i8 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + goto finish; + } + memset(out, 0, outSize); + for (size_t i = 0; i < dimensions; i++) { + out[i] = ((i8 *)a)[i] + ((i8 *)b)[i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8); + goto finish; + } + } +finish: + aCleanup(a); + bCleanup(b); + return; +} +static void vec_sub(sqlite3_context *context, int argc, sqlite3_value **argv) { + assert(argc == 2); + int rc; + void *a = NULL, *b = NULL; + size_t dimensions; + vector_cleanup aCleanup, bCleanup; + char *error; + enum VectorElementType elementType; + rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions, + &aCleanup, &bCleanup, &error); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, error, -1); + sqlite3_free(error); + return; + } + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + sqlite3_result_error(context, "Cannot subtract two bitvectors together.", + -1); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + size_t outSize = dimensions * sizeof(f32); + f32 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + goto finish; + } + memset(out, 0, outSize); + for (size_t i = 0; i < dimensions; i++) { + out[i] = ((f32 *)a)[i] - ((f32 *)b)[i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32); + goto finish; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + size_t outSize = dimensions * sizeof(i8); + i8 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + goto finish; + } + memset(out, 0, outSize); + for (size_t i = 0; i < dimensions; i++) { + out[i] = ((i8 *)a)[i] - ((i8 *)b)[i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8); + goto finish; + } + } +finish: + aCleanup(a); + bCleanup(b); + return; +} +static void vec_slice(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 3); + + void *vector; + size_t dimensions; + vector_cleanup cleanup; + char *err; + enum VectorElementType elementType; + + int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, + &cleanup, &err); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, err, -1); + sqlite3_free(err); + return; + } + + int start = sqlite3_value_int(argv[1]); + int end = sqlite3_value_int(argv[2]); + + if (start < 0) { + sqlite3_result_error(context, + "slice 'start' index must be a postive number.", -1); + goto done; + } + if (end < 0) { + sqlite3_result_error(context, "slice 'end' index must be a postive number.", + -1); + goto done; + } + if (((size_t)start) > dimensions) { + sqlite3_result_error( + context, "slice 'start' index is greater than the number of dimensions", + -1); + goto done; + } + if (((size_t)end) > dimensions) { + sqlite3_result_error( + context, "slice 'end' index is greater than the number of dimensions", + -1); + goto done; + } + if (start > end) { + sqlite3_result_error(context, + "slice 'start' index is greater than 'end' index", -1); + goto done; + } + if (start == end) { + sqlite3_result_error(context, + "slice 'start' index is equal to the 'end' index, " + "vectors must have non-zero length", + -1); + goto done; + } + size_t n = end - start; + + switch (elementType) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + int outSize = n * sizeof(f32); + f32 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + goto done; + } + memset(out, 0, outSize); + for (size_t i = 0; i < n; i++) { + out[i] = ((f32 *)vector)[start + i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32); + goto done; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + int outSize = n * sizeof(i8); + i8 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + return; + } + memset(out, 0, outSize); + for (size_t i = 0; i < n; i++) { + out[i] = ((i8 *)vector)[start + i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8); + goto done; + } + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + if ((start % CHAR_BIT) != 0) { + sqlite3_result_error(context, "start index must be divisible by 8.", -1); + goto done; + } + if ((end % CHAR_BIT) != 0) { + sqlite3_result_error(context, "end index must be divisible by 8.", -1); + goto done; + } + int outSize = n / CHAR_BIT; + u8 *out = sqlite3_malloc(outSize); + if (!out) { + sqlite3_result_error_nomem(context); + return; + } + memset(out, 0, outSize); + for (size_t i = 0; i < n / CHAR_BIT; i++) { + out[i] = ((u8 *)vector)[(start / CHAR_BIT) + i]; + } + sqlite3_result_blob(context, out, outSize, sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT); + goto done; + } + } +done: + cleanup(vector); +} + +static void vec_to_json(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 1); + void *vector; + size_t dimensions; + vector_cleanup cleanup; + char *err; + enum VectorElementType elementType; + + int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, + &cleanup, &err); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, err, -1); + sqlite3_free(err); + return; + } + + sqlite3_str *str = sqlite3_str_new(sqlite3_context_db_handle(context)); + sqlite3_str_appendall(str, "["); + for (size_t i = 0; i < dimensions; i++) { + if (i != 0) { + sqlite3_str_appendall(str, ","); + } + if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) { + f32 value = ((f32 *)vector)[i]; + if (isnan(value)) { + sqlite3_str_appendall(str, "null"); + } else { + sqlite3_str_appendf(str, "%f", value); + } + + } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_INT8) { + sqlite3_str_appendf(str, "%d", ((i8 *)vector)[i]); + } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) { + u8 b = (((u8 *)vector)[i / 8] >> (i % CHAR_BIT)) & 1; + sqlite3_str_appendf(str, "%d", b); + } + } + sqlite3_str_appendall(str, "]"); + int len = sqlite3_str_length(str); + char *s = sqlite3_str_finish(str); + if (s) { + sqlite3_result_text(context, s, len, sqlite3_free); + sqlite3_result_subtype(context, JSON_SUBTYPE); + } else { + sqlite3_result_error_nomem(context); + } + cleanup(vector); +} + +static void vec_normalize(sqlite3_context *context, int argc, + sqlite3_value **argv) { + assert(argc == 1); + void *vector; + size_t dimensions; + vector_cleanup cleanup; + char *err; + enum VectorElementType elementType; + + int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, + &cleanup, &err); + if (rc != SQLITE_OK) { + sqlite3_result_error(context, err, -1); + sqlite3_free(err); + return; + } + + if (elementType != SQLITE_VEC_ELEMENT_TYPE_FLOAT32) { + sqlite3_result_error( + context, "only float32 vectors are supported when normalizing", -1); + cleanup(vector); + return; + } + + int outSize = dimensions * sizeof(f32); + f32 *out = sqlite3_malloc(outSize); + if (!out) { + cleanup(vector); + sqlite3_result_error_code(context, SQLITE_NOMEM); + return; + } + memset(out, 0, outSize); + + f32 *v = (f32 *)vector; + + f32 norm = 0; + for (size_t i = 0; i < dimensions; i++) { + norm += v[i] * v[i]; + } + norm = sqrt(norm); + for (size_t i = 0; i < dimensions; i++) { + out[i] = v[i] / norm; + } + + sqlite3_result_blob(context, out, dimensions * sizeof(f32), sqlite3_free); + sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32); + cleanup(vector); +} + +static void _static_text_func(sqlite3_context *context, int argc, + sqlite3_value **argv) { + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + sqlite3_result_text(context, sqlite3_user_data(context), -1, SQLITE_STATIC); +} + +#pragma endregion + +enum Vec0TokenType { + TOKEN_TYPE_IDENTIFIER, + TOKEN_TYPE_DIGIT, + TOKEN_TYPE_LBRACKET, + TOKEN_TYPE_RBRACKET, + TOKEN_TYPE_PLUS, + TOKEN_TYPE_EQ, + TOKEN_TYPE_LPAREN, + TOKEN_TYPE_RPAREN, + TOKEN_TYPE_COMMA, +}; +struct Vec0Token { + enum Vec0TokenType token_type; + char *start; + char *end; +}; + +int is_alpha(char x) { + return (x >= 'a' && x <= 'z') || (x >= 'A' && x <= 'Z'); +} +int is_digit(char x) { return (x >= '0' && x <= '9'); } +int is_whitespace(char x) { + return x == ' ' || x == '\t' || x == '\n' || x == '\r'; +} + +#define VEC0_TOKEN_RESULT_EOF 1 +#define VEC0_TOKEN_RESULT_SOME 2 +#define VEC0_TOKEN_RESULT_ERROR 3 + +int vec0_token_next(char *start, char *end, struct Vec0Token *out) { + char *ptr = start; + while (ptr < end) { + char curr = *ptr; + if (is_whitespace(curr)) { + ptr++; + continue; + } else if (curr == '+') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_PLUS; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '[') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_LBRACKET; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ']') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_RBRACKET; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '=') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_EQ; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '(') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_LPAREN; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ')') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_RPAREN; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ',') { + ptr++; + out->start = ptr; + out->end = ptr; + out->token_type = TOKEN_TYPE_COMMA; + return VEC0_TOKEN_RESULT_SOME; + } else if (is_alpha(curr)) { + char *start = ptr; + while (ptr < end && (is_alpha(*ptr) || is_digit(*ptr) || *ptr == '_')) { + ptr++; + } + out->start = start; + out->end = ptr; + out->token_type = TOKEN_TYPE_IDENTIFIER; + return VEC0_TOKEN_RESULT_SOME; + } else if (is_digit(curr)) { + char *start = ptr; + while (ptr < end && (is_digit(*ptr))) { + ptr++; + } + out->start = start; + out->end = ptr; + out->token_type = TOKEN_TYPE_DIGIT; + return VEC0_TOKEN_RESULT_SOME; + } else { + return VEC0_TOKEN_RESULT_ERROR; + } + } + return VEC0_TOKEN_RESULT_EOF; +} + +struct Vec0Scanner { + char *start; + char *end; + char *ptr; +}; + +void vec0_scanner_init(struct Vec0Scanner *scanner, const char *source, + int source_length) { + scanner->start = (char *)source; + scanner->end = (char *)source + source_length; + scanner->ptr = (char *)source; +} +int vec0_scanner_next(struct Vec0Scanner *scanner, struct Vec0Token *out) { + int rc = vec0_token_next(scanner->start, scanner->end, out); + if (rc == VEC0_TOKEN_RESULT_SOME) { + scanner->start = out->end; + } + return rc; +} + +int vec0_parse_table_option(const char *source, int source_length, + char **out_key, int *out_key_length, + char **out_value, int *out_value_length) { + int rc; + struct Vec0Scanner scanner; + struct Vec0Token token; + char *key; + char *value; + int keyLength, valueLength; + + vec0_scanner_init(&scanner, source, source_length); + + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + key = token.start; + keyLength = token.end - token.start; + + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) { + return SQLITE_EMPTY; + } + + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + !((token.token_type == TOKEN_TYPE_IDENTIFIER) || + (token.token_type == TOKEN_TYPE_DIGIT))) { + return SQLITE_ERROR; + } + value = token.start; + valueLength = token.end - token.start; + + rc = vec0_scanner_next(&scanner, &token); + if (rc == VEC0_TOKEN_RESULT_EOF) { + *out_key = key; + *out_key_length = keyLength; + *out_value = value; + *out_value_length = valueLength; + return SQLITE_OK; + } + return SQLITE_ERROR; +} +/** + * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if + * it's a PARTITION KEY definition. + * + * @param source: argv[i] source string + * @param source_length: length of the source string + * @param out_column_name: If it is a partition key, the output column name. Same lifetime + * as source, points to specific char * + * @param out_column_name_length: Length of out_column_name in bytes + * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER. + * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is. + */ +int vec0_parse_partition_key_definition(const char *source, int source_length, + char **out_column_name, + int *out_column_name_length, + int *out_column_type) { + struct Vec0Scanner scanner; + struct Vec0Token token; + char *column_name; + int column_name_length; + int column_type; + vec0_scanner_init(&scanner, source, source_length); + + // Check first token is identifier, will be the column name + int rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + + column_name = token.start; + column_name_length = token.end - token.start; + + // Check the next token matches "text" or "integer", as column type + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) { + column_type = SQLITE_TEXT; + } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) == + 0 || + sqlite3_strnicmp(token.start, "integer", + token.end - token.start) == 0) { + column_type = SQLITE_INTEGER; + } else { + return SQLITE_EMPTY; + } + + // Check the next token is identifier and matches "partition" + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "partition", token.end - token.start) != 0) { + return SQLITE_EMPTY; + } + + // Check the next token is identifier and matches "key" + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) { + return SQLITE_EMPTY; + } + + *out_column_name = column_name; + *out_column_name_length = column_name_length; + *out_column_type = column_type; + + return SQLITE_OK; +} + +/** + * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if + * it's an auxiliar column definition, ie `+[name] [type]` like `+contents text` + * + * @param source: argv[i] source string + * @param source_length: length of the source string + * @param out_column_name: If it is a partition key, the output column name. Same lifetime + * as source, points to specific char * + * @param out_column_name_length: Length of out_column_name in bytes + * @param out_column_type: SQLITE_TEXT, SQLITE_INTEGER, SQLITE_FLOAT, or SQLITE_BLOB. + * @return int: SQLITE_EMPTY if not an aux column, SQLITE_OK if it is. + */ +int vec0_parse_auxiliary_column_definition(const char *source, int source_length, + char **out_column_name, + int *out_column_name_length, + int *out_column_type) { + struct Vec0Scanner scanner; + struct Vec0Token token; + char *column_name; + int column_name_length; + int column_type; + vec0_scanner_init(&scanner, source, source_length); + + // Check first token is '+', which denotes aux columns + int rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME || + token.token_type != TOKEN_TYPE_PLUS) { + return SQLITE_EMPTY; + } + + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + + column_name = token.start; + column_name_length = token.end - token.start; + + // Check the next token matches "text" or "integer", as column type + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) { + column_type = SQLITE_TEXT; + } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) == + 0 || + sqlite3_strnicmp(token.start, "integer", + token.end - token.start) == 0) { + column_type = SQLITE_INTEGER; + } else if (sqlite3_strnicmp(token.start, "float", token.end - token.start) == + 0 || + sqlite3_strnicmp(token.start, "double", + token.end - token.start) == 0) { + column_type = SQLITE_FLOAT; + } else if (sqlite3_strnicmp(token.start, "blob", token.end - token.start) ==0) { + column_type = SQLITE_BLOB; + } else { + return SQLITE_EMPTY; + } + + *out_column_name = column_name; + *out_column_name_length = column_name_length; + *out_column_type = column_type; + + return SQLITE_OK; +} + +typedef enum { + VEC0_METADATA_COLUMN_KIND_BOOLEAN, + VEC0_METADATA_COLUMN_KIND_INTEGER, + VEC0_METADATA_COLUMN_KIND_FLOAT, + VEC0_METADATA_COLUMN_KIND_TEXT, + // future: blob, date, datetime +} vec0_metadata_column_kind; + +/** + * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if + * it's an metadata column definition, ie `[name] [type]` like `is_released boolean` + * + * @param source: argv[i] source string + * @param source_length: length of the source string + * @param out_column_name: If it is a metadata column, the output column name. Same lifetime + * as source, points to specific char * + * @param out_column_name_length: Length of out_column_name in bytes + * @param out_column_type: one of vec0_metadata_column_kind + * @return int: SQLITE_EMPTY if not an metadata column, SQLITE_OK if it is. + */ +int vec0_parse_metadata_column_definition(const char *source, int source_length, + char **out_column_name, + int *out_column_name_length, + vec0_metadata_column_kind *out_column_type) { + struct Vec0Scanner scanner; + struct Vec0Token token; + char *column_name; + int column_name_length; + vec0_metadata_column_kind column_type; + int rc; + vec0_scanner_init(&scanner, source, source_length); + + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME || + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + + column_name = token.start; + column_name_length = token.end - token.start; + + // Check the next token matches a valid metadata type + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME || + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + char * t = token.start; + int n = token.end - token.start; + if (sqlite3_strnicmp(t, "boolean", n) == 0 || sqlite3_strnicmp(t, "bool", n) == 0) { + column_type = VEC0_METADATA_COLUMN_KIND_BOOLEAN; + }else if (sqlite3_strnicmp(t, "int64", n) == 0 || sqlite3_strnicmp(t, "integer64", n) == 0 || sqlite3_strnicmp(t, "integer", n) == 0 || sqlite3_strnicmp(t, "int", n) == 0) { + column_type = VEC0_METADATA_COLUMN_KIND_INTEGER; + }else if (sqlite3_strnicmp(t, "float", n) == 0 || sqlite3_strnicmp(t, "double", n) == 0 || sqlite3_strnicmp(t, "float64", n) == 0 || sqlite3_strnicmp(t, "f64", n) == 0) { + column_type = VEC0_METADATA_COLUMN_KIND_FLOAT; + } else if (sqlite3_strnicmp(t, "text", n) == 0) { + column_type = VEC0_METADATA_COLUMN_KIND_TEXT; + } else { + return SQLITE_EMPTY; + } + + *out_column_name = column_name; + *out_column_name_length = column_name_length; + *out_column_type = column_type; + + return SQLITE_OK; +} + +/** + * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if + * it's a PRIMARY KEY definition. + * + * @param source: argv[i] source string + * @param source_length: length of the source string + * @param out_column_name: If it is a PK, the output column name. Same lifetime + * as source, points to specific char * + * @param out_column_name_length: Length of out_column_name in bytes + * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER. + * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is. + */ +int vec0_parse_primary_key_definition(const char *source, int source_length, + char **out_column_name, + int *out_column_name_length, + int *out_column_type) { + struct Vec0Scanner scanner; + struct Vec0Token token; + char *column_name; + int column_name_length; + int column_type; + vec0_scanner_init(&scanner, source, source_length); + + // Check first token is identifier, will be the column name + int rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + + column_name = token.start; + column_name_length = token.end - token.start; + + // Check the next token matches "text" or "integer", as column type + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) { + column_type = SQLITE_TEXT; + } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) == + 0 || + sqlite3_strnicmp(token.start, "integer", + token.end - token.start) == 0) { + column_type = SQLITE_INTEGER; + } else { + return SQLITE_EMPTY; + } + + // Check the next token is identifier and matches "primary" + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "primary", token.end - token.start) != 0) { + return SQLITE_EMPTY; + } + + // Check the next token is identifier and matches "key" + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) { + return SQLITE_EMPTY; + } + + *out_column_name = column_name; + *out_column_name_length = column_name_length; + *out_column_type = column_type; + + return SQLITE_OK; +} + +enum Vec0DistanceMetrics { + VEC0_DISTANCE_METRIC_L2 = 1, + VEC0_DISTANCE_METRIC_COSINE = 2, + VEC0_DISTANCE_METRIC_L1 = 3, +}; + +struct VectorColumnDefinition { + char *name; + int name_length; + size_t dimensions; + enum VectorElementType element_type; + enum Vec0DistanceMetrics distance_metric; +}; + +struct Vec0PartitionColumnDefinition { + int type; + char * name; + int name_length; +}; + +struct Vec0AuxiliaryColumnDefinition { + int type; + char * name; + int name_length; +}; +struct Vec0MetadataColumnDefinition { + vec0_metadata_column_kind kind; + char * name; + int name_length; +}; + +size_t vector_byte_size(enum VectorElementType element_type, + size_t dimensions) { + switch (element_type) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: + return dimensions * sizeof(f32); + case SQLITE_VEC_ELEMENT_TYPE_INT8: + return dimensions * sizeof(i8); + case SQLITE_VEC_ELEMENT_TYPE_BIT: + return dimensions / CHAR_BIT; + } + return 0; +} + +size_t vector_column_byte_size(struct VectorColumnDefinition column) { + return vector_byte_size(column.element_type, column.dimensions); +} + +/** + * @brief Parse an vec0 vtab argv[i] column definition and see if + * it's a vector column defintion, ex `contents_embedding float[768]`. + * + * @param source vec0 argv[i] item + * @param source_length length of source in bytes + * @param outColumn Output the parse vector column to this struct, if success + * @return int SQLITE_OK on success, SQLITE_EMPTY is it's not a vector column + * definition, SQLITE_ERROR on error. + */ +int vec0_parse_vector_column(const char *source, int source_length, + struct VectorColumnDefinition *outColumn) { + // parses a vector column definition like so: + // "abc float[123]", "abc_123 bit[1234]", eetc. + // https://github.com/asg017/sqlite-vec/issues/46 + int rc; + struct Vec0Scanner scanner; + struct Vec0Token token; + + char *name; + int nameLength; + enum VectorElementType elementType; + enum Vec0DistanceMetrics distanceMetric = VEC0_DISTANCE_METRIC_L2; + int dimensions; + + vec0_scanner_init(&scanner, source, source_length); + + // starts with an identifier + rc = vec0_scanner_next(&scanner, &token); + + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + + name = token.start; + nameLength = token.end - token.start; + + // vector column type comes next: float, int, or bit + rc = vec0_scanner_next(&scanner, &token); + + if (rc != VEC0_TOKEN_RESULT_SOME || + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_EMPTY; + } + if (sqlite3_strnicmp(token.start, "float", 5) == 0 || + sqlite3_strnicmp(token.start, "f32", 3) == 0) { + elementType = SQLITE_VEC_ELEMENT_TYPE_FLOAT32; + } else if (sqlite3_strnicmp(token.start, "int8", 4) == 0 || + sqlite3_strnicmp(token.start, "i8", 2) == 0) { + elementType = SQLITE_VEC_ELEMENT_TYPE_INT8; + } else if (sqlite3_strnicmp(token.start, "bit", 3) == 0) { + elementType = SQLITE_VEC_ELEMENT_TYPE_BIT; + } else { + return SQLITE_EMPTY; + } + + // left '[' bracket + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_LBRACKET) { + return SQLITE_EMPTY; + } + + // digit, for vector dimension length + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_DIGIT) { + return SQLITE_ERROR; + } + dimensions = atoi(token.start); + if (dimensions <= 0) { + return SQLITE_ERROR; + } + + // // right ']' bracket + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_RBRACKET) { + return SQLITE_ERROR; + } + + // any other tokens left should be column-level options , ex `key=value` + // ex `distance_metric=L2 distance_metric=cosine` should error + while (1) { + // should be EOF or identifier (option key) + rc = vec0_scanner_next(&scanner, &token); + if (rc == VEC0_TOKEN_RESULT_EOF) { + break; + } + + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_ERROR; + } + + char *key = token.start; + int keyLength = token.end - token.start; + + if (sqlite3_strnicmp(key, "distance_metric", keyLength) == 0) { + + if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) { + return SQLITE_ERROR; + } + // ensure equal sign after distance_metric + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) { + return SQLITE_ERROR; + } + + // distance_metric value, an identifier (L2, cosine, etc) + rc = vec0_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME && + token.token_type != TOKEN_TYPE_IDENTIFIER) { + return SQLITE_ERROR; + } + + char *value = token.start; + int valueLength = token.end - token.start; + if (sqlite3_strnicmp(value, "l2", valueLength) == 0) { + distanceMetric = VEC0_DISTANCE_METRIC_L2; + } else if (sqlite3_strnicmp(value, "l1", valueLength) == 0) { + distanceMetric = VEC0_DISTANCE_METRIC_L1; + } else if (sqlite3_strnicmp(value, "cosine", valueLength) == 0) { + distanceMetric = VEC0_DISTANCE_METRIC_COSINE; + } else { + return SQLITE_ERROR; + } + } + // unknown key + else { + return SQLITE_ERROR; + } + } + + outColumn->name = sqlite3_mprintf("%.*s", nameLength, name); + if (!outColumn->name) { + return SQLITE_ERROR; + } + outColumn->name_length = nameLength; + outColumn->distance_metric = distanceMetric; + outColumn->element_type = elementType; + outColumn->dimensions = dimensions; + return SQLITE_OK; +} + +#pragma region vec_each table function + +typedef struct vec_each_vtab vec_each_vtab; +struct vec_each_vtab { + sqlite3_vtab base; +}; + +typedef struct vec_each_cursor vec_each_cursor; +struct vec_each_cursor { + sqlite3_vtab_cursor base; + i64 iRowid; + enum VectorElementType vector_type; + void *vector; + size_t dimensions; + vector_cleanup cleanup; +}; + +static int vec_eachConnect(sqlite3 *db, void *pAux, int argc, + const char *const *argv, sqlite3_vtab **ppVtab, + char **pzErr) { + UNUSED_PARAMETER(pAux); + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + UNUSED_PARAMETER(pzErr); + vec_each_vtab *pNew; + int rc; + + rc = sqlite3_declare_vtab(db, "CREATE TABLE x(value, vector hidden)"); +#define VEC_EACH_COLUMN_VALUE 0 +#define VEC_EACH_COLUMN_VECTOR 1 + if (rc == SQLITE_OK) { + pNew = sqlite3_malloc(sizeof(*pNew)); + *ppVtab = (sqlite3_vtab *)pNew; + if (pNew == 0) + return SQLITE_NOMEM; + memset(pNew, 0, sizeof(*pNew)); + } + return rc; +} + +static int vec_eachDisconnect(sqlite3_vtab *pVtab) { + vec_each_vtab *p = (vec_each_vtab *)pVtab; + sqlite3_free(p); + return SQLITE_OK; +} + +static int vec_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) { + UNUSED_PARAMETER(p); + vec_each_cursor *pCur; + pCur = sqlite3_malloc(sizeof(*pCur)); + if (pCur == 0) + return SQLITE_NOMEM; + memset(pCur, 0, sizeof(*pCur)); + *ppCursor = &pCur->base; + return SQLITE_OK; +} + +static int vec_eachClose(sqlite3_vtab_cursor *cur) { + vec_each_cursor *pCur = (vec_each_cursor *)cur; + if(pCur->vector) { + pCur->cleanup(pCur->vector); + } + sqlite3_free(pCur); + return SQLITE_OK; +} + +static int vec_eachBestIndex(sqlite3_vtab *pVTab, + sqlite3_index_info *pIdxInfo) { + UNUSED_PARAMETER(pVTab); + int hasVector = 0; + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i]; + // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn, + // pCons->op, pCons->usable); + switch (pCons->iColumn) { + case VEC_EACH_COLUMN_VECTOR: { + if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) { + hasVector = 1; + pIdxInfo->aConstraintUsage[i].argvIndex = 1; + pIdxInfo->aConstraintUsage[i].omit = 1; + } + break; + } + } + } + if (!hasVector) { + return SQLITE_CONSTRAINT; + } + + pIdxInfo->estimatedCost = (double)100000; + pIdxInfo->estimatedRows = 100000; + + return SQLITE_OK; +} + +static int vec_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, + const char *idxStr, int argc, sqlite3_value **argv) { + UNUSED_PARAMETER(idxNum); + UNUSED_PARAMETER(idxStr); + assert(argc == 1); + vec_each_cursor *pCur = (vec_each_cursor *)pVtabCursor; + + if (pCur->vector) { + pCur->cleanup(pCur->vector); + pCur->vector = NULL; + } + + char *pzErrMsg; + int rc = vector_from_value(argv[0], &pCur->vector, &pCur->dimensions, + &pCur->vector_type, &pCur->cleanup, &pzErrMsg); + if (rc != SQLITE_OK) { + sqlite3_free(pzErrMsg); + return SQLITE_ERROR; + } + pCur->iRowid = 0; + return SQLITE_OK; +} + +static int vec_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { + vec_each_cursor *pCur = (vec_each_cursor *)cur; + *pRowid = pCur->iRowid; + return SQLITE_OK; +} + +static int vec_eachEof(sqlite3_vtab_cursor *cur) { + vec_each_cursor *pCur = (vec_each_cursor *)cur; + return pCur->iRowid >= (i64)pCur->dimensions; +} + +static int vec_eachNext(sqlite3_vtab_cursor *cur) { + vec_each_cursor *pCur = (vec_each_cursor *)cur; + pCur->iRowid++; + return SQLITE_OK; +} + +static int vec_eachColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context, + int i) { + vec_each_cursor *pCur = (vec_each_cursor *)cur; + switch (i) { + case VEC_EACH_COLUMN_VALUE: + switch (pCur->vector_type) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + sqlite3_result_double(context, ((f32 *)pCur->vector)[pCur->iRowid]); + break; + } + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + u8 x = ((u8 *)pCur->vector)[pCur->iRowid / CHAR_BIT]; + sqlite3_result_int(context, + (x & (0b10000000 >> ((pCur->iRowid % CHAR_BIT)))) > 0); + break; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + sqlite3_result_int(context, ((i8 *)pCur->vector)[pCur->iRowid]); + break; + } + } + + break; + } + return SQLITE_OK; +} + +static sqlite3_module vec_eachModule = { + /* iVersion */ 0, + /* xCreate */ 0, + /* xConnect */ vec_eachConnect, + /* xBestIndex */ vec_eachBestIndex, + /* xDisconnect */ vec_eachDisconnect, + /* xDestroy */ 0, + /* xOpen */ vec_eachOpen, + /* xClose */ vec_eachClose, + /* xFilter */ vec_eachFilter, + /* xNext */ vec_eachNext, + /* xEof */ vec_eachEof, + /* xColumn */ vec_eachColumn, + /* xRowid */ vec_eachRowid, + /* xUpdate */ 0, + /* xBegin */ 0, + /* xSync */ 0, + /* xCommit */ 0, + /* xRollback */ 0, + /* xFindMethod */ 0, + /* xRename */ 0, + /* xSavepoint */ 0, + /* xRelease */ 0, + /* xRollbackTo */ 0, + /* xShadowName */ 0, +#if SQLITE_VERSION_NUMBER >= 3044000 + /* xIntegrity */ 0 +#endif +}; + +#pragma endregion + +#pragma region vec_npy_each table function + +enum NpyTokenType { + NPY_TOKEN_TYPE_IDENTIFIER, + NPY_TOKEN_TYPE_NUMBER, + NPY_TOKEN_TYPE_LPAREN, + NPY_TOKEN_TYPE_RPAREN, + NPY_TOKEN_TYPE_LBRACE, + NPY_TOKEN_TYPE_RBRACE, + NPY_TOKEN_TYPE_COLON, + NPY_TOKEN_TYPE_COMMA, + NPY_TOKEN_TYPE_STRING, + NPY_TOKEN_TYPE_FALSE, +}; + +struct NpyToken { + enum NpyTokenType token_type; + unsigned char *start; + unsigned char *end; +}; + +int npy_token_next(unsigned char *start, unsigned char *end, + struct NpyToken *out) { + unsigned char *ptr = start; + while (ptr < end) { + unsigned char curr = *ptr; + if (is_whitespace(curr)) { + ptr++; + continue; + } else if (curr == '(') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_LPAREN; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ')') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_RPAREN; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '{') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_LBRACE; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '}') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_RBRACE; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ':') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_COLON; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == ',') { + out->start = ptr++; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_COMMA; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == '\'') { + unsigned char *start = ptr; + ptr++; + while (ptr < end) { + if ((*ptr) == '\'') { + break; + } + ptr++; + } + if (ptr >= end || (*ptr) != '\'') { + return VEC0_TOKEN_RESULT_ERROR; + } + out->start = start; + out->end = ++ptr; + out->token_type = NPY_TOKEN_TYPE_STRING; + return VEC0_TOKEN_RESULT_SOME; + } else if (curr == 'F' && + strncmp((char *)ptr, "False", strlen("False")) == 0) { + out->start = ptr; + out->end = (ptr + (int)strlen("False")); + ptr = out->end; + out->token_type = NPY_TOKEN_TYPE_FALSE; + return VEC0_TOKEN_RESULT_SOME; + } else if (is_digit(curr)) { + unsigned char *start = ptr; + while (ptr < end && (is_digit(*ptr))) { + ptr++; + } + out->start = start; + out->end = ptr; + out->token_type = NPY_TOKEN_TYPE_NUMBER; + return VEC0_TOKEN_RESULT_SOME; + } else { + return VEC0_TOKEN_RESULT_ERROR; + } + } + return VEC0_TOKEN_RESULT_ERROR; +} + +struct NpyScanner { + unsigned char *start; + unsigned char *end; + unsigned char *ptr; +}; + +void npy_scanner_init(struct NpyScanner *scanner, const unsigned char *source, + int source_length) { + scanner->start = (unsigned char *)source; + scanner->end = (unsigned char *)source + source_length; + scanner->ptr = (unsigned char *)source; +} + +int npy_scanner_next(struct NpyScanner *scanner, struct NpyToken *out) { + int rc = npy_token_next(scanner->start, scanner->end, out); + if (rc == VEC0_TOKEN_RESULT_SOME) { + scanner->start = out->end; + } + return rc; +} + +#define NPY_PARSE_ERROR "Error parsing numpy array: " +int parse_npy_header(sqlite3_vtab *pVTab, const unsigned char *header, + size_t headerLength, + enum VectorElementType *out_element_type, + int *fortran_order, size_t *numElements, + size_t *numDimensions) { + + struct NpyScanner scanner; + struct NpyToken token; + int rc; + npy_scanner_init(&scanner, header, headerLength); + + if (npy_scanner_next(&scanner, &token) != VEC0_TOKEN_RESULT_SOME && + token.token_type != NPY_TOKEN_TYPE_LBRACE) { + vtab_set_error(pVTab, + NPY_PARSE_ERROR "numpy header did not start with '{'"); + return SQLITE_ERROR; + } + while (1) { + rc = npy_scanner_next(&scanner, &token); + if (rc != VEC0_TOKEN_RESULT_SOME) { + vtab_set_error(pVTab, NPY_PARSE_ERROR "expected key in numpy header"); + return SQLITE_ERROR; + } + + if (token.token_type == NPY_TOKEN_TYPE_RBRACE) { + break; + } + if (token.token_type != NPY_TOKEN_TYPE_STRING) { + vtab_set_error(pVTab, NPY_PARSE_ERROR + "expected a string as key in numpy header"); + return SQLITE_ERROR; + } + unsigned char *key = token.start; + + rc = npy_scanner_next(&scanner, &token); + if ((rc != VEC0_TOKEN_RESULT_SOME) || + (token.token_type != NPY_TOKEN_TYPE_COLON)) { + vtab_set_error(pVTab, NPY_PARSE_ERROR + "expected a ':' after key in numpy header"); + return SQLITE_ERROR; + } + + if (strncmp((char *)key, "'descr'", strlen("'descr'")) == 0) { + rc = npy_scanner_next(&scanner, &token); + if ((rc != VEC0_TOKEN_RESULT_SOME) || + (token.token_type != NPY_TOKEN_TYPE_STRING)) { + vtab_set_error(pVTab, NPY_PARSE_ERROR + "expected a string value after 'descr' key"); + return SQLITE_ERROR; + } + if (strncmp((char *)token.start, "'maxChunks = 1024; + pCur->chunksBufferSize = + (vector_byte_size(element_type, numDimensions)) * pCur->maxChunks; + pCur->chunksBuffer = sqlite3_malloc(pCur->chunksBufferSize); + if (pCur->chunksBufferSize && !pCur->chunksBuffer) { + return SQLITE_NOMEM; + } + + pCur->currentChunkSize = + fread(pCur->chunksBuffer, vector_byte_size(element_type, numDimensions), + pCur->maxChunks, file); + + pCur->currentChunkIndex = 0; + pCur->elementType = element_type; + pCur->nElements = numElements; + pCur->nDimensions = numDimensions; + pCur->input_type = VEC_NPY_EACH_INPUT_FILE; + + pCur->eof = pCur->currentChunkSize == 0; + pCur->file = file; + return SQLITE_OK; +} +#endif + +int parse_npy_buffer(sqlite3_vtab *pVTab, const unsigned char *buffer, + int bufferLength, void **data, size_t *numElements, + size_t *numDimensions, + enum VectorElementType *element_type) { + + if (bufferLength < 10) { + // IMP: V03312_20150 + vtab_set_error(pVTab, "numpy array too short"); + return SQLITE_ERROR; + } + if (memcmp(NPY_MAGIC, buffer, sizeof(NPY_MAGIC)) != 0) { + // V11954_28792 + vtab_set_error(pVTab, "numpy array does not contain the 'magic' header"); + return SQLITE_ERROR; + } + + u8 major = buffer[6]; + u8 minor = buffer[7]; + uint16_t headerLength = 0; + memcpy(&headerLength, &buffer[8], sizeof(uint16_t)); + + i32 totalHeaderLength = sizeof(NPY_MAGIC) + sizeof(major) + sizeof(minor) + + sizeof(headerLength) + headerLength; + i32 dataSize = bufferLength - totalHeaderLength; + + if (dataSize < 0) { + vtab_set_error(pVTab, "numpy array header length is invalid"); + return SQLITE_ERROR; + } + + const unsigned char *header = &buffer[10]; + int fortran_order; + + int rc = parse_npy_header(pVTab, header, headerLength, element_type, + &fortran_order, numElements, numDimensions); + if (rc != SQLITE_OK) { + return rc; + } + + i32 expectedDataSize = + (*numElements * vector_byte_size(*element_type, *numDimensions)); + if (expectedDataSize != dataSize) { + vtab_set_error(pVTab, + "numpy array error: Expected a data size of %d, found %d", + expectedDataSize, dataSize); + return SQLITE_ERROR; + } + + *data = (void *)&buffer[totalHeaderLength]; + return SQLITE_OK; +} + +static int vec_npy_eachConnect(sqlite3 *db, void *pAux, int argc, + const char *const *argv, sqlite3_vtab **ppVtab, + char **pzErr) { + UNUSED_PARAMETER(pAux); + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + UNUSED_PARAMETER(pzErr); + vec_npy_each_vtab *pNew; + int rc; + + rc = sqlite3_declare_vtab(db, "CREATE TABLE x(vector, input hidden)"); +#define VEC_NPY_EACH_COLUMN_VECTOR 0 +#define VEC_NPY_EACH_COLUMN_INPUT 1 + if (rc == SQLITE_OK) { + pNew = sqlite3_malloc(sizeof(*pNew)); + *ppVtab = (sqlite3_vtab *)pNew; + if (pNew == 0) + return SQLITE_NOMEM; + memset(pNew, 0, sizeof(*pNew)); + } + return rc; +} + +static int vec_npy_eachDisconnect(sqlite3_vtab *pVtab) { + vec_npy_each_vtab *p = (vec_npy_each_vtab *)pVtab; + sqlite3_free(p); + return SQLITE_OK; +} + +static int vec_npy_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) { + UNUSED_PARAMETER(p); + vec_npy_each_cursor *pCur; + pCur = sqlite3_malloc(sizeof(*pCur)); + if (pCur == 0) + return SQLITE_NOMEM; + memset(pCur, 0, sizeof(*pCur)); + *ppCursor = &pCur->base; + return SQLITE_OK; +} + +static int vec_npy_eachClose(sqlite3_vtab_cursor *cur) { + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur; +#ifndef SQLITE_VEC_OMIT_FS + if (pCur->file) { + fclose(pCur->file); + pCur->file = NULL; + } +#endif + if (pCur->chunksBuffer) { + sqlite3_free(pCur->chunksBuffer); + pCur->chunksBuffer = NULL; + } + if (pCur->vector) { + pCur->vector = NULL; + } + sqlite3_free(pCur); + return SQLITE_OK; +} + +static int vec_npy_eachBestIndex(sqlite3_vtab *pVTab, + sqlite3_index_info *pIdxInfo) { + int hasInput; + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i]; + // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn, + // pCons->op, pCons->usable); + switch (pCons->iColumn) { + case VEC_NPY_EACH_COLUMN_INPUT: { + if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) { + hasInput = 1; + pIdxInfo->aConstraintUsage[i].argvIndex = 1; + pIdxInfo->aConstraintUsage[i].omit = 1; + } + break; + } + } + } + if (!hasInput) { + pVTab->zErrMsg = sqlite3_mprintf("input argument is required"); + return SQLITE_ERROR; + } + + pIdxInfo->estimatedCost = (double)100000; + pIdxInfo->estimatedRows = 100000; + + return SQLITE_OK; +} + +static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, + const char *idxStr, int argc, + sqlite3_value **argv) { + UNUSED_PARAMETER(idxNum); + UNUSED_PARAMETER(idxStr); + assert(argc == 1); + int rc; + + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)pVtabCursor; + +#ifndef SQLITE_VEC_OMIT_FS + if (pCur->file) { + fclose(pCur->file); + pCur->file = NULL; + } +#endif + if (pCur->chunksBuffer) { + sqlite3_free(pCur->chunksBuffer); + pCur->chunksBuffer = NULL; + } + if (pCur->vector) { + pCur->vector = NULL; + } + +#ifndef SQLITE_VEC_OMIT_FS + struct VecNpyFile *f = NULL; + if ((f = sqlite3_value_pointer(argv[0], SQLITE_VEC_NPY_FILE_NAME))) { + FILE *file = fopen(f->path, "r"); + if (!file) { + vtab_set_error(pVtabCursor->pVtab, "Could not open numpy file"); + return SQLITE_ERROR; + } + + rc = parse_npy_file(pVtabCursor->pVtab, file, pCur); + if (rc != SQLITE_OK) { +#ifndef SQLITE_VEC_OMIT_FS + fclose(file); +#endif + return rc; + } + + } else +#endif + { + + const unsigned char *input = sqlite3_value_blob(argv[0]); + int inputLength = sqlite3_value_bytes(argv[0]); + void *data; + size_t numElements; + size_t numDimensions; + enum VectorElementType element_type; + + rc = parse_npy_buffer(pVtabCursor->pVtab, input, inputLength, &data, + &numElements, &numDimensions, &element_type); + if (rc != SQLITE_OK) { + return rc; + } + + pCur->vector = data; + pCur->elementType = element_type; + pCur->nElements = numElements; + pCur->nDimensions = numDimensions; + pCur->input_type = VEC_NPY_EACH_INPUT_BUFFER; + } + + pCur->iRowid = 0; + return SQLITE_OK; +} + +static int vec_npy_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur; + *pRowid = pCur->iRowid; + return SQLITE_OK; +} + +static int vec_npy_eachEof(sqlite3_vtab_cursor *cur) { + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur; + if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) { + return (!pCur->nElements) || (size_t)pCur->iRowid >= pCur->nElements; + } + return pCur->eof; +} + +static int vec_npy_eachNext(sqlite3_vtab_cursor *cur) { + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur; + pCur->iRowid++; + if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) { + return SQLITE_OK; + } + +#ifndef SQLITE_VEC_OMIT_FS + // else: input is a file + pCur->currentChunkIndex++; + if (pCur->currentChunkIndex >= pCur->currentChunkSize) { + pCur->currentChunkSize = + fread(pCur->chunksBuffer, + vector_byte_size(pCur->elementType, pCur->nDimensions), + pCur->maxChunks, pCur->file); + if (!pCur->currentChunkSize) { + pCur->eof = 1; + } + pCur->currentChunkIndex = 0; + } +#endif + return SQLITE_OK; +} + +static int vec_npy_eachColumnBuffer(vec_npy_each_cursor *pCur, + sqlite3_context *context, int i) { + switch (i) { + case VEC_NPY_EACH_COLUMN_VECTOR: { + sqlite3_result_subtype(context, pCur->elementType); + switch (pCur->elementType) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + sqlite3_result_blob( + context, + &((unsigned char *) + pCur->vector)[pCur->iRowid * pCur->nDimensions * sizeof(f32)], + pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT); + + break; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + // https://github.com/asg017/sqlite-vec/issues/42 + sqlite3_result_error(context, + "vec_npy_each only supports float32 vectors", -1); + break; + } + } + + break; + } + } + return SQLITE_OK; +} +static int vec_npy_eachColumnFile(vec_npy_each_cursor *pCur, + sqlite3_context *context, int i) { + switch (i) { + case VEC_NPY_EACH_COLUMN_VECTOR: { + switch (pCur->elementType) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + sqlite3_result_blob( + context, + &((unsigned char *) + pCur->chunksBuffer)[pCur->currentChunkIndex * + pCur->nDimensions * sizeof(f32)], + pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT); + break; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + // https://github.com/asg017/sqlite-vec/issues/42 + sqlite3_result_error(context, + "vec_npy_each only supports float32 vectors", -1); + break; + } + } + break; + } + } + return SQLITE_OK; +} +static int vec_npy_eachColumn(sqlite3_vtab_cursor *cur, + sqlite3_context *context, int i) { + vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur; + switch (pCur->input_type) { + case VEC_NPY_EACH_INPUT_BUFFER: + return vec_npy_eachColumnBuffer(pCur, context, i); + case VEC_NPY_EACH_INPUT_FILE: + return vec_npy_eachColumnFile(pCur, context, i); + } + return SQLITE_ERROR; +} + +static sqlite3_module vec_npy_eachModule = { + /* iVersion */ 0, + /* xCreate */ 0, + /* xConnect */ vec_npy_eachConnect, + /* xBestIndex */ vec_npy_eachBestIndex, + /* xDisconnect */ vec_npy_eachDisconnect, + /* xDestroy */ 0, + /* xOpen */ vec_npy_eachOpen, + /* xClose */ vec_npy_eachClose, + /* xFilter */ vec_npy_eachFilter, + /* xNext */ vec_npy_eachNext, + /* xEof */ vec_npy_eachEof, + /* xColumn */ vec_npy_eachColumn, + /* xRowid */ vec_npy_eachRowid, + /* xUpdate */ 0, + /* xBegin */ 0, + /* xSync */ 0, + /* xCommit */ 0, + /* xRollback */ 0, + /* xFindMethod */ 0, + /* xRename */ 0, + /* xSavepoint */ 0, + /* xRelease */ 0, + /* xRollbackTo */ 0, + /* xShadowName */ 0, +#if SQLITE_VERSION_NUMBER >= 3044000 + /* xIntegrity */ 0, +#endif +}; + +#pragma endregion + +#pragma region vec0 virtual table + +#define VEC0_COLUMN_ID 0 +#define VEC0_COLUMN_USERN_START 1 +#define VEC0_COLUMN_OFFSET_DISTANCE 1 +#define VEC0_COLUMN_OFFSET_K 2 + +#define VEC0_SHADOW_INFO_NAME "\"%w\".\"%w_info\"" + +#define VEC0_SHADOW_CHUNKS_NAME "\"%w\".\"%w_chunks\"" +/// 1) schema, 2) original vtab table name +#define VEC0_SHADOW_CHUNKS_CREATE \ + "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(" \ + "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," \ + "size INTEGER NOT NULL," \ + "validity BLOB NOT NULL," \ + "rowids BLOB NOT NULL" \ + ");" + +#define VEC0_SHADOW_ROWIDS_NAME "\"%w\".\"%w_rowids\"" +/// 1) schema, 2) original vtab table name +#define VEC0_SHADOW_ROWIDS_CREATE_BASIC \ + "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \ + "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \ + "id," \ + "chunk_id INTEGER," \ + "chunk_offset INTEGER" \ + ");" + +// vec0 tables with a text primary keys are still backed by int64 primary keys, +// since a fixed-length rowid is required for vec0 chunks. But we add a new 'id +// text unique' column to emulate a text primary key interface. +#define VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT \ + "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \ + "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \ + "id TEXT UNIQUE NOT NULL," \ + "chunk_id INTEGER," \ + "chunk_offset INTEGER" \ + ");" + +/// 1) schema, 2) original vtab table name +#define VEC0_SHADOW_VECTOR_N_NAME "\"%w\".\"%w_vector_chunks%02d\"" + +/// 1) schema, 2) original vtab table name +// +// IMPORTANT: "rowid" is declared as PRIMARY KEY but WITHOUT the INTEGER type. +// This means it is NOT a true SQLite rowid alias — the user-defined "rowid" +// column and the internal SQLite rowid (_rowid_) are two separate values. +// When inserting, both must be set explicitly to keep them in sync. See the +// _rowid_ bindings in vec0_new_chunk() and the explanation in +// SHADOW_TABLE_ROWID_QUIRK below. +#define VEC0_SHADOW_VECTOR_N_CREATE \ + "CREATE TABLE " VEC0_SHADOW_VECTOR_N_NAME "(" \ + "rowid PRIMARY KEY," \ + "vectors BLOB NOT NULL" \ + ");" + +#define VEC0_SHADOW_AUXILIARY_NAME "\"%w\".\"%w_auxiliary\"" + +#define VEC0_SHADOW_METADATA_N_NAME "\"%w\".\"%w_metadatachunks%02d\"" +#define VEC0_SHADOW_METADATA_TEXT_DATA_NAME "\"%w\".\"%w_metadatatext%02d\"" + +#define VEC_INTERAL_ERROR "Internal sqlite-vec error: " +#define REPORT_URL "https://github.com/asg017/sqlite-vec/issues/new" + +typedef struct vec0_vtab vec0_vtab; + +#define VEC0_MAX_VECTOR_COLUMNS 16 +#define VEC0_MAX_PARTITION_COLUMNS 4 +#define VEC0_MAX_AUXILIARY_COLUMNS 16 +#define VEC0_MAX_METADATA_COLUMNS 16 + +#define SQLITE_VEC_VEC0_MAX_DIMENSIONS 8192 +#define VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH 16 +#define VEC0_METADATA_TEXT_VIEW_DATA_LENGTH 12 + +typedef enum { + // vector column, ie "contents_embedding float[1024]" + SQLITE_VEC0_USER_COLUMN_KIND_VECTOR = 1, + + // partition key column, ie "user_id integer partition key" + SQLITE_VEC0_USER_COLUMN_KIND_PARTITION = 2, + + // + SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY = 3, + + // metadata column that can be filtered, ie "genre text" + SQLITE_VEC0_USER_COLUMN_KIND_METADATA = 4, +} vec0_user_column_kind; + +struct vec0_vtab { + sqlite3_vtab base; + + // the SQLite connection of the host database + sqlite3 *db; + + // True if the primary key of the vec0 table has a column type TEXT. + // Will change the schema of the _rowids table, and insert/query logic. + int pkIsText; + + // number of defined vector columns. + int numVectorColumns; + + // number of defined PARTITION KEY columns. + int numPartitionColumns; + + // number of defined auxiliary columns + int numAuxiliaryColumns; + + // number of defined metadata columns + int numMetadataColumns; + + + // Name of the schema the table exists on. + // Must be freed with sqlite3_free() + char *schemaName; + + // Name of the table the table exists on. + // Must be freed with sqlite3_free() + char *tableName; + + // Name of the _rowids shadow table. + // Must be freed with sqlite3_free() + char *shadowRowidsName; + + // Name of the _chunks shadow table. + // Must be freed with sqlite3_free() + char *shadowChunksName; + + // contains enum vec0_user_column_kind values for up to + // numVectorColumns + numPartitionColumns entries + vec0_user_column_kind user_column_kinds[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS]; + + uint8_t user_column_idxs[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS]; + + + // Name of all the vector chunk shadow tables. + // Ex '_vector_chunks00' + // Only the first numVectorColumns entries will be available. + // The first numVectorColumns entries must be freed with sqlite3_free() + char *shadowVectorChunksNames[VEC0_MAX_VECTOR_COLUMNS]; + + // Name of all metadata chunk shadow tables, ie `_metadatachunks00` + // Only the first numMetadataColumns entries will be available. + // The first numMetadataColumns entries must be freed with sqlite3_free() + char *shadowMetadataChunksNames[VEC0_MAX_METADATA_COLUMNS]; + + struct VectorColumnDefinition vector_columns[VEC0_MAX_VECTOR_COLUMNS]; + struct Vec0PartitionColumnDefinition paritition_columns[VEC0_MAX_PARTITION_COLUMNS]; + struct Vec0AuxiliaryColumnDefinition auxiliary_columns[VEC0_MAX_AUXILIARY_COLUMNS]; + struct Vec0MetadataColumnDefinition metadata_columns[VEC0_MAX_METADATA_COLUMNS]; + + int chunk_size; + + // select latest chunk from _chunks, getting chunk_id + sqlite3_stmt *stmtLatestChunk; + + /** + * Statement to insert a row into the _rowids table, with a rowid. + * Parameters: + * 1: int64, rowid to insert + * Result columns: none + * SQL: "INSERT INTO _rowids(rowid) VALUES (?)" + * + * Must be cleaned up with sqlite3_finalize(). + */ + sqlite3_stmt *stmtRowidsInsertRowid; + + /** + * Statement to insert a row into the _rowids table, with an id. + * The id column isn't a tradition primary key, but instead a unique + * column to handle "text primary key" vec0 tables. The true int64 rowid + * can be retrieved after inserting with sqlite3_last_rowid(). + * + * Parameters: + * 1: text or null, id to insert + * Result columns: none + * + * Must be cleaned up with sqlite3_finalize(). + */ + sqlite3_stmt *stmtRowidsInsertId; + + /** + * Statement to update the "position" columns chunk_id and chunk_offset for + * a given _rowids row. Used when the "next available" chunk position is found + * for a vector. + * + * Parameters: + * 1: int64, chunk_id value + * 2: int64, chunk_offset value + * 3: int64, rowid value + * Result columns: none + * + * Must be cleaned up with sqlite3_finalize(). + */ + sqlite3_stmt *stmtRowidsUpdatePosition; + + /** + * Statement to quickly find the chunk_id + chunk_offset of a given row. + * Parameters: + * 1: rowid of the row/vector to lookup + * Result columns: + * 0: chunk_id (i64) + * 1: chunk_offset (i64) + * SQL: "SELECT id, chunk_id, chunk_offset FROM _rowids WHERE rowid = ?"" + * + * Must be cleaned up with sqlite3_finalize(). + */ + sqlite3_stmt *stmtRowidsGetChunkPosition; +}; + +/** + * @brief Finalize all the sqlite3_stmt members in a vec0_vtab. + * + * @param p vec0_vtab pointer + */ +void vec0_free_resources(vec0_vtab *p) { + sqlite3_finalize(p->stmtLatestChunk); + p->stmtLatestChunk = NULL; + sqlite3_finalize(p->stmtRowidsInsertRowid); + p->stmtRowidsInsertRowid = NULL; + sqlite3_finalize(p->stmtRowidsInsertId); + p->stmtRowidsInsertId = NULL; + sqlite3_finalize(p->stmtRowidsUpdatePosition); + p->stmtRowidsUpdatePosition = NULL; + sqlite3_finalize(p->stmtRowidsGetChunkPosition); + p->stmtRowidsGetChunkPosition = NULL; +} + +/** + * @brief Free all memory and sqlite3_stmt members of a vec0_vtab + * + * @param p vec0_vtab pointer + */ +void vec0_free(vec0_vtab *p) { + vec0_free_resources(p); + + sqlite3_free(p->schemaName); + p->schemaName = NULL; + sqlite3_free(p->tableName); + p->tableName = NULL; + sqlite3_free(p->shadowChunksName); + p->shadowChunksName = NULL; + sqlite3_free(p->shadowRowidsName); + p->shadowRowidsName = NULL; + + for (int i = 0; i < p->numVectorColumns; i++) { + sqlite3_free(p->shadowVectorChunksNames[i]); + p->shadowVectorChunksNames[i] = NULL; + + sqlite3_free(p->vector_columns[i].name); + p->vector_columns[i].name = NULL; + } + + for (int i = 0; i < p->numPartitionColumns; i++) { + sqlite3_free(p->paritition_columns[i].name); + p->paritition_columns[i].name = NULL; + } + + for (int i = 0; i < p->numAuxiliaryColumns; i++) { + sqlite3_free(p->auxiliary_columns[i].name); + p->auxiliary_columns[i].name = NULL; + } + + for (int i = 0; i < p->numMetadataColumns; i++) { + sqlite3_free(p->metadata_columns[i].name); + p->metadata_columns[i].name = NULL; + } +} + +int vec0_num_defined_user_columns(vec0_vtab *p) { + return p->numVectorColumns + p->numPartitionColumns + p->numAuxiliaryColumns + p->numMetadataColumns; +} + +/** + * @brief Returns the index of the distance hidden column for the given vec0 + * table. + * + * @param p vec0 table + * @return int + */ +int vec0_column_distance_idx(vec0_vtab *p) { + return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) + + VEC0_COLUMN_OFFSET_DISTANCE; +} + +/** + * @brief Returns the index of the k hidden column for the given vec0 table. + * + * @param p vec0 table + * @return int k column index + */ +int vec0_column_k_idx(vec0_vtab *p) { + return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) + + VEC0_COLUMN_OFFSET_K; +} + +/** + * Returns 1 if the given column-based index is a valid vector column, + * 0 otherwise. + */ +int vec0_column_idx_is_vector(vec0_vtab *pVtab, int column_idx) { + return column_idx >= VEC0_COLUMN_USERN_START && + column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) && + pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_VECTOR; +} + +/** + * Returns the vector index of the given user column index. + * ONLY call if validated with vec0_column_idx_is_vector before + */ +int vec0_column_idx_to_vector_idx(vec0_vtab *pVtab, int column_idx) { + UNUSED_PARAMETER(pVtab); + return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START]; +} +/** + * Returns 1 if the given column-based index is a "partition key" column, + * 0 otherwise. + */ +int vec0_column_idx_is_partition(vec0_vtab *pVtab, int column_idx) { + return column_idx >= VEC0_COLUMN_USERN_START && + column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) && + pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_PARTITION; +} + +/** + * Returns the partition column index of the given user column index. + * ONLY call if validated with vec0_column_idx_is_vector before + */ +int vec0_column_idx_to_partition_idx(vec0_vtab *pVtab, int column_idx) { + UNUSED_PARAMETER(pVtab); + return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START]; +} + +/** + * Returns 1 if the given column-based index is a auxiliary column, + * 0 otherwise. + */ +int vec0_column_idx_is_auxiliary(vec0_vtab *pVtab, int column_idx) { + return column_idx >= VEC0_COLUMN_USERN_START && + column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) && + pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY; +} + +/** + * Returns the auxiliary column index of the given user column index. + * ONLY call if validated with vec0_column_idx_to_partition_idx before + */ +int vec0_column_idx_to_auxiliary_idx(vec0_vtab *pVtab, int column_idx) { + UNUSED_PARAMETER(pVtab); + return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START]; +} + +/** + * Returns 1 if the given column-based index is a metadata column, + * 0 otherwise. + */ +int vec0_column_idx_is_metadata(vec0_vtab *pVtab, int column_idx) { + return column_idx >= VEC0_COLUMN_USERN_START && + column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) && + pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_METADATA; +} + +/** + * Returns the metadata column index of the given user column index. + * ONLY call if validated with vec0_column_idx_is_metadata before + */ +int vec0_column_idx_to_metadata_idx(vec0_vtab *pVtab, int column_idx) { + UNUSED_PARAMETER(pVtab); + return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START]; +} + +/** + * @brief Retrieve the chunk_id, chunk_offset, and possible "id" value + * of a vec0_vtab row with the provided rowid + * + * @param p vec0_vtab + * @param rowid the rowid of the row to query + * @param id output, optional sqlite3_value to provide the id. + * Useful for text PK rows. Must be freed with sqlite3_value_free() + * @param chunk_id output, the chunk_id the row belongs to + * @param chunk_offset output, the offset within the chunk the row belongs to + * @return SQLITE_ROW on success, error code otherwise. SQLITE_EMPTY if row DNE + */ +int vec0_get_chunk_position(vec0_vtab *p, i64 rowid, sqlite3_value **id, + i64 *chunk_id, i64 *chunk_offset) { + int rc; + + if (!p->stmtRowidsGetChunkPosition) { + const char *zSql = + sqlite3_mprintf("SELECT id, chunk_id, chunk_offset " + "FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsGetChunkPosition, 0); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + vtab_set_error( + &p->base, VEC_INTERAL_ERROR + "could not initialize 'rowids get chunk position' statement"); + goto cleanup; + } + } + + sqlite3_bind_int64(p->stmtRowidsGetChunkPosition, 1, rowid); + rc = sqlite3_step(p->stmtRowidsGetChunkPosition); + // special case: when no results, return SQLITE_EMPTY to convey "that chunk + // position doesnt exist" + if (rc == SQLITE_DONE) { + rc = SQLITE_EMPTY; + goto cleanup; + } + if (rc != SQLITE_ROW) { + goto cleanup; + } + + if (id) { + sqlite3_value *value = + sqlite3_column_value(p->stmtRowidsGetChunkPosition, 0); + *id = sqlite3_value_dup(value); + if (!*id) { + rc = SQLITE_NOMEM; + goto cleanup; + } + } + + if (chunk_id) { + *chunk_id = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 1); + } + if (chunk_offset) { + *chunk_offset = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 2); + } + + rc = SQLITE_OK; + +cleanup: + sqlite3_reset(p->stmtRowidsGetChunkPosition); + sqlite3_clear_bindings(p->stmtRowidsGetChunkPosition); + return rc; +} + +/** + * @brief Return the id value from the _rowids table where _rowids.rowid = + * rowid. + * + * @param pVtab: vec0 table to query + * @param rowid: rowid of the row to query. + * @param out: A dup'ed sqlite3_value of the id column. Might be null. + * Must be cleaned up with sqlite3_value_free(). + * @returns SQLITE_OK on success, error code on failure + */ +int vec0_get_id_value_from_rowid(vec0_vtab *pVtab, i64 rowid, + sqlite3_value **out) { + // PERF: different strategy than get_chunk_position? + return vec0_get_chunk_position((vec0_vtab *)pVtab, rowid, out, NULL, NULL); +} + +int vec0_rowid_from_id(vec0_vtab *p, sqlite3_value *valueId, i64 *rowid) { + sqlite3_stmt *stmt = NULL; + int rc; + char *zSql; + zSql = sqlite3_mprintf("SELECT rowid" + " FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE id = ?", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_bind_value(stmt, 1, valueId); + rc = sqlite3_step(stmt); + if (rc == SQLITE_DONE) { + rc = SQLITE_EMPTY; + goto cleanup; + } + if (rc != SQLITE_ROW) { + goto cleanup; + } + *rowid = sqlite3_column_int64(stmt, 0); + rc = sqlite3_step(stmt); + if (rc != SQLITE_DONE) { + goto cleanup; + } + + rc = SQLITE_OK; + +cleanup: + sqlite3_finalize(stmt); + return rc; +} + +int vec0_result_id(vec0_vtab *p, sqlite3_context *context, i64 rowid) { + if (!p->pkIsText) { + sqlite3_result_int64(context, rowid); + return SQLITE_OK; + } + sqlite3_value *valueId; + int rc = vec0_get_id_value_from_rowid(p, rowid, &valueId); + if (rc != SQLITE_OK) { + return rc; + } + if (!valueId) { + sqlite3_result_error_nomem(context); + } else { + sqlite3_result_value(context, valueId); + sqlite3_value_free(valueId); + } + return SQLITE_OK; +} + +/** + * @brief + * + * @param pVtab: virtual table to query + * @param rowid: row to lookup + * @param vector_column_idx: which vector column to query + * @param outVector: Output pointer to the vector buffer. + * Must be sqlite3_free()'ed. + * @param outVectorSize: Pointer to a int where the size of outVector + * will be stored. + * @return int SQLITE_OK on success. + */ +int vec0_get_vector_data(vec0_vtab *pVtab, i64 rowid, int vector_column_idx, + void **outVector, int *outVectorSize) { + vec0_vtab *p = pVtab; + int rc, brc; + i64 chunk_id; + i64 chunk_offset; + size_t size; + void *buf = NULL; + int blobOffset; + sqlite3_blob *vectorBlob = NULL; + assert((vector_column_idx >= 0) && + (vector_column_idx < pVtab->numVectorColumns)); + + rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset); + if (rc == SQLITE_EMPTY) { + vtab_set_error(&pVtab->base, "Could not find a row with rowid %lld", rowid); + goto cleanup; + } + if (rc != SQLITE_OK) { + goto cleanup; + } + + rc = sqlite3_blob_open(p->db, p->schemaName, + p->shadowVectorChunksNames[vector_column_idx], + "vectors", chunk_id, 0, &vectorBlob); + + if (rc != SQLITE_OK) { + vtab_set_error(&pVtab->base, + "Could not fetch vector data for %lld, opening blob failed", + rowid); + rc = SQLITE_ERROR; + goto cleanup; + } + + size = vector_column_byte_size(pVtab->vector_columns[vector_column_idx]); + blobOffset = chunk_offset * size; + + buf = sqlite3_malloc(size); + if (!buf) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + rc = sqlite3_blob_read(vectorBlob, buf, size, blobOffset); + if (rc != SQLITE_OK) { + sqlite3_free(buf); + buf = NULL; + vtab_set_error( + &pVtab->base, + "Could not fetch vector data for %lld, reading from blob failed", + rowid); + rc = SQLITE_ERROR; + goto cleanup; + } + + *outVector = buf; + if (outVectorSize) { + *outVectorSize = size; + } + rc = SQLITE_OK; + +cleanup: + brc = sqlite3_blob_close(vectorBlob); + if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) { + vtab_set_error( + &p->base, VEC_INTERAL_ERROR + "unknown error, could not close vector blob, please file an issue"); + return brc; + } + + return rc; +} + +/** + * @brief Retrieve the sqlite3_value of the i'th partition value for the given row. + * + * @param pVtab - the vec0_vtab in questions + * @param rowid - rowid of target row + * @param partition_idx - which partition column to retrieve + * @param outValue - output sqlite3_value + * @return int - SQLITE_OK on success, otherwise error code + */ +int vec0_get_partition_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int partition_idx, sqlite3_value ** outValue) { + int rc; + i64 chunk_id; + i64 chunk_offset; + rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset); + if(rc != SQLITE_OK) { + return rc; + } + sqlite3_stmt * stmt = NULL; + char * zSql = sqlite3_mprintf("SELECT partition%02d FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE chunk_id = ?", partition_idx, pVtab->schemaName, pVtab->tableName); + if(!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if(rc != SQLITE_OK) { + return rc; + } + sqlite3_bind_int64(stmt, 1, chunk_id); + rc = sqlite3_step(stmt); + if(rc != SQLITE_ROW) { + rc = SQLITE_ERROR; + goto done; + } + *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0)); + if(!*outValue) { + rc = SQLITE_NOMEM; + goto done; + } + rc = SQLITE_OK; + + done: + sqlite3_finalize(stmt); + return rc; + +} + +/** + * @brief Get the value of an auxiliary column for the given rowid + * + * @param pVtab vec0_vtab + * @param rowid the rowid of the row to lookup + * @param auxiliary_idx aux index of the column we care about + * @param outValue Output sqlite3_value to store + * @return int SQLITE_OK on success, error code otherwise + */ +int vec0_get_auxiliary_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int auxiliary_idx, sqlite3_value ** outValue) { + int rc; + sqlite3_stmt * stmt = NULL; + char * zSql = sqlite3_mprintf("SELECT value%02d FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?", auxiliary_idx, pVtab->schemaName, pVtab->tableName); + if(!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if(rc != SQLITE_OK) { + return rc; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + if(rc != SQLITE_ROW) { + rc = SQLITE_ERROR; + goto done; + } + *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0)); + if(!*outValue) { + rc = SQLITE_NOMEM; + goto done; + } + rc = SQLITE_OK; + + done: + sqlite3_finalize(stmt); + return rc; +} + +/** + * @brief Result the given metadata value for the given row and metadata column index. + * Will traverse the metadatachunksNN table with BLOB I/0 for the given rowid. + * + * @param p + * @param rowid + * @param metadata_idx + * @param context + * @return int + */ +int vec0_result_metadata_value_for_rowid(vec0_vtab *p, i64 rowid, int metadata_idx, sqlite3_context * context) { + int rc; + i64 chunk_id; + i64 chunk_offset; + rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset); + if(rc != SQLITE_OK) { + return rc; + } + sqlite3_blob * blobValue; + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &blobValue); + if(rc != SQLITE_OK) { + return rc; + } + + switch(p->metadata_columns[metadata_idx].kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + u8 block; + rc = sqlite3_blob_read(blobValue, &block, sizeof(block), chunk_offset / CHAR_BIT); + if(rc != SQLITE_OK) { + goto done; + } + int value = block >> ((chunk_offset % CHAR_BIT)) & 1; + sqlite3_result_int(context, value); + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + i64 value; + rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64)); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_result_int64(context, value); + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + double value; + rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(double)); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_result_double(context, value); + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + rc = sqlite3_blob_read(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + if(rc != SQLITE_OK) { + goto done; + } + int length = ((int *)view)[0]; + if(length <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + sqlite3_result_text(context, (const char*) (view + 4), length, SQLITE_TRANSIENT); + } + else { + sqlite3_stmt * stmt; + const char * zSql = sqlite3_mprintf("SELECT data FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx); + if(!zSql) { + rc = SQLITE_ERROR; + goto done; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free((void *) zSql); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + if(rc != SQLITE_ROW) { + sqlite3_finalize(stmt); + rc = SQLITE_ERROR; + goto done; + } + sqlite3_result_value(context, sqlite3_column_value(stmt, 0)); + sqlite3_finalize(stmt); + rc = SQLITE_OK; + } + break; + } + } + done: + // blobValue is read-only, will not fail on close + sqlite3_blob_close(blobValue); + return rc; + +} + +int vec0_get_latest_chunk_rowid(vec0_vtab *p, i64 *chunk_rowid, sqlite3_value ** partitionKeyValues) { + int rc; + const char *zSql; + // lazy initialize stmtLatestChunk when needed. May be cleared during xSync() + if (!p->stmtLatestChunk) { + if(p->numPartitionColumns > 0) { + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE ", + p->schemaName, p->tableName); + + for(int i = 0; i < p->numPartitionColumns; i++) { + if(i != 0) { + sqlite3_str_appendall(s, " AND "); + } + sqlite3_str_appendf(s, " partition%02d = ? ", i); + } + zSql = sqlite3_str_finish(s); + }else { + zSql = sqlite3_mprintf("SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME, + p->schemaName, p->tableName); + } + + if (!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtLatestChunk, 0); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + // IMP: V21406_05476 + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "could not initialize 'latest chunk' statement"); + goto cleanup; + } + } + + for(int i = 0; i < p->numPartitionColumns; i++) { + sqlite3_bind_value(p->stmtLatestChunk, i+1, (partitionKeyValues[i])); + } + + rc = sqlite3_step(p->stmtLatestChunk); + if (rc != SQLITE_ROW) { + // IMP: V31559_15629 + vtab_set_error(&p->base, VEC_INTERAL_ERROR "Could not find latest chunk"); + rc = SQLITE_ERROR; + goto cleanup; + } + if(sqlite3_column_type(p->stmtLatestChunk, 0) == SQLITE_NULL){ + rc = SQLITE_EMPTY; + goto cleanup; + } + *chunk_rowid = sqlite3_column_int64(p->stmtLatestChunk, 0); + rc = sqlite3_step(p->stmtLatestChunk); + if (rc != SQLITE_DONE) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "unknown result code when closing out stmtLatestChunk. " + "Please file an issue: " REPORT_URL, + p->schemaName, p->shadowChunksName); + goto cleanup; + } + rc = SQLITE_OK; + +cleanup: + if (p->stmtLatestChunk) { + sqlite3_reset(p->stmtLatestChunk); + sqlite3_clear_bindings(p->stmtLatestChunk); + } + return rc; +} + +int vec0_rowids_insert_rowid(vec0_vtab *p, i64 rowid) { + int rc = SQLITE_OK; + int entered = 0; + UNUSED_PARAMETER(entered); // temporary + if (!p->stmtRowidsInsertRowid) { + const char *zSql = + sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(rowid)" + "VALUES (?);", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertRowid, 0); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "could not initialize 'insert rowids' statement"); + goto cleanup; + } + } + +#if SQLITE_THREADSAFE + if (sqlite3_mutex_enter) { + sqlite3_mutex_enter(sqlite3_db_mutex(p->db)); + entered = 1; + } +#endif + sqlite3_bind_int64(p->stmtRowidsInsertRowid, 1, rowid); + rc = sqlite3_step(p->stmtRowidsInsertRowid); + + if (rc != SQLITE_DONE) { + if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_PRIMARYKEY) { + // IMP: V17090_01160 + vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key", + p->tableName); + } else { + // IMP: V04679_21517 + vtab_set_error(&p->base, + "Error inserting rowid into rowids shadow table: %s", + sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId))); + } + rc = SQLITE_ERROR; + goto cleanup; + } + + rc = SQLITE_OK; + +cleanup: + if (p->stmtRowidsInsertRowid) { + sqlite3_reset(p->stmtRowidsInsertRowid); + sqlite3_clear_bindings(p->stmtRowidsInsertRowid); + } + +#if SQLITE_THREADSAFE + if (sqlite3_mutex_leave && entered) { + sqlite3_mutex_leave(sqlite3_db_mutex(p->db)); + } +#endif + return rc; +} + +int vec0_rowids_insert_id(vec0_vtab *p, sqlite3_value *idValue, i64 *rowid) { + int rc = SQLITE_OK; + int entered = 0; + UNUSED_PARAMETER(entered); // temporary + if (!p->stmtRowidsInsertId) { + const char *zSql = + sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(id)" + "VALUES (?);", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto complete; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertId, 0); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "could not initialize 'insert rowids id' statement"); + goto complete; + } + } + +#if SQLITE_THREADSAFE + if (sqlite3_mutex_enter) { + sqlite3_mutex_enter(sqlite3_db_mutex(p->db)); + entered = 1; + } +#endif + + if (idValue) { + sqlite3_bind_value(p->stmtRowidsInsertId, 1, idValue); + } + rc = sqlite3_step(p->stmtRowidsInsertId); + + if (rc != SQLITE_DONE) { + if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_UNIQUE) { + // IMP: V20497_04568 + vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key", + p->tableName); + } else { + // IMP: V24016_08086 + // IMP: V15177_32015 + vtab_set_error(&p->base, + "Error inserting id into rowids shadow table: %s", + sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId))); + } + rc = SQLITE_ERROR; + goto complete; + } + + *rowid = sqlite3_last_insert_rowid(p->db); + rc = SQLITE_OK; + +complete: + if (p->stmtRowidsInsertId) { + sqlite3_reset(p->stmtRowidsInsertId); + sqlite3_clear_bindings(p->stmtRowidsInsertId); + } + +#if SQLITE_THREADSAFE + if (sqlite3_mutex_leave && entered) { + sqlite3_mutex_leave(sqlite3_db_mutex(p->db)); + } +#endif + return rc; +} + +int vec0_metadata_chunk_size(vec0_metadata_column_kind kind, int chunk_size) { + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: + return chunk_size / 8; + case VEC0_METADATA_COLUMN_KIND_INTEGER: + return chunk_size * sizeof(i64); + case VEC0_METADATA_COLUMN_KIND_FLOAT: + return chunk_size * sizeof(double); + case VEC0_METADATA_COLUMN_KIND_TEXT: + return chunk_size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH; + } + return 0; +} + +int vec0_rowids_update_position(vec0_vtab *p, i64 rowid, i64 chunk_rowid, + i64 chunk_offset) { + int rc = SQLITE_OK; + + if (!p->stmtRowidsUpdatePosition) { + const char *zSql = sqlite3_mprintf(" UPDATE " VEC0_SHADOW_ROWIDS_NAME + " SET chunk_id = ?, chunk_offset = ?" + " WHERE rowid = ?", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsUpdatePosition, 0); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "could not initialize 'update rowids position' statement"); + goto cleanup; + } + } + + sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 1, chunk_rowid); + sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 2, chunk_offset); + sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 3, rowid); + + rc = sqlite3_step(p->stmtRowidsUpdatePosition); + if (rc != SQLITE_DONE) { + // IMP: V21925_05995 + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "could not update rowids position for rowid=%lld, " + "chunk_rowid=%lld, chunk_offset=%lld", + rowid, chunk_rowid, chunk_offset); + rc = SQLITE_ERROR; + goto cleanup; + } + rc = SQLITE_OK; + +cleanup: + if (p->stmtRowidsUpdatePosition) { + sqlite3_reset(p->stmtRowidsUpdatePosition); + sqlite3_clear_bindings(p->stmtRowidsUpdatePosition); + } + + return rc; +} + +/** + * @brief Adds a new chunk for the vec0 table, and the corresponding vector + * chunks. + * + * Inserts a new row into the _chunks table, with blank data, and uses that new + * rowid to insert new blank rows into _vector_chunksXX tables. + * + * @param p: vec0 table to add new chunk + * @param paritionKeyValues: Array of partition key valeus for the new chunk, if available + * @param chunk_rowid: Output pointer, if not NULL, then will be filled with the + * new chunk rowid. + * @return int SQLITE_OK on success, error code otherwise. + */ +int vec0_new_chunk(vec0_vtab *p, sqlite3_value ** partitionKeyValues, i64 *chunk_rowid) { + int rc; + char *zSql; + sqlite3_stmt *stmt; + i64 rowid; + + // Step 1: Insert a new row in _chunks, capture that new rowid + if(p->numPartitionColumns > 0) { + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_CHUNKS_NAME, p->schemaName, p->tableName); + sqlite3_str_appendall(s, "(size, validity, rowids"); + for(int i = 0; i < p->numPartitionColumns; i++) { + sqlite3_str_appendf(s, ", partition%02d", i); + } + sqlite3_str_appendall(s, ") VALUES (?, ?, ?"); + for(int i = 0; i < p->numPartitionColumns; i++) { + sqlite3_str_appendall(s, ", ?"); + } + sqlite3_str_appendall(s, ")"); + + zSql = sqlite3_str_finish(s); + }else { + zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_CHUNKS_NAME + "(size, validity, rowids) " + "VALUES (?, ?, ?);", + p->schemaName, p->tableName); + } + + if (!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) { + sqlite3_finalize(stmt); + return rc; + } + +#if SQLITE_THREADSAFE + if (sqlite3_mutex_enter) { + sqlite3_mutex_enter(sqlite3_db_mutex(p->db)); + } +#endif + + sqlite3_bind_int64(stmt, 1, p->chunk_size); // size + sqlite3_bind_zeroblob(stmt, 2, p->chunk_size / CHAR_BIT); // validity bitmap + sqlite3_bind_zeroblob(stmt, 3, p->chunk_size * sizeof(i64)); // rowids + + for(int i = 0; i < p->numPartitionColumns; i++) { + sqlite3_bind_value(stmt, 4 + i, partitionKeyValues[i]); + } + + rc = sqlite3_step(stmt); + int failed = rc != SQLITE_DONE; + rowid = sqlite3_last_insert_rowid(p->db); +#if SQLITE_THREADSAFE + if (sqlite3_mutex_leave) { + sqlite3_mutex_leave(sqlite3_db_mutex(p->db)); + } +#endif + sqlite3_finalize(stmt); + if (failed) { + return SQLITE_ERROR; + } + + // Step 2: Create new vector chunks for each vector column, with + // that new chunk_rowid. + // + // SHADOW_TABLE_ROWID_QUIRK: The _vector_chunksNN and _metadatachunksNN + // shadow tables declare "rowid PRIMARY KEY" without the INTEGER type, so + // the user-defined "rowid" column is NOT an alias for the internal SQLite + // rowid (_rowid_). When only appending rows these two happen to stay in + // sync, but after a chunk is deleted (vec0Update_Delete_DeleteChunkIfEmpty) + // and a new one is created, the auto-assigned _rowid_ can diverge from the + // user "rowid" value. Since sqlite3_blob_open() addresses rows by internal + // _rowid_, we must explicitly set BOTH _rowid_ and "rowid" to the same + // value so that later blob operations can find the row. + // + // The correct long-term fix is changing the schema to + // "rowid INTEGER PRIMARY KEY" + // which makes it a true alias, but that would break existing databases. + + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) { + continue; + } + int vector_column_idx = p->user_column_idxs[i]; + i64 vectorsSize = + p->chunk_size * vector_column_byte_size(p->vector_columns[vector_column_idx]); + + // See SHADOW_TABLE_ROWID_QUIRK above for why _rowid_ and rowid are both set. + zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_VECTOR_N_NAME + "(_rowid_, rowid, vectors)" + "VALUES (?, ?, ?)", + p->schemaName, p->tableName, vector_column_idx); + if (!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + + if (rc != SQLITE_OK) { + sqlite3_finalize(stmt); + return rc; + } + + sqlite3_bind_int64(stmt, 1, rowid); // _rowid_ (internal SQLite rowid) + sqlite3_bind_int64(stmt, 2, rowid); // rowid (user-defined column) + sqlite3_bind_zeroblob64(stmt, 3, vectorsSize); + + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if (rc != SQLITE_DONE) { + return rc; + } + } + + // Step 3: Create new metadata chunks for each metadata column + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) { + continue; + } + int metadata_column_idx = p->user_column_idxs[i]; + // See SHADOW_TABLE_ROWID_QUIRK above for why _rowid_ and rowid are both set. + zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_N_NAME + "(_rowid_, rowid, data)" + "VALUES (?, ?, ?)", + p->schemaName, p->tableName, metadata_column_idx); + if (!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + + if (rc != SQLITE_OK) { + sqlite3_finalize(stmt); + return rc; + } + + sqlite3_bind_int64(stmt, 1, rowid); // _rowid_ (internal SQLite rowid) + sqlite3_bind_int64(stmt, 2, rowid); // rowid (user-defined column) + sqlite3_bind_zeroblob64(stmt, 3, vec0_metadata_chunk_size(p->metadata_columns[metadata_column_idx].kind, p->chunk_size)); + + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if (rc != SQLITE_DONE) { + return rc; + } + } + + + if (chunk_rowid) { + *chunk_rowid = rowid; + } + + return SQLITE_OK; +} + +struct vec0_query_fullscan_data { + sqlite3_stmt *rowids_stmt; + i8 done; +}; +void vec0_query_fullscan_data_clear( + struct vec0_query_fullscan_data *fullscan_data) { + if (!fullscan_data) + return; + + if (fullscan_data->rowids_stmt) { + sqlite3_finalize(fullscan_data->rowids_stmt); + fullscan_data->rowids_stmt = NULL; + } +} + +struct vec0_query_knn_data { + i64 k; + i64 k_used; + // Array of rowids of size k. Must be freed with sqlite3_free(). + i64 *rowids; + // Array of distances of size k. Must be freed with sqlite3_free(). + f32 *distances; + i64 current_idx; +}; +void vec0_query_knn_data_clear(struct vec0_query_knn_data *knn_data) { + if (!knn_data) + return; + + if (knn_data->rowids) { + sqlite3_free(knn_data->rowids); + knn_data->rowids = NULL; + } + if (knn_data->distances) { + sqlite3_free(knn_data->distances); + knn_data->distances = NULL; + } +} + +struct vec0_query_point_data { + i64 rowid; + void *vectors[VEC0_MAX_VECTOR_COLUMNS]; + int done; +}; +void vec0_query_point_data_clear(struct vec0_query_point_data *point_data) { + if (!point_data) + return; + for (int i = 0; i < VEC0_MAX_VECTOR_COLUMNS; i++) { + sqlite3_free(point_data->vectors[i]); + point_data->vectors[i] = NULL; + } +} + +typedef enum { + // If any values are updated, please update the ARCHITECTURE.md docs accordingly! + + VEC0_QUERY_PLAN_FULLSCAN = '1', + VEC0_QUERY_PLAN_POINT = '2', + VEC0_QUERY_PLAN_KNN = '3', +} vec0_query_plan; + +typedef struct vec0_cursor vec0_cursor; +struct vec0_cursor { + sqlite3_vtab_cursor base; + + vec0_query_plan query_plan; + struct vec0_query_fullscan_data *fullscan_data; + struct vec0_query_knn_data *knn_data; + struct vec0_query_point_data *point_data; +}; + +void vec0_cursor_clear(vec0_cursor *pCur) { + if (pCur->fullscan_data) { + vec0_query_fullscan_data_clear(pCur->fullscan_data); + sqlite3_free(pCur->fullscan_data); + pCur->fullscan_data = NULL; + } + if (pCur->knn_data) { + vec0_query_knn_data_clear(pCur->knn_data); + sqlite3_free(pCur->knn_data); + pCur->knn_data = NULL; + } + if (pCur->point_data) { + vec0_query_point_data_clear(pCur->point_data); + sqlite3_free(pCur->point_data); + pCur->point_data = NULL; + } +} + +#define VEC_CONSTRUCTOR_ERROR "vec0 constructor error: " +static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv, + sqlite3_vtab **ppVtab, char **pzErr, bool isCreate) { + UNUSED_PARAMETER(pAux); + vec0_vtab *pNew; + int rc; + const char *zSql; + + pNew = sqlite3_malloc(sizeof(*pNew)); + if (pNew == 0) + return SQLITE_NOMEM; + memset(pNew, 0, sizeof(*pNew)); + + // Declared chunk_size=N for entire table. + // -1 to use the defualt, otherwise will get re-assigned on `chunk_size=N` + // option + int chunk_size = -1; + int numVectorColumns = 0; + int numPartitionColumns = 0; + int numAuxiliaryColumns = 0; + int numMetadataColumns = 0; + int user_column_idx = 0; + + // track if a "primary key" column is defined + char *pkColumnName = NULL; + int pkColumnNameLength; + int pkColumnType = SQLITE_INTEGER; + + for (int i = 3; i < argc; i++) { + struct VectorColumnDefinition vecColumn; + struct Vec0PartitionColumnDefinition partitionColumn; + struct Vec0AuxiliaryColumnDefinition auxColumn; + struct Vec0MetadataColumnDefinition metadataColumn; + char *cName = NULL; + int cNameLength; + int cType; + + // Scenario #1: Constructor argument is a vector column definition, ie `foo float[1024]` + rc = vec0_parse_vector_column(argv[i], strlen(argv[i]), &vecColumn); + if (rc == SQLITE_ERROR) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR "could not parse vector column '%s'", argv[i]); + goto error; + } + if (rc == SQLITE_OK) { + if (numVectorColumns >= VEC0_MAX_VECTOR_COLUMNS) { + sqlite3_free(vecColumn.name); + *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "Too many provided vector columns, maximum %d", + VEC0_MAX_VECTOR_COLUMNS); + goto error; + } + + if (vecColumn.dimensions > SQLITE_VEC_VEC0_MAX_DIMENSIONS) { + sqlite3_free(vecColumn.name); + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR + "Dimension on vector column too large, provided %lld, maximum %lld", + (i64)vecColumn.dimensions, SQLITE_VEC_VEC0_MAX_DIMENSIONS); + goto error; + } + pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_VECTOR; + pNew->user_column_idxs[user_column_idx] = numVectorColumns; + memcpy(&pNew->vector_columns[numVectorColumns], &vecColumn, sizeof(vecColumn)); + numVectorColumns++; + pNew->numVectorColumns = numVectorColumns; + user_column_idx++; + + continue; + } + + // Scenario #2: Constructor argument is a partition key column definition, ie `user_id text partition key` + rc = vec0_parse_partition_key_definition(argv[i], strlen(argv[i]), &cName, + &cNameLength, &cType); + if (rc == SQLITE_OK) { + if (numPartitionColumns >= VEC0_MAX_PARTITION_COLUMNS) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR + "More than %d partition key columns were provided", + VEC0_MAX_PARTITION_COLUMNS); + goto error; + } + partitionColumn.type = cType; + partitionColumn.name_length = cNameLength; + partitionColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName); + if(!partitionColumn.name) { + rc = SQLITE_NOMEM; + goto error; + } + + pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_PARTITION; + pNew->user_column_idxs[user_column_idx] = numPartitionColumns; + memcpy(&pNew->paritition_columns[numPartitionColumns], &partitionColumn, sizeof(partitionColumn)); + numPartitionColumns++; + pNew->numPartitionColumns = numPartitionColumns; + user_column_idx++; + continue; + } + + // Scenario #3: Constructor argument is a primary key column definition, ie `article_id text primary key` + rc = vec0_parse_primary_key_definition(argv[i], strlen(argv[i]), &cName, + &cNameLength, &cType); + if (rc == SQLITE_OK) { + if (pkColumnName) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR + "More than one primary key definition was provided, vec0 only " + "suports a single primary key column", + argv[i]); + goto error; + } + pkColumnName = cName; + pkColumnNameLength = cNameLength; + pkColumnType = cType; + continue; + } + + // Scenario #4: Constructor argument is a auxiliary column definition, ie `+contents text` + rc = vec0_parse_auxiliary_column_definition(argv[i], strlen(argv[i]), &cName, + &cNameLength, &cType); + if(rc == SQLITE_OK) { + if (numAuxiliaryColumns >= VEC0_MAX_AUXILIARY_COLUMNS) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR + "More than %d auxiliary columns were provided", + VEC0_MAX_AUXILIARY_COLUMNS); + goto error; + } + auxColumn.type = cType; + auxColumn.name_length = cNameLength; + auxColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName); + if(!auxColumn.name) { + rc = SQLITE_NOMEM; + goto error; + } + + pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY; + pNew->user_column_idxs[user_column_idx] = numAuxiliaryColumns; + memcpy(&pNew->auxiliary_columns[numAuxiliaryColumns], &auxColumn, sizeof(auxColumn)); + numAuxiliaryColumns++; + pNew->numAuxiliaryColumns = numAuxiliaryColumns; + user_column_idx++; + continue; + } + + vec0_metadata_column_kind kind; + rc = vec0_parse_metadata_column_definition(argv[i], strlen(argv[i]), &cName, + &cNameLength, &kind); + if(rc == SQLITE_OK) { + if (numMetadataColumns >= VEC0_MAX_METADATA_COLUMNS) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR + "More than %d metadata columns were provided", + VEC0_MAX_METADATA_COLUMNS); + goto error; + } + metadataColumn.kind = kind; + metadataColumn.name_length = cNameLength; + metadataColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName); + if(!metadataColumn.name) { + rc = SQLITE_NOMEM; + goto error; + } + + pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_METADATA; + pNew->user_column_idxs[user_column_idx] = numMetadataColumns; + memcpy(&pNew->metadata_columns[numMetadataColumns], &metadataColumn, sizeof(metadataColumn)); + numMetadataColumns++; + pNew->numMetadataColumns = numMetadataColumns; + user_column_idx++; + continue; + } + + // Scenario #4: Constructor argument is a table-level option, ie `chunk_size` + + char *key; + char *value; + int keyLength, valueLength; + rc = vec0_parse_table_option(argv[i], strlen(argv[i]), &key, &keyLength, + &value, &valueLength); + if (rc == SQLITE_ERROR) { + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR "could not parse table option '%s'", argv[i]); + goto error; + } + if (rc == SQLITE_OK) { + if (sqlite3_strnicmp(key, "chunk_size", keyLength) == 0) { + chunk_size = atoi(value); + if (chunk_size <= 0) { + // IMP: V01931_18769 + *pzErr = + sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "chunk_size must be a non-zero positive integer"); + goto error; + } + if ((chunk_size % 8) != 0) { + // IMP: V14110_30948 + *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "chunk_size must be divisible by 8"); + goto error; + } +#define SQLITE_VEC_CHUNK_SIZE_MAX 4096 + if (chunk_size > SQLITE_VEC_CHUNK_SIZE_MAX) { + *pzErr = + sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "chunk_size too large"); + goto error; + } + } else { + // IMP: V27642_11712 + *pzErr = sqlite3_mprintf( + VEC_CONSTRUCTOR_ERROR "Unknown table option: %.*s", keyLength, key); + goto error; + } + continue; + } + + // Scenario #5: Unknown constructor argument + *pzErr = + sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "Could not parse '%s'", argv[i]); + goto error; + } + + if (chunk_size < 0) { + chunk_size = 1024; + } + + if (numVectorColumns <= 0) { + *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "At least one vector column is required"); + goto error; + } + + sqlite3_str *createStr = sqlite3_str_new(NULL); + sqlite3_str_appendall(createStr, "CREATE TABLE x("); + if (pkColumnName) { + sqlite3_str_appendf(createStr, "\"%.*w\" primary key, ", pkColumnNameLength, + pkColumnName); + } else { + sqlite3_str_appendall(createStr, "rowid, "); + } + for (int i = 0; i < numVectorColumns + numPartitionColumns + numAuxiliaryColumns + numMetadataColumns; i++) { + switch(pNew->user_column_kinds[i]) { + case SQLITE_VEC0_USER_COLUMN_KIND_VECTOR: { + int vector_idx = pNew->user_column_idxs[i]; + sqlite3_str_appendf(createStr, "\"%.*w\", ", + pNew->vector_columns[vector_idx].name_length, + pNew->vector_columns[vector_idx].name); + break; + } + case SQLITE_VEC0_USER_COLUMN_KIND_PARTITION: { + int partition_idx = pNew->user_column_idxs[i]; + sqlite3_str_appendf(createStr, "\"%.*w\", ", + pNew->paritition_columns[partition_idx].name_length, + pNew->paritition_columns[partition_idx].name); + break; + } + case SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY: { + int auxiliary_idx = pNew->user_column_idxs[i]; + sqlite3_str_appendf(createStr, "\"%.*w\", ", + pNew->auxiliary_columns[auxiliary_idx].name_length, + pNew->auxiliary_columns[auxiliary_idx].name); + break; + } + case SQLITE_VEC0_USER_COLUMN_KIND_METADATA: { + int metadata_idx = pNew->user_column_idxs[i]; + sqlite3_str_appendf(createStr, "\"%.*w\", ", + pNew->metadata_columns[metadata_idx].name_length, + pNew->metadata_columns[metadata_idx].name); + break; + } + } + + } + sqlite3_str_appendall(createStr, " distance hidden, k hidden) "); + if (pkColumnName) { + sqlite3_str_appendall(createStr, "without rowid "); + } + zSql = sqlite3_str_finish(createStr); + if (!zSql) { + goto error; + } + rc = sqlite3_declare_vtab(db, zSql); + sqlite3_free((void *)zSql); + if (rc != SQLITE_OK) { + *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "could not declare virtual table, '%s'", + sqlite3_errmsg(db)); + goto error; + } + + const char *schemaName = argv[1]; + const char *tableName = argv[2]; + + pNew->db = db; + pNew->pkIsText = pkColumnType == SQLITE_TEXT; + pNew->schemaName = sqlite3_mprintf("%s", schemaName); + if (!pNew->schemaName) { + goto error; + } + pNew->tableName = sqlite3_mprintf("%s", tableName); + if (!pNew->tableName) { + goto error; + } + pNew->shadowRowidsName = sqlite3_mprintf("%s_rowids", tableName); + if (!pNew->shadowRowidsName) { + goto error; + } + pNew->shadowChunksName = sqlite3_mprintf("%s_chunks", tableName); + if (!pNew->shadowChunksName) { + goto error; + } + pNew->numVectorColumns = numVectorColumns; + pNew->numPartitionColumns = numPartitionColumns; + pNew->numAuxiliaryColumns = numAuxiliaryColumns; + pNew->numMetadataColumns = numMetadataColumns; + + for (int i = 0; i < pNew->numVectorColumns; i++) { + pNew->shadowVectorChunksNames[i] = + sqlite3_mprintf("%s_vector_chunks%02d", tableName, i); + if (!pNew->shadowVectorChunksNames[i]) { + goto error; + } + } + for (int i = 0; i < pNew->numMetadataColumns; i++) { + pNew->shadowMetadataChunksNames[i] = + sqlite3_mprintf("%s_metadatachunks%02d", tableName, i); + if (!pNew->shadowMetadataChunksNames[i]) { + goto error; + } + } + pNew->chunk_size = chunk_size; + + // if xCreate, then create the necessary shadow tables + if (isCreate) { + sqlite3_stmt *stmt; + int rc; + + char * zCreateInfo = sqlite3_mprintf("CREATE TABLE "VEC0_SHADOW_INFO_NAME " (key text primary key, value any)", pNew->schemaName, pNew->tableName); + if(!zCreateInfo) { + goto error; + } + rc = sqlite3_prepare_v2(db, zCreateInfo, -1, &stmt, NULL); + + sqlite3_free((void *) zCreateInfo); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + // TODO(IMP) + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf("Could not create '_info' shadow table: %s", + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + char * zSeedInfo = sqlite3_mprintf( + "INSERT INTO "VEC0_SHADOW_INFO_NAME "(key, value) VALUES " + "(?1, ?2), (?3, ?4), (?5, ?6), (?7, ?8) ", + pNew->schemaName, pNew->tableName + ); + if(!zSeedInfo) { + goto error; + } + rc = sqlite3_prepare_v2(db, zSeedInfo, -1, &stmt, NULL); + sqlite3_free((void *) zSeedInfo); + if (rc != SQLITE_OK) { + // TODO(IMP) + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s", + sqlite3_errmsg(db)); + goto error; + } + sqlite3_bind_text(stmt, 1, "CREATE_VERSION", -1, SQLITE_STATIC); + sqlite3_bind_text(stmt, 2, SQLITE_VEC_VERSION, -1, SQLITE_STATIC); + sqlite3_bind_text(stmt, 3, "CREATE_VERSION_MAJOR", -1, SQLITE_STATIC); + sqlite3_bind_int(stmt, 4, SQLITE_VEC_VERSION_MAJOR); + sqlite3_bind_text(stmt, 5, "CREATE_VERSION_MINOR", -1, SQLITE_STATIC); + sqlite3_bind_int(stmt, 6, SQLITE_VEC_VERSION_MINOR); + sqlite3_bind_text(stmt, 7, "CREATE_VERSION_PATCH", -1, SQLITE_STATIC); + sqlite3_bind_int(stmt, 8, SQLITE_VEC_VERSION_PATCH); + + if(sqlite3_step(stmt) != SQLITE_DONE) { + // TODO(IMP) + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s", + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + + + // create the _chunks shadow table + char *zCreateShadowChunks = NULL; + if(pNew->numPartitionColumns) { + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(", pNew->schemaName, pNew->tableName); + sqlite3_str_appendall(s, "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," "size INTEGER NOT NULL,"); + sqlite3_str_appendall(s, "sequence_id integer,"); + for(int i = 0; i < pNew->numPartitionColumns;i++) { + sqlite3_str_appendf(s, "partition%02d,", i); + } + sqlite3_str_appendall(s, "validity BLOB NOT NULL, rowids BLOB NOT NULL);"); + zCreateShadowChunks = sqlite3_str_finish(s); + }else { + zCreateShadowChunks = sqlite3_mprintf(VEC0_SHADOW_CHUNKS_CREATE, + pNew->schemaName, pNew->tableName); + } + if (!zCreateShadowChunks) { + goto error; + } + rc = sqlite3_prepare_v2(db, zCreateShadowChunks, -1, &stmt, 0); + sqlite3_free((void *)zCreateShadowChunks); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + // IMP: V17740_01811 + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf("Could not create '_chunks' shadow table: %s", + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + // create the _rowids shadow table + char *zCreateShadowRowids; + if (pNew->pkIsText) { + // adds a "text unique not null" constraint to the id column + zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT, + pNew->schemaName, pNew->tableName); + } else { + zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_BASIC, + pNew->schemaName, pNew->tableName); + } + if (!zCreateShadowRowids) { + goto error; + } + rc = sqlite3_prepare_v2(db, zCreateShadowRowids, -1, &stmt, 0); + sqlite3_free((void *)zCreateShadowRowids); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + // IMP: V11631_28470 + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf("Could not create '_rowids' shadow table: %s", + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + for (int i = 0; i < pNew->numVectorColumns; i++) { + char *zSql = sqlite3_mprintf(VEC0_SHADOW_VECTOR_N_CREATE, + pNew->schemaName, pNew->tableName, i); + if (!zSql) { + goto error; + } + rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + // IMP: V25919_09989 + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf( + "Could not create '_vector_chunks%02d' shadow table: %s", i, + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + } + + // See SHADOW_TABLE_ROWID_QUIRK in vec0_new_chunk() — same "rowid PRIMARY KEY" + // without INTEGER type issue applies here. + for (int i = 0; i < pNew->numMetadataColumns; i++) { + char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_N_NAME "(rowid PRIMARY KEY, data BLOB NOT NULL);", + pNew->schemaName, pNew->tableName, i); + if (!zSql) { + goto error; + } + rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf( + "Could not create '_metata_chunks%02d' shadow table: %s", i, + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + if(pNew->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) { + char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME "(rowid PRIMARY KEY, data TEXT);", + pNew->schemaName, pNew->tableName, i); + if (!zSql) { + goto error; + } + rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf( + "Could not create '_metadatatext%02d' shadow table: %s", i, + sqlite3_errmsg(db)); + goto error; + } + sqlite3_finalize(stmt); + + } + } + + if(pNew->numAuxiliaryColumns > 0) { + sqlite3_stmt * stmt; + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_AUXILIARY_NAME "( rowid integer PRIMARY KEY ", pNew->schemaName, pNew->tableName); + for(int i = 0; i < pNew->numAuxiliaryColumns; i++) { + sqlite3_str_appendf(s, ", value%02d", i); + } + sqlite3_str_appendall(s, ")"); + char *zSql = sqlite3_str_finish(s); + if(!zSql) { + goto error; + } + rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + sqlite3_finalize(stmt); + *pzErr = sqlite3_mprintf( + "Could not create auxiliary shadow table: %s", + sqlite3_errmsg(db)); + + goto error; + } + sqlite3_finalize(stmt); + } + } + + *ppVtab = (sqlite3_vtab *)pNew; + return SQLITE_OK; + +error: + vec0_free(pNew); + sqlite3_free(pNew); + return SQLITE_ERROR; +} + +static int vec0Create(sqlite3 *db, void *pAux, int argc, + const char *const *argv, sqlite3_vtab **ppVtab, + char **pzErr) { + return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, true); +} +static int vec0Connect(sqlite3 *db, void *pAux, int argc, + const char *const *argv, sqlite3_vtab **ppVtab, + char **pzErr) { + return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, false); +} + +static int vec0Disconnect(sqlite3_vtab *pVtab) { + vec0_vtab *p = (vec0_vtab *)pVtab; + vec0_free(p); + sqlite3_free(p); + return SQLITE_OK; +} +static int vec0Destroy(sqlite3_vtab *pVtab) { + vec0_vtab *p = (vec0_vtab *)pVtab; + sqlite3_stmt *stmt; + int rc; + const char *zSql; + + // Free up any sqlite3_stmt, otherwise DROPs on those tables will fail + vec0_free_resources(p); + + // TODO(test) later: can't evidence-of here, bc always gives "SQL logic error" instead of + // provided error + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_CHUNKS_NAME, p->schemaName, + p->tableName); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + vtab_set_error(pVtab, "could not drop chunks shadow table"); + goto done; + } + sqlite3_finalize(stmt); + + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_INFO_NAME, p->schemaName, + p->tableName); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + vtab_set_error(pVtab, "could not drop info shadow table"); + goto done; + } + sqlite3_finalize(stmt); + + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_ROWIDS_NAME, p->schemaName, + p->tableName); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + goto done; + } + sqlite3_finalize(stmt); + + for (int i = 0; i < p->numVectorColumns; i++) { + zSql = sqlite3_mprintf("DROP TABLE \"%w\".\"%w\"", p->schemaName, + p->shadowVectorChunksNames[i]); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + goto done; + } + sqlite3_finalize(stmt); + } + + if(p->numAuxiliaryColumns > 0) { + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_AUXILIARY_NAME, p->schemaName, p->tableName); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + goto done; + } + sqlite3_finalize(stmt); + } + + + for (int i = 0; i < p->numMetadataColumns; i++) { + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_N_NAME, p->schemaName,p->tableName, i); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + goto done; + } + sqlite3_finalize(stmt); + + if(p->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) { + zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME, p->schemaName,p->tableName, i); + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0); + sqlite3_free((void *)zSql); + if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) { + rc = SQLITE_ERROR; + goto done; + } + sqlite3_finalize(stmt); + } + } + + stmt = NULL; + rc = SQLITE_OK; + +done: + sqlite3_finalize(stmt); + vec0_free(p); + // If there was an error + if (rc == SQLITE_OK) { + sqlite3_free(p); + } + return rc; +} + +static int vec0Open(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) { + UNUSED_PARAMETER(p); + vec0_cursor *pCur; + pCur = sqlite3_malloc(sizeof(*pCur)); + if (pCur == 0) + return SQLITE_NOMEM; + memset(pCur, 0, sizeof(*pCur)); + *ppCursor = &pCur->base; + return SQLITE_OK; +} + +static int vec0Close(sqlite3_vtab_cursor *cur) { + vec0_cursor *pCur = (vec0_cursor *)cur; + vec0_cursor_clear(pCur); + sqlite3_free(pCur); + return SQLITE_OK; +} + +// All the different type of "values" provided to argv/argc in vec0Filter. +// These enums denote the use and purpose of all of them. +typedef enum { + // If any values are updated, please update the ARCHITECTURE.md docs accordingly! + + // ~~~ KNN QUERIES ~~~ // + VEC0_IDXSTR_KIND_KNN_MATCH = '{', + VEC0_IDXSTR_KIND_KNN_K = '}', + VEC0_IDXSTR_KIND_KNN_ROWID_IN = '[', + // argv[i] is a constraint on a PARTITON KEY column in a KNN query + // + VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT = ']', + + // argv[i] is a constraint on the distance column in a KNN query + VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT = '*', + + // ~~~ POINT QUERIES ~~~ // + VEC0_IDXSTR_KIND_POINT_ID = '!', + + // ~~~ ??? ~~~ // + VEC0_IDXSTR_KIND_METADATA_CONSTRAINT = '&', +} vec0_idxstr_kind; + +// The different SQLITE_INDEX_CONSTRAINT values that vec0 partition key columns +// support, but as characters that fit nicely in idxstr. +typedef enum { + // If any values are updated, please update the ARCHITECTURE.md docs accordingly! + + // Equality constraint on a PARTITON KEY column, ex `user_id = 123` + VEC0_PARTITION_OPERATOR_EQ = 'a', + + // "Greater than" constraint on a PARTITON KEY column, ex `year > 2024` + VEC0_PARTITION_OPERATOR_GT = 'b', + + // "Less than or equal to" constraint on a PARTITON KEY column, ex `year <= 2024` + VEC0_PARTITION_OPERATOR_LE = 'c', + + // "Less than" constraint on a PARTITON KEY column, ex `year < 2024` + VEC0_PARTITION_OPERATOR_LT = 'd', + + // "Greater than or equal to" constraint on a PARTITON KEY column, ex `year >= 2024` + VEC0_PARTITION_OPERATOR_GE = 'e', + + // "Not equal to" constraint on a PARTITON KEY column, ex `year != 2024` + VEC0_PARTITION_OPERATOR_NE = 'f', +} vec0_partition_operator; +typedef enum { + VEC0_METADATA_OPERATOR_EQ = 'a', + VEC0_METADATA_OPERATOR_GT = 'b', + VEC0_METADATA_OPERATOR_LE = 'c', + VEC0_METADATA_OPERATOR_LT = 'd', + VEC0_METADATA_OPERATOR_GE = 'e', + VEC0_METADATA_OPERATOR_NE = 'f', + VEC0_METADATA_OPERATOR_IN = 'g', +} vec0_metadata_operator; + + +typedef enum { + + VEC0_DISTANCE_CONSTRAINT_GT = 'a', + VEC0_DISTANCE_CONSTRAINT_GE = 'b', + VEC0_DISTANCE_CONSTRAINT_LT = 'c', + VEC0_DISTANCE_CONSTRAINT_LE = 'd', +} vec0_distance_constraint_operator; + +static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { + vec0_vtab *p = (vec0_vtab *)pVTab; + /** + * Possible query plans are: + * 1. KNN when: + * a) An `MATCH` op on vector column + * b) ORDER BY on distance column + * c) LIMIT + * d) rowid in (...) OPTIONAL + * 2. Point when: + * a) An `EQ` op on rowid column + * 3. else: fullscan + * + */ + int iMatchTerm = -1; + int iMatchVectorTerm = -1; + int iLimitTerm = -1; + int iRowidTerm = -1; + int iKTerm = -1; + int iRowidInTerm = -1; + int hasAuxConstraint = 0; + +#ifdef SQLITE_VEC_DEBUG + printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint); +#endif + + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + u8 vtabIn = 0; + +#if COMPILER_SUPPORTS_VTAB_IN + if (sqlite3_libversion_number() >= 3038000) { + vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1); + } +#endif + +#ifdef SQLITE_VEC_DEBUG + printf("xBestIndex [%d] usable=%d iColumn=%d op=%d vtabin=%d\n", i, + pIdxInfo->aConstraint[i].usable, pIdxInfo->aConstraint[i].iColumn, + pIdxInfo->aConstraint[i].op, vtabIn); +#endif + if (!pIdxInfo->aConstraint[i].usable) + continue; + + int iColumn = pIdxInfo->aConstraint[i].iColumn; + int op = pIdxInfo->aConstraint[i].op; + + if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) { + iLimitTerm = i; + } + if (op == SQLITE_INDEX_CONSTRAINT_MATCH && + vec0_column_idx_is_vector(p, iColumn)) { + if (iMatchTerm > -1) { + vtab_set_error( + pVTab, "only 1 MATCH operator is allowed in a single vec0 query"); + return SQLITE_ERROR; + } + iMatchTerm = i; + iMatchVectorTerm = vec0_column_idx_to_vector_idx(p, iColumn); + } + if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == VEC0_COLUMN_ID) { + if (vtabIn) { + if (iRowidInTerm != -1) { + vtab_set_error(pVTab, "only 1 'rowid in (..)' operator is allowed in " + "a single vec0 query"); + return SQLITE_ERROR; + } + iRowidInTerm = i; + + } else { + iRowidTerm = i; + } + } + if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == vec0_column_k_idx(p)) { + iKTerm = i; + } + if( + (op != SQLITE_INDEX_CONSTRAINT_LIMIT && op != SQLITE_INDEX_CONSTRAINT_OFFSET) + && vec0_column_idx_is_auxiliary(p, iColumn)) { + hasAuxConstraint = 1; + } + } + + sqlite3_str *idxStr = sqlite3_str_new(NULL); + int rc; + + if (iMatchTerm >= 0) { + if (iLimitTerm < 0 && iKTerm < 0) { + vtab_set_error( + pVTab, + "A LIMIT or 'k = ?' constraint is required on vec0 knn queries."); + rc = SQLITE_ERROR; + goto done; + } + if (iLimitTerm >= 0 && iKTerm >= 0) { + vtab_set_error(pVTab, "Only LIMIT or 'k =?' can be provided, not both"); + rc = SQLITE_ERROR; + goto done; + } + + if (pIdxInfo->nOrderBy) { + if (pIdxInfo->nOrderBy > 1) { + vtab_set_error(pVTab, "Only a single 'ORDER BY distance' clause is " + "allowed on vec0 KNN queries"); + rc = SQLITE_ERROR; + goto done; + } + if (pIdxInfo->aOrderBy[0].iColumn != vec0_column_distance_idx(p)) { + vtab_set_error(pVTab, + "Only a single 'ORDER BY distance' clause is allowed on " + "vec0 KNN queries, not on other columns"); + rc = SQLITE_ERROR; + goto done; + } + if (pIdxInfo->aOrderBy[0].desc) { + vtab_set_error( + pVTab, "Only ascending in ORDER BY distance clause is supported, " + "DESC is not supported yet."); + rc = SQLITE_ERROR; + goto done; + } + } + + if(hasAuxConstraint) { + // IMP: V25623_09693 + vtab_set_error(pVTab, "An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query."); + rc = SQLITE_ERROR; + goto done; + } + + sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_KNN); + + int argvIndex = 1; + pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_MATCH); + sqlite3_str_appendchar(idxStr, 3, '_'); + + if (iLimitTerm >= 0) { + pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1; + } else { + pIdxInfo->aConstraintUsage[iKTerm].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[iKTerm].omit = 1; + } + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_K); + sqlite3_str_appendchar(idxStr, 3, '_'); + +#if COMPILER_SUPPORTS_VTAB_IN + if (iRowidInTerm >= 0) { + // already validated as >= SQLite 3.38 bc iRowidInTerm is only >= 0 when + // vtabIn == 1 + sqlite3_vtab_in(pIdxInfo, iRowidInTerm, 1); + pIdxInfo->aConstraintUsage[iRowidInTerm].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[iRowidInTerm].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_ROWID_IN); + sqlite3_str_appendchar(idxStr, 3, '_'); + } +#endif + + // find any PARTITION KEY column constraints + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + if (!pIdxInfo->aConstraint[i].usable) + continue; + + int iColumn = pIdxInfo->aConstraint[i].iColumn; + int op = pIdxInfo->aConstraint[i].op; + if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) { + continue; + } + if(!vec0_column_idx_is_partition(p, iColumn)) { + continue; + } + + int partition_idx = vec0_column_idx_to_partition_idx(p, iColumn); + char value = 0; + + switch(op) { + case SQLITE_INDEX_CONSTRAINT_EQ: { + value = VEC0_PARTITION_OPERATOR_EQ; + break; + } + case SQLITE_INDEX_CONSTRAINT_GT: { + value = VEC0_PARTITION_OPERATOR_GT; + break; + } + case SQLITE_INDEX_CONSTRAINT_LE: { + value = VEC0_PARTITION_OPERATOR_LE; + break; + } + case SQLITE_INDEX_CONSTRAINT_LT: { + value = VEC0_PARTITION_OPERATOR_LT; + break; + } + case SQLITE_INDEX_CONSTRAINT_GE: { + value = VEC0_PARTITION_OPERATOR_GE; + break; + } + case SQLITE_INDEX_CONSTRAINT_NE: { + value = VEC0_PARTITION_OPERATOR_NE; + break; + } + } + + if(value) { + pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[i].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT); + sqlite3_str_appendchar(idxStr, 1, 'A' + partition_idx); + sqlite3_str_appendchar(idxStr, 1, value); + sqlite3_str_appendchar(idxStr, 1, '_'); + } + + } + + // find any metadata column constraints + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + if (!pIdxInfo->aConstraint[i].usable) + continue; + + int iColumn = pIdxInfo->aConstraint[i].iColumn; + int op = pIdxInfo->aConstraint[i].op; + if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) { + continue; + } + if(!vec0_column_idx_is_metadata(p, iColumn)) { + continue; + } + + int metadata_idx = vec0_column_idx_to_metadata_idx(p, iColumn); + char value = 0; + + switch(op) { + case SQLITE_INDEX_CONSTRAINT_EQ: { + int vtabIn = 0; + #if COMPILER_SUPPORTS_VTAB_IN + if (sqlite3_libversion_number() >= 3038000) { + vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1); + } + if(vtabIn) { + switch(p->metadata_columns[metadata_idx].kind) { + case VEC0_METADATA_COLUMN_KIND_FLOAT: + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + // IMP: V15248_32086 + rc = SQLITE_ERROR; + vtab_set_error(pVTab, "'xxx in (...)' is only available on INTEGER or TEXT metadata columns."); + goto done; + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: + case VEC0_METADATA_COLUMN_KIND_TEXT: { + break; + } + } + value = VEC0_METADATA_OPERATOR_IN; + sqlite3_vtab_in(pIdxInfo, i, 1); + }else + #endif + { + value = VEC0_PARTITION_OPERATOR_EQ; + } + break; + } + case SQLITE_INDEX_CONSTRAINT_GT: { + value = VEC0_METADATA_OPERATOR_GT; + break; + } + case SQLITE_INDEX_CONSTRAINT_LE: { + value = VEC0_METADATA_OPERATOR_LE; + break; + } + case SQLITE_INDEX_CONSTRAINT_LT: { + value = VEC0_METADATA_OPERATOR_LT; + break; + } + case SQLITE_INDEX_CONSTRAINT_GE: { + value = VEC0_METADATA_OPERATOR_GE; + break; + } + case SQLITE_INDEX_CONSTRAINT_NE: { + value = VEC0_METADATA_OPERATOR_NE; + break; + } + default: { + // IMP: V16511_00582 + rc = SQLITE_ERROR; + vtab_set_error(pVTab, + "An illegal WHERE constraint was provided on a vec0 metadata column in a KNN query. " + "Only one of EQUALS, GREATER_THAN, LESS_THAN_OR_EQUAL, LESS_THAN, GREATER_THAN_OR_EQUAL, NOT_EQUALS is allowed." + ); + goto done; + } + } + + if(p->metadata_columns[metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_BOOLEAN) { + if(!(value == VEC0_METADATA_OPERATOR_EQ || value == VEC0_METADATA_OPERATOR_NE)) { + // IMP: V10145_26984 + rc = SQLITE_ERROR; + vtab_set_error(pVTab, "ONLY EQUALS (=) or NOT_EQUALS (!=) operators are allowed on boolean metadata columns."); + goto done; + } + } + + pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[i].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_METADATA_CONSTRAINT); + sqlite3_str_appendchar(idxStr, 1, 'A' + metadata_idx); + sqlite3_str_appendchar(idxStr, 1, value); + sqlite3_str_appendchar(idxStr, 1, '_'); + + } + + // find any distance column constraints + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + if (!pIdxInfo->aConstraint[i].usable) + continue; + + int iColumn = pIdxInfo->aConstraint[i].iColumn; + int op = pIdxInfo->aConstraint[i].op; + if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) { + continue; + } + if(vec0_column_distance_idx(p) != iColumn) { + continue; + } + + char value = 0; + switch(op) { + case SQLITE_INDEX_CONSTRAINT_GT: { + value = VEC0_DISTANCE_CONSTRAINT_GT; + break; + } + case SQLITE_INDEX_CONSTRAINT_GE: { + value = VEC0_DISTANCE_CONSTRAINT_GE; + break; + } + case SQLITE_INDEX_CONSTRAINT_LT: { + value = VEC0_DISTANCE_CONSTRAINT_LT; + break; + } + case SQLITE_INDEX_CONSTRAINT_LE: { + value = VEC0_DISTANCE_CONSTRAINT_LE; + break; + } + default: { + // IMP TODO + rc = SQLITE_ERROR; + vtab_set_error( + pVTab, + "Illegal WHERE constraint on distance column in a KNN query. " + "Only one of GT, GE, LT, LE constraints are allowed." + ); + goto done; + } + } + + pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++; + pIdxInfo->aConstraintUsage[i].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT); + sqlite3_str_appendchar(idxStr, 1, value); + sqlite3_str_appendchar(idxStr, 1, '_'); + sqlite3_str_appendchar(idxStr, 1, '_'); + } + + + + pIdxInfo->idxNum = iMatchVectorTerm; + pIdxInfo->estimatedCost = 30.0; + pIdxInfo->estimatedRows = 10; + + } else if (iRowidTerm >= 0) { + sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_POINT); + pIdxInfo->aConstraintUsage[iRowidTerm].argvIndex = 1; + pIdxInfo->aConstraintUsage[iRowidTerm].omit = 1; + sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_POINT_ID); + sqlite3_str_appendchar(idxStr, 3, '_'); + pIdxInfo->idxNum = pIdxInfo->colUsed; + pIdxInfo->estimatedCost = 10.0; + pIdxInfo->estimatedRows = 1; + } else { + sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_FULLSCAN); + pIdxInfo->estimatedCost = 3000000.0; + pIdxInfo->estimatedRows = 100000; + } + pIdxInfo->idxStr = sqlite3_str_finish(idxStr); + idxStr = NULL; + if (!pIdxInfo->idxStr) { + rc = SQLITE_OK; + goto done; + } + pIdxInfo->needToFreeIdxStr = 1; + + rc = SQLITE_OK; + + done: + if(idxStr) { + sqlite3_str_finish(idxStr); + } + return rc; +} + +// forward delcaration bc vec0Filter uses it +static int vec0Next(sqlite3_vtab_cursor *cur); + +void merge_sorted_lists(f32 *a, i64 *a_rowids, i64 a_length, f32 *b, + i64 *b_rowids, i32 *b_top_idxs, i64 b_length, f32 *out, + i64 *out_rowids, i64 out_length, i64 *out_used) { + // assert((a_length >= out_length) || (b_length >= out_length)); + i64 ptrA = 0; + i64 ptrB = 0; + for (int i = 0; i < out_length; i++) { + if ((ptrA >= a_length) && (ptrB >= b_length)) { + *out_used = i; + return; + } + if (ptrA >= a_length) { + out[i] = b[b_top_idxs[ptrB]]; + out_rowids[i] = b_rowids[b_top_idxs[ptrB]]; + ptrB++; + } else if (ptrB >= b_length) { + out[i] = a[ptrA]; + out_rowids[i] = a_rowids[ptrA]; + ptrA++; + } else { + if (a[ptrA] <= b[b_top_idxs[ptrB]]) { + out[i] = a[ptrA]; + out_rowids[i] = a_rowids[ptrA]; + ptrA++; + } else { + out[i] = b[b_top_idxs[ptrB]]; + out_rowids[i] = b_rowids[b_top_idxs[ptrB]]; + ptrB++; + } + } + } + + *out_used = out_length; +} + +u8 *bitmap_new(i32 n) { + assert(n % 8 == 0); + u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT); + if (p) { + memset(p, 0, n * sizeof(u8) / CHAR_BIT); + } + return p; +} +u8 *bitmap_new_from(i32 n, u8 *from) { + assert(n % 8 == 0); + u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT); + if (p) { + memcpy(p, from, n / CHAR_BIT); + } + return p; +} + +void bitmap_copy(u8 *base, u8 *from, i32 n) { + assert(n % 8 == 0); + memcpy(base, from, n / CHAR_BIT); +} + +void bitmap_and_inplace(u8 *base, u8 *other, i32 n) { + assert((n % 8) == 0); + for (int i = 0; i < n / CHAR_BIT; i++) { + base[i] = base[i] & other[i]; + } +} + +void bitmap_set(u8 *bitmap, i32 position, int value) { + if (value) { + bitmap[position / CHAR_BIT] |= 1 << (position % CHAR_BIT); + } else { + bitmap[position / CHAR_BIT] &= ~(1 << (position % CHAR_BIT)); + } +} + +int bitmap_get(u8 *bitmap, i32 position) { + return (((bitmap[position / CHAR_BIT]) >> (position % CHAR_BIT)) & 1); +} + +void bitmap_clear(u8 *bitmap, i32 n) { + assert((n % 8) == 0); + memset(bitmap, 0, n / CHAR_BIT); +} + +void bitmap_fill(u8 *bitmap, i32 n) { + assert((n % 8) == 0); + memset(bitmap, 0xFF, n / CHAR_BIT); +} + +/** + * @brief Finds the minimum k items in distances, and writes the indicies to + * out. + * + * @param distances input f32 array of size n, the items to consider. + * @param n: size of distances array. + * @param out: Output array of size k, will contain at most k element indicies + * @param k: Size of output array + * @return int + */ +int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k, + u8 *bTaken, i32 *k_used) { + assert(k > 0); + assert(k <= n); + + bitmap_clear(bTaken, n); + + for (int ik = 0; ik < k; ik++) { + int min_idx = 0; + while (min_idx < n && + (bitmap_get(bTaken, min_idx) || !bitmap_get(candidates, min_idx))) { + min_idx++; + } + if (min_idx >= n) { + *k_used = ik; + return SQLITE_OK; + } + + for (int i = 0; i < n; i++) { + if (distances[i] <= distances[min_idx] && !bitmap_get(bTaken, i) && + (bitmap_get(candidates, i))) { + min_idx = i; + } + } + + out[ik] = min_idx; + bitmap_set(bTaken, min_idx, 1); + } + *k_used = k; + return SQLITE_OK; +} + +int vec0_get_metadata_text_long_value( + vec0_vtab * p, + sqlite3_stmt ** stmt, + int metadata_idx, + i64 rowid, + int *n, + char ** s) { + int rc; + if(!(*stmt)) { + const char * zSql = sqlite3_mprintf("select data from " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " where rowid = ?", p->schemaName, p->tableName, metadata_idx); + if(!zSql) { + rc = SQLITE_NOMEM; + goto done; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, stmt, NULL); + sqlite3_free( (void *) zSql); + if(rc != SQLITE_OK) { + goto done; + } + } + + sqlite3_reset(*stmt); + sqlite3_bind_int64(*stmt, 1, rowid); + rc = sqlite3_step(*stmt); + if(rc != SQLITE_ROW) { + rc = SQLITE_ERROR; + goto done; + } + *s = (char *) sqlite3_column_text(*stmt, 0); + *n = sqlite3_column_bytes(*stmt, 0); + rc = SQLITE_OK; + done: + return rc; +} + +/** + * @brief Crete at "iterator" (sqlite3_stmt) of chunks with the given constraints + * + * Any VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT values in idxStr/argv will be applied + * as WHERE constraints in the underlying stmt SQL, and any consumer of the stmt + * can freely step through the stmt with all constraints satisfied. + * + * @param p - vec0_vtab + * @param idxStr - the xBestIndex/xFilter idxstr containing VEC0_IDXSTR values + * @param argc - number of argv values from xFilter + * @param argv - array of sqlite3_value from xFilter + * @param outStmt - output sqlite3_stmt of chunks with all filters applied + * @return int SQLITE_OK on success, error code otherwise + */ +int vec0_chunks_iter(vec0_vtab * p, const char * idxStr, int argc, sqlite3_value ** argv, sqlite3_stmt** outStmt) { + // always null terminated, enforced by SQLite + int idxStrLength = strlen(idxStr); + // "1" refers to the initial vec0_query_plan char, 4 is the number of chars per "element" + int numValueEntries = (idxStrLength-1) / 4; + assert(argc == numValueEntries); + + int rc; + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "select chunk_id, validity, rowids " + " from " VEC0_SHADOW_CHUNKS_NAME, + p->schemaName, p->tableName); + + int appendedWhere = 0; + for(int i = 0; i < numValueEntries; i++) { + int idx = 1 + (i * 4); + char kind = idxStr[idx + 0]; + if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) { + continue; + } + + int partition_idx = idxStr[idx + 1] - 'A'; + int operator = idxStr[idx + 2]; + // idxStr[idx + 3] is just null, a '_' placeholder + + if(!appendedWhere) { + sqlite3_str_appendall(s, " WHERE "); + appendedWhere = 1; + }else { + sqlite3_str_appendall(s, " AND "); + } + switch(operator) { + case VEC0_PARTITION_OPERATOR_EQ: + sqlite3_str_appendf(s, " partition%02d = ? ", partition_idx); + break; + case VEC0_PARTITION_OPERATOR_GT: + sqlite3_str_appendf(s, " partition%02d > ? ", partition_idx); + break; + case VEC0_PARTITION_OPERATOR_LE: + sqlite3_str_appendf(s, " partition%02d <= ? ", partition_idx); + break; + case VEC0_PARTITION_OPERATOR_LT: + sqlite3_str_appendf(s, " partition%02d < ? ", partition_idx); + break; + case VEC0_PARTITION_OPERATOR_GE: + sqlite3_str_appendf(s, " partition%02d >= ? ", partition_idx); + break; + case VEC0_PARTITION_OPERATOR_NE: + sqlite3_str_appendf(s, " partition%02d != ? ", partition_idx); + break; + default: { + char * zSql = sqlite3_str_finish(s); + sqlite3_free(zSql); + return SQLITE_ERROR; + } + + } + + } + + char *zSql = sqlite3_str_finish(s); + if (!zSql) { + return SQLITE_NOMEM; + } + + rc = sqlite3_prepare_v2(p->db, zSql, -1, outStmt, NULL); + sqlite3_free(zSql); + if(rc != SQLITE_OK) { + return rc; + } + + int n = 1; + for(int i = 0; i < numValueEntries; i++) { + int idx = 1 + (i * 4); + char kind = idxStr[idx + 0]; + if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) { + continue; + } + sqlite3_bind_value(*outStmt, n++, argv[i]); + } + + return rc; +} + +// a single `xxx in (...)` constraint on a metadata column. TEXT or INTEGER only for now. +struct Vec0MetadataIn{ + // index of argv[i]` the constraint is on + int argv_idx; + // metadata column index of the constraint, derived from idxStr + argv_idx + int metadata_idx; + // array of the copied `(...)` values from sqlite3_vtab_in_first()/sqlite3_vtab_in_next() + struct Array array; +}; + +// Array elements for `xxx in (...)` values for a text column. basically just a string +struct Vec0MetadataInTextEntry { + int n; + char * zString; +}; + + +int vec0_metadata_filter_text(vec0_vtab * p, sqlite3_value * value, const void * buffer, int size, vec0_metadata_operator op, u8* b, int metadata_idx, int chunk_rowid, struct Array * aMetadataIn, int argv_idx) { + int rc; + sqlite3_stmt * stmt = NULL; + i64 * rowids = NULL; + sqlite3_blob * rowidsBlob; + const char * sTarget = (const char *) sqlite3_value_text(value); + int nTarget = sqlite3_value_bytes(value); + + + // TODO(perf): only text metadata news the rowids BLOB. Make it so that + // rowids BLOB is re-used when multiple fitlers on text columns, + // ex "name BETWEEN 'a' and 'b'"" + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids", chunk_rowid, 0, &rowidsBlob); + if(rc != SQLITE_OK) { + return rc; + } + assert(sqlite3_blob_bytes(rowidsBlob) % sizeof(i64) == 0); + assert((sqlite3_blob_bytes(rowidsBlob) / sizeof(i64)) == size); + + rowids = sqlite3_malloc(sqlite3_blob_bytes(rowidsBlob)); + if(!rowids) { + sqlite3_blob_close(rowidsBlob); + return SQLITE_NOMEM; + } + + rc = sqlite3_blob_read(rowidsBlob, rowids, sqlite3_blob_bytes(rowidsBlob), 0); + if(rc != SQLITE_OK) { + sqlite3_blob_close(rowidsBlob); + return rc; + } + sqlite3_blob_close(rowidsBlob); + + switch(op) { + int nPrefix; + char * sPrefix; + char *sFull; + int nFull; + u8 * view; + case VEC0_METADATA_OPERATOR_EQ: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + + // for EQ the text lengths must match + if(nPrefix != nTarget) { + bitmap_set(b, i, 0); + continue; + } + int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)); + + // for short strings, use the prefix comparison direclty + if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + bitmap_set(b, i, cmpPrefix == 0); + continue; + } + // for EQ on longs strings, the prefix must match + if(cmpPrefix) { + bitmap_set(b, i, 0); + continue; + } + // consult the full string + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) == 0); + } + break; + } + case VEC0_METADATA_OPERATOR_NE: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + + // for NE if text lengths dont match, it never will + if(nPrefix != nTarget) { + bitmap_set(b, i, 1); + continue; + } + + int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)); + + // for short strings, use the prefix comparison direclty + if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + bitmap_set(b, i, cmpPrefix != 0); + continue; + } + // for NE on longs strings, if prefixes dont match, then long string wont + if(cmpPrefix) { + bitmap_set(b, i, 1); + continue; + } + // consult the full string + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) != 0); + } + break; + } + case VEC0_METADATA_OPERATOR_GT: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget)); + + if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + // if prefix match, check which is longer + if(cmpPrefix == 0) { + bitmap_set(b, i, nPrefix > nTarget); + } + else { + bitmap_set(b, i, cmpPrefix > 0); + } + continue; + } + // TODO(perf): may not need to compare full text in some cases + + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) > 0); + } + break; + } + case VEC0_METADATA_OPERATOR_GE: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget)); + + if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + // if prefix match, check which is longer + if(cmpPrefix == 0) { + bitmap_set(b, i, nPrefix >= nTarget); + } + else { + bitmap_set(b, i, cmpPrefix >= 0); + } + continue; + } + // TODO(perf): may not need to compare full text in some cases + + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) >= 0); + } + break; + } + case VEC0_METADATA_OPERATOR_LE: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget)); + + if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + // if prefix match, check which is longer + if(cmpPrefix == 0) { + bitmap_set(b, i, nPrefix <= nTarget); + } + else { + bitmap_set(b, i, cmpPrefix <= 0); + } + continue; + } + // TODO(perf): may not need to compare full text in some cases + + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) <= 0); + } + break; + } + case VEC0_METADATA_OPERATOR_LT: { + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget)); + + if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + // if prefix match, check which is longer + if(cmpPrefix == 0) { + bitmap_set(b, i, nPrefix < nTarget); + } + else { + bitmap_set(b, i, cmpPrefix < 0); + } + continue; + } + // TODO(perf): may not need to compare full text in some cases + + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + bitmap_set(b, i, strncmp(sFull, sTarget, nFull) < 0); + } + break; + } + + case VEC0_METADATA_OPERATOR_IN: { + size_t metadataInIdx = -1; + for(size_t i = 0; i < aMetadataIn->length; i++) { + struct Vec0MetadataIn * metadataIn = &(((struct Vec0MetadataIn *) aMetadataIn->z)[i]); + if(metadataIn->argv_idx == argv_idx) { + metadataInIdx = i; + break; + } + } + if(metadataInIdx < 0) { + rc = SQLITE_ERROR; + goto done; + } + + struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx]; + struct Array * aTarget = &(metadataIn->array); + + + int nPrefix; + char * sPrefix; + char *sFull; + int nFull; + u8 * view; + for(int i = 0; i < size; i++) { + view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + nPrefix = ((int*) view)[0]; + sPrefix = (char *) &view[4]; + for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) { + struct Vec0MetadataInTextEntry * entry = &(((struct Vec0MetadataInTextEntry*)aTarget->z)[target_idx]); + if(entry->n != nPrefix) { + continue; + } + int cmpPrefix = strncmp(sPrefix, entry->zString, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)); + if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + if(cmpPrefix == 0) { + bitmap_set(b, i, 1); + break; + } + continue; + } + if(cmpPrefix) { + continue; + } + + rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull); + if(rc != SQLITE_OK) { + goto done; + } + if(nPrefix != nFull) { + rc = SQLITE_ERROR; + goto done; + } + if(strncmp(sFull, entry->zString, nFull) == 0) { + bitmap_set(b, i, 1); + break; + } + } + } + break; + } + + } + rc = SQLITE_OK; + + done: + sqlite3_finalize(stmt); + sqlite3_free(rowids); + return rc; + +} + +/** + * @brief Fill in bitmap of chunk values, whether or not the values match a metadata constraint + * + * @param p vec0_vtab + * @param metadata_idx index of the metatadata column to perfrom constraints on + * @param value sqlite3_value of the constraints value + * @param blob sqlite3_blob that is already opened on the metdata column's shadow chunk table + * @param chunk_rowid rowid of the chunk to calculate on + * @param b pre-allocated and zero'd out bitmap to write results to + * @param size size of the chunk + * @return int SQLITE_OK on success, error code otherwise + */ +int vec0_set_metadata_filter_bitmap( + vec0_vtab *p, + int metadata_idx, + vec0_metadata_operator op, + sqlite3_value * value, + sqlite3_blob * blob, + i64 chunk_rowid, + u8* b, + int size, + struct Array * aMetadataIn, int argv_idx) { + // TODO: shouldn't this skip in-valid entries from the chunk's validity bitmap? + + int rc; + rc = sqlite3_blob_reopen(blob, chunk_rowid); + if(rc != SQLITE_OK) { + return rc; + } + + vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind; + int szMatch = 0; + int blobSize = sqlite3_blob_bytes(blob); + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + szMatch = blobSize == size / CHAR_BIT; + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + szMatch = blobSize == size * sizeof(i64); + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + szMatch = blobSize == size * sizeof(double); + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + szMatch = blobSize == size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH; + break; + } + } + if(!szMatch) { + return SQLITE_ERROR; + } + void * buffer = sqlite3_malloc(blobSize); + if(!buffer) { + return SQLITE_NOMEM; + } + rc = sqlite3_blob_read(blob, buffer, blobSize, 0); + if(rc != SQLITE_OK) { + goto done; + } + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + int target = sqlite3_value_int(value); + if( (target && op == VEC0_METADATA_OPERATOR_EQ) || (!target && op == VEC0_METADATA_OPERATOR_NE)) { + for(int i = 0; i < size; i++) { bitmap_set(b, i, bitmap_get((u8*) buffer, i)); } + } + else { + for(int i = 0; i < size; i++) { bitmap_set(b, i, !bitmap_get((u8*) buffer, i)); } + } + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + i64 * array = (i64*) buffer; + i64 target = sqlite3_value_int64(value); + switch(op) { + case VEC0_METADATA_OPERATOR_EQ: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); } + break; + } + case VEC0_METADATA_OPERATOR_GT: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); } + break; + } + case VEC0_METADATA_OPERATOR_LE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); } + break; + } + case VEC0_METADATA_OPERATOR_LT: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); } + break; + } + case VEC0_METADATA_OPERATOR_GE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); } + break; + } + case VEC0_METADATA_OPERATOR_NE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); } + break; + } + case VEC0_METADATA_OPERATOR_IN: { + int metadataInIdx = -1; + for(size_t i = 0; i < aMetadataIn->length; i++) { + struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[i]; + if(metadataIn->argv_idx == argv_idx) { + metadataInIdx = i; + break; + } + } + if(metadataInIdx < 0) { + rc = SQLITE_ERROR; + goto done; + } + struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx]; + struct Array * aTarget = &(metadataIn->array); + + for(int i = 0; i < size; i++) { + for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) { + if( ((i64*)aTarget->z)[target_idx] == array[i]) { + bitmap_set(b, i, 1); + break; + } + } + } + break; + } + } + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + double * array = (double*) buffer; + double target = sqlite3_value_double(value); + switch(op) { + case VEC0_METADATA_OPERATOR_EQ: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); } + break; + } + case VEC0_METADATA_OPERATOR_GT: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); } + break; + } + case VEC0_METADATA_OPERATOR_LE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); } + break; + } + case VEC0_METADATA_OPERATOR_LT: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); } + break; + } + case VEC0_METADATA_OPERATOR_GE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); } + break; + } + case VEC0_METADATA_OPERATOR_NE: { + for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); } + break; + } + case VEC0_METADATA_OPERATOR_IN: { + // should never be reached + break; + } + } + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + rc = vec0_metadata_filter_text(p, value, buffer, size, op, b, metadata_idx, chunk_rowid, aMetadataIn, argv_idx); + if(rc != SQLITE_OK) { + goto done; + } + break; + } + } + done: + sqlite3_free(buffer); + return rc; +} + +int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks, + struct VectorColumnDefinition *vector_column, + int vectorColumnIdx, struct Array *arrayRowidsIn, + struct Array * aMetadataIn, + const char * idxStr, int argc, sqlite3_value ** argv, + void *queryVector, i64 k, i64 **out_topk_rowids, + f32 **out_topk_distances, i64 *out_used) { + // for each chunk, get top min(k, chunk_size) rowid + distances to query vec. + // then reconcile all topk_chunks for a true top k. + // output only rowids + distances for now + + int rc = SQLITE_OK; + sqlite3_blob *blobVectors = NULL; + + void *baseVectors = NULL; // memory: chunk_size * dimensions * element_size + + // OWNED BY CALLER ON SUCCESS + i64 *topk_rowids = NULL; // memory: k * 4 + // OWNED BY CALLER ON SUCCESS + f32 *topk_distances = NULL; // memory: k * 4 + + i64 *tmp_topk_rowids = NULL; // memory: k * 4 + f32 *tmp_topk_distances = NULL; // memory: k * 4 + f32 *chunk_distances = NULL; // memory: chunk_size * 4 + u8 *b = NULL; // memory: chunk_size / 8 + u8 *bTaken = NULL; // memory: chunk_size / 8 + i32 *chunk_topk_idxs = NULL; // memory: k * 4 + u8 *bmRowids = NULL; // memory: chunk_size / 8 + u8 *bmMetadata = NULL; // memory: chunk_size / 8 + // // total: a lot??? + + // 6 * (k * 4) + (k * 2) + (chunk_size / 8) + (chunk_size * dimensions * 4) + + topk_rowids = sqlite3_malloc(k * sizeof(i64)); + if (!topk_rowids) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(topk_rowids, 0, k * sizeof(i64)); + + topk_distances = sqlite3_malloc(k * sizeof(f32)); + if (!topk_distances) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(topk_distances, 0, k * sizeof(f32)); + + tmp_topk_rowids = sqlite3_malloc(k * sizeof(i64)); + if (!tmp_topk_rowids) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(tmp_topk_rowids, 0, k * sizeof(i64)); + + tmp_topk_distances = sqlite3_malloc(k * sizeof(f32)); + if (!tmp_topk_distances) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(tmp_topk_distances, 0, k * sizeof(f32)); + + i64 k_used = 0; + i64 baseVectorsSize = p->chunk_size * vector_column_byte_size(*vector_column); + baseVectors = sqlite3_malloc(baseVectorsSize); + if (!baseVectors) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + chunk_distances = sqlite3_malloc(p->chunk_size * sizeof(f32)); + if (!chunk_distances) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + b = bitmap_new(p->chunk_size); + if (!b) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + bTaken = bitmap_new(p->chunk_size); + if (!bTaken) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + chunk_topk_idxs = sqlite3_malloc(k * sizeof(i32)); + if (!chunk_topk_idxs) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + bmRowids = arrayRowidsIn ? bitmap_new(p->chunk_size) : NULL; + if (arrayRowidsIn && !bmRowids) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + sqlite3_blob * metadataBlobs[VEC0_MAX_METADATA_COLUMNS]; + memset(metadataBlobs, 0, sizeof(sqlite3_blob*) * VEC0_MAX_METADATA_COLUMNS); + + bmMetadata = bitmap_new(p->chunk_size); + if(!bmMetadata) { + rc = SQLITE_NOMEM; + goto cleanup; + } + + int idxStrLength = strlen(idxStr); + int numValueEntries = (idxStrLength-1) / 4; + assert(numValueEntries == argc); + int hasMetadataFilters = 0; + int hasDistanceConstraints = 0; + for(int i = 0; i < argc; i++) { + int idx = 1 + (i * 4); + char kind = idxStr[idx + 0]; + if(kind == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) { + hasMetadataFilters = 1; + } + else if(kind == VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) { + hasDistanceConstraints = 1; + } + } + + while (true) { + rc = sqlite3_step(stmtChunks); + if (rc == SQLITE_DONE) { + break; + } + if (rc != SQLITE_ROW) { + vtab_set_error(&p->base, "chunks iter error"); + rc = SQLITE_ERROR; + goto cleanup; + } + memset(chunk_distances, 0, p->chunk_size * sizeof(f32)); + memset(chunk_topk_idxs, 0, k * sizeof(i32)); + bitmap_clear(b, p->chunk_size); + + i64 chunk_id = sqlite3_column_int64(stmtChunks, 0); + unsigned char *chunkValidity = + (unsigned char *)sqlite3_column_blob(stmtChunks, 1); + i64 validitySize = sqlite3_column_bytes(stmtChunks, 1); + if (validitySize != p->chunk_size / CHAR_BIT) { + // IMP: V05271_22109 + vtab_set_error( + &p->base, + "chunk validity size doesn't match - expected %lld, found %lld", + p->chunk_size / CHAR_BIT, validitySize); + rc = SQLITE_ERROR; + goto cleanup; + } + + i64 *chunkRowids = (i64 *)sqlite3_column_blob(stmtChunks, 2); + i64 rowidsSize = sqlite3_column_bytes(stmtChunks, 2); + if (rowidsSize != p->chunk_size * sizeof(i64)) { + // IMP: V02796_19635 + vtab_set_error(&p->base, "rowids size doesn't match"); + vtab_set_error( + &p->base, + "chunk rowids size doesn't match - expected %lld, found %lld", + p->chunk_size * sizeof(i64), rowidsSize); + rc = SQLITE_ERROR; + goto cleanup; + } + + // open the vector chunk blob for the current chunk + rc = sqlite3_blob_open(p->db, p->schemaName, + p->shadowVectorChunksNames[vectorColumnIdx], + "vectors", chunk_id, 0, &blobVectors); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "could not open vectors blob for chunk %lld", + chunk_id); + rc = SQLITE_ERROR; + goto cleanup; + } + + i64 currentBaseVectorsSize = sqlite3_blob_bytes(blobVectors); + i64 expectedBaseVectorsSize = + p->chunk_size * vector_column_byte_size(*vector_column); + if (currentBaseVectorsSize != expectedBaseVectorsSize) { + // IMP: V16465_00535 + vtab_set_error( + &p->base, + "vectors blob size doesn't match - expected %lld, found %lld", + expectedBaseVectorsSize, currentBaseVectorsSize); + rc = SQLITE_ERROR; + goto cleanup; + } + rc = sqlite3_blob_read(blobVectors, baseVectors, currentBaseVectorsSize, 0); + + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "vectors blob read error for %lld", chunk_id); + rc = SQLITE_ERROR; + goto cleanup; + } + + bitmap_copy(b, chunkValidity, p->chunk_size); + if (arrayRowidsIn) { + bitmap_clear(bmRowids, p->chunk_size); + + for (int i = 0; i < p->chunk_size; i++) { + if (!bitmap_get(chunkValidity, i)) { + continue; + } + i64 rowid = chunkRowids[i]; + void *in = bsearch(&rowid, arrayRowidsIn->z, arrayRowidsIn->length, + sizeof(i64), _cmp); + bitmap_set(bmRowids, i, in ? 1 : 0); + } + bitmap_and_inplace(b, bmRowids, p->chunk_size); + } + + if(hasMetadataFilters) { + for(int i = 0; i < argc; i++) { + int idx = 1 + (i * 4); + char kind = idxStr[idx + 0]; + if(kind != VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) { + continue; + } + int metadata_idx = idxStr[idx + 1] - 'A'; + int operator = idxStr[idx + 2]; + + if(!metadataBlobs[metadata_idx]) { + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &metadataBlobs[metadata_idx]); + vtab_set_error(&p->base, "Could not open metadata blob"); + if(rc != SQLITE_OK) { + goto cleanup; + } + } + + bitmap_clear(bmMetadata, p->chunk_size); + rc = vec0_set_metadata_filter_bitmap(p, metadata_idx, operator, argv[i], metadataBlobs[metadata_idx], chunk_id, bmMetadata, p->chunk_size, aMetadataIn, i); + if(rc != SQLITE_OK) { + vtab_set_error(&p->base, "Could not filter metadata fields"); + if(rc != SQLITE_OK) { + goto cleanup; + } + } + bitmap_and_inplace(b, bmMetadata, p->chunk_size); + } + } + + + for (int i = 0; i < p->chunk_size; i++) { + if (!bitmap_get(b, i)) { + continue; + }; + + f32 result; + switch (vector_column->element_type) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: { + const f32 *base_i = + ((f32 *)baseVectors) + (i * vector_column->dimensions); + switch (vector_column->distance_metric) { + case VEC0_DISTANCE_METRIC_L2: { + result = distance_l2_sqr_float(base_i, (f32 *)queryVector, + &vector_column->dimensions); + break; + } + case VEC0_DISTANCE_METRIC_L1: { + result = distance_l1_f32(base_i, (f32 *)queryVector, + &vector_column->dimensions); + break; + } + case VEC0_DISTANCE_METRIC_COSINE: { + result = distance_cosine_float(base_i, (f32 *)queryVector, + &vector_column->dimensions); + break; + } + } + break; + } + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + const i8 *base_i = + ((i8 *)baseVectors) + (i * vector_column->dimensions); + switch (vector_column->distance_metric) { + case VEC0_DISTANCE_METRIC_L2: { + result = distance_l2_sqr_int8(base_i, (i8 *)queryVector, + &vector_column->dimensions); + break; + } + case VEC0_DISTANCE_METRIC_L1: { + result = distance_l1_int8(base_i, (i8 *)queryVector, + &vector_column->dimensions); + break; + } + case VEC0_DISTANCE_METRIC_COSINE: { + result = distance_cosine_int8(base_i, (i8 *)queryVector, + &vector_column->dimensions); + break; + } + } + + break; + } + case SQLITE_VEC_ELEMENT_TYPE_BIT: { + const u8 *base_i = + ((u8 *)baseVectors) + (i * (vector_column->dimensions / CHAR_BIT)); + result = distance_hamming(base_i, (u8 *)queryVector, + &vector_column->dimensions); + break; + } + } + + chunk_distances[i] = result; + } + + if(hasDistanceConstraints) { + for(int i = 0; i < argc; i++) { + int idx = 1 + (i * 4); + char kind = idxStr[idx + 0]; + // TODO casts f64 to f32, is that a problem? + f32 target = (f32) sqlite3_value_double(argv[i]); + + if(kind != VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) { + continue; + } + vec0_distance_constraint_operator op = idxStr[idx + 1]; + + switch(op) { + case VEC0_DISTANCE_CONSTRAINT_GE: { + for(int i = 0; i < p->chunk_size;i++) { + if(bitmap_get(b, i) && !(chunk_distances[i] >= target)) { + bitmap_set(b, i, 0); + } + } + break; + } + case VEC0_DISTANCE_CONSTRAINT_GT: { + for(int i = 0; i < p->chunk_size;i++) { + if(bitmap_get(b, i) && !(chunk_distances[i] > target)) { + bitmap_set(b, i, 0); + } + } + break; + } + case VEC0_DISTANCE_CONSTRAINT_LE: { + for(int i = 0; i < p->chunk_size;i++) { + if(bitmap_get(b, i) && !(chunk_distances[i] <= target)) { + bitmap_set(b, i, 0); + } + } + break; + } + case VEC0_DISTANCE_CONSTRAINT_LT: { + for(int i = 0; i < p->chunk_size;i++) { + if(bitmap_get(b, i) && !(chunk_distances[i] < target)) { + bitmap_set(b, i, 0); + } + } + break; + } + } + } + } + + int used1; + min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs, + min(k, p->chunk_size), bTaken, &used1); + + i64 used; + merge_sorted_lists(topk_distances, topk_rowids, k_used, chunk_distances, + chunkRowids, chunk_topk_idxs, + min(min(k, p->chunk_size), used1), tmp_topk_distances, + tmp_topk_rowids, k, &used); + + for (int i = 0; i < used; i++) { + topk_rowids[i] = tmp_topk_rowids[i]; + topk_distances[i] = tmp_topk_distances[i]; + } + k_used = used; + // blobVectors is always opened with read-only permissions, so this never + // fails. + sqlite3_blob_close(blobVectors); + blobVectors = NULL; + } + + *out_topk_rowids = topk_rowids; + *out_topk_distances = topk_distances; + *out_used = k_used; + rc = SQLITE_OK; + +cleanup: + if (rc != SQLITE_OK) { + sqlite3_free(topk_rowids); + sqlite3_free(topk_distances); + } + sqlite3_free(chunk_topk_idxs); + sqlite3_free(tmp_topk_rowids); + sqlite3_free(tmp_topk_distances); + sqlite3_free(b); + sqlite3_free(bTaken); + sqlite3_free(bmRowids); + sqlite3_free(baseVectors); + sqlite3_free(chunk_distances); + sqlite3_free(bmMetadata); + for(int i = 0; i < VEC0_MAX_METADATA_COLUMNS; i++) { + sqlite3_blob_close(metadataBlobs[i]); + } + // blobVectors is always opened with read-only permissions, so this never + // fails. + sqlite3_blob_close(blobVectors); + return rc; +} + +int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum, + const char *idxStr, int argc, sqlite3_value **argv) { + assert(argc == (strlen(idxStr)-1) / 4); + int rc; + struct vec0_query_knn_data *knn_data; + + int vectorColumnIdx = idxNum; + struct VectorColumnDefinition *vector_column = + &p->vector_columns[vectorColumnIdx]; + + struct Array *arrayRowidsIn = NULL; + sqlite3_stmt *stmtChunks = NULL; + void *queryVector; + size_t dimensions; + enum VectorElementType elementType; + vector_cleanup queryVectorCleanup = vector_cleanup_noop; + char *pzError; + knn_data = sqlite3_malloc(sizeof(*knn_data)); + if (!knn_data) { + return SQLITE_NOMEM; + } + memset(knn_data, 0, sizeof(*knn_data)); + // array of `struct Vec0MetadataIn`, IF there are any `xxx in (...)` metadata constraints + struct Array * aMetadataIn = NULL; + + int query_idx =-1; + int k_idx = -1; + int rowid_in_idx = -1; + for(int i = 0; i < argc; i++) { + if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_MATCH) { + query_idx = i; + } + if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_K) { + k_idx = i; + } + if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_ROWID_IN) { + rowid_in_idx = i; + } + } + assert(query_idx >= 0); + assert(k_idx >= 0); + + // make sure the query vector matches the vector column (type dimensions etc.) + rc = vector_from_value(argv[query_idx], &queryVector, &dimensions, &elementType, + &queryVectorCleanup, &pzError); + + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + "Query vector on the \"%.*s\" column is invalid: %z", + vector_column->name_length, vector_column->name, pzError); + rc = SQLITE_ERROR; + goto cleanup; + } + if (elementType != vector_column->element_type) { + vtab_set_error( + &p->base, + "Query vector for the \"%.*s\" column is expected to be of type " + "%s, but a %s vector was provided.", + vector_column->name_length, vector_column->name, + vector_subtype_name(vector_column->element_type), + vector_subtype_name(elementType)); + rc = SQLITE_ERROR; + goto cleanup; + } + if (dimensions != vector_column->dimensions) { + vtab_set_error( + &p->base, + "Dimension mismatch for query vector for the \"%.*s\" column. " + "Expected %d dimensions but received %d.", + vector_column->name_length, vector_column->name, + vector_column->dimensions, dimensions); + rc = SQLITE_ERROR; + goto cleanup; + } + + i64 k = sqlite3_value_int64(argv[k_idx]); + if (k < 0) { + vtab_set_error( + &p->base, "k value in knn queries must be greater than or equal to 0."); + rc = SQLITE_ERROR; + goto cleanup; + } +#define SQLITE_VEC_VEC0_K_MAX 4096 + if (k > SQLITE_VEC_VEC0_K_MAX) { + vtab_set_error( + &p->base, + "k value in knn query too large, provided %lld and the limit is %lld", + k, SQLITE_VEC_VEC0_K_MAX); + rc = SQLITE_ERROR; + goto cleanup; + } + + if (k == 0) { + knn_data->k = 0; + pCur->knn_data = knn_data; + pCur->query_plan = VEC0_QUERY_PLAN_KNN; + rc = SQLITE_OK; + goto cleanup; + } + +// handle when a `rowid in (...)` operation was provided +// Array of all the rowids that appear in any `rowid in (...)` constraint. +// NULL if none were provided, which means a "full" scan. +#if COMPILER_SUPPORTS_VTAB_IN + if (rowid_in_idx >= 0) { + sqlite3_value *item; + int rc; + arrayRowidsIn = sqlite3_malloc(sizeof(*arrayRowidsIn)); + if (!arrayRowidsIn) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(arrayRowidsIn, 0, sizeof(*arrayRowidsIn)); + + rc = array_init(arrayRowidsIn, sizeof(i64), 32); + if (rc != SQLITE_OK) { + goto cleanup; + } + for (rc = sqlite3_vtab_in_first(argv[rowid_in_idx], &item); rc == SQLITE_OK && item; + rc = sqlite3_vtab_in_next(argv[rowid_in_idx], &item)) { + i64 rowid; + if (p->pkIsText) { + rc = vec0_rowid_from_id(p, item, &rowid); + if (rc != SQLITE_OK) { + goto cleanup; + } + } else { + rowid = sqlite3_value_int64(item); + } + rc = array_append(arrayRowidsIn, &rowid); + if (rc != SQLITE_OK) { + goto cleanup; + } + } + if (rc != SQLITE_DONE) { + vtab_set_error(&p->base, "error processing rowid in (...) array"); + goto cleanup; + } + qsort(arrayRowidsIn->z, arrayRowidsIn->length, arrayRowidsIn->element_size, + _cmp); + } +#endif + + #if COMPILER_SUPPORTS_VTAB_IN + for(int i = 0; i < argc; i++) { + if(!(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT && idxStr[1 + (i*4) + 2] == VEC0_METADATA_OPERATOR_IN)) { + continue; + } + int metadata_idx = idxStr[1 + (i*4) + 1] - 'A'; + if(!aMetadataIn) { + aMetadataIn = sqlite3_malloc(sizeof(*aMetadataIn)); + if(!aMetadataIn) { + rc = SQLITE_NOMEM; + goto cleanup; + } + memset(aMetadataIn, 0, sizeof(*aMetadataIn)); + rc = array_init(aMetadataIn, sizeof(struct Vec0MetadataIn), 8); + if(rc != SQLITE_OK) { + goto cleanup; + } + } + + struct Vec0MetadataIn item; + memset(&item, 0, sizeof(item)); + item.metadata_idx=metadata_idx; + item.argv_idx = i; + + switch(p->metadata_columns[metadata_idx].kind) { + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + rc = array_init(&item.array, sizeof(i64), 16); + if(rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_value *entry; + for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) { + i64 v = sqlite3_value_int64(entry); + rc = array_append(&item.array, &v); + if (rc != SQLITE_OK) { + goto cleanup; + } + } + + if (rc != SQLITE_DONE) { + vtab_set_error(&p->base, "Error fetching next value in `x in (...)` integer expression"); + goto cleanup; + } + + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + rc = array_init(&item.array, sizeof(struct Vec0MetadataInTextEntry), 16); + if(rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_value *entry; + for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) { + const char * s = (const char *) sqlite3_value_text(entry); + int n = sqlite3_value_bytes(entry); + + struct Vec0MetadataInTextEntry entry; + entry.zString = sqlite3_mprintf("%.*s", n, s); + if(!entry.zString) { + rc = SQLITE_NOMEM; + goto cleanup; + } + entry.n = n; + rc = array_append(&item.array, &entry); + if (rc != SQLITE_OK) { + goto cleanup; + } + } + + if (rc != SQLITE_DONE) { + vtab_set_error(&p->base, "Error fetching next value in `x in (...)` text expression"); + goto cleanup; + } + + break; + } + default: { + vtab_set_error(&p->base, "Internal sqlite-vec error"); + goto cleanup; + } + } + + rc = array_append(aMetadataIn, &item); + if(rc != SQLITE_OK) { + goto cleanup; + } + } + #endif + + rc = vec0_chunks_iter(p, idxStr, argc, argv, &stmtChunks); + if (rc != SQLITE_OK) { + // IMP: V06942_23781 + vtab_set_error(&p->base, "Error preparing stmtChunk: %s", + sqlite3_errmsg(p->db)); + goto cleanup; + } + + i64 *topk_rowids = NULL; + f32 *topk_distances = NULL; + i64 k_used = 0; + rc = vec0Filter_knn_chunks_iter(p, stmtChunks, vector_column, vectorColumnIdx, + arrayRowidsIn, aMetadataIn, idxStr, argc, argv, queryVector, k, &topk_rowids, + &topk_distances, &k_used); + if (rc != SQLITE_OK) { + goto cleanup; + } + + knn_data->current_idx = 0; + knn_data->k = k; + knn_data->rowids = topk_rowids; + knn_data->distances = topk_distances; + knn_data->k_used = k_used; + + pCur->knn_data = knn_data; + pCur->query_plan = VEC0_QUERY_PLAN_KNN; + rc = SQLITE_OK; + +cleanup: + sqlite3_finalize(stmtChunks); + array_cleanup(arrayRowidsIn); + sqlite3_free(arrayRowidsIn); + queryVectorCleanup(queryVector); + if(aMetadataIn) { + for(size_t i = 0; i < aMetadataIn->length; i++) { + struct Vec0MetadataIn* item = &((struct Vec0MetadataIn *) aMetadataIn->z)[i]; + for(size_t j = 0; j < item->array.length; j++) { + if(p->metadata_columns[item->metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_TEXT) { + struct Vec0MetadataInTextEntry entry = ((struct Vec0MetadataInTextEntry*)item->array.z)[j]; + sqlite3_free(entry.zString); + } + } + array_cleanup(&item->array); + } + array_cleanup(aMetadataIn); + } + + sqlite3_free(aMetadataIn); + + if (rc != SQLITE_OK) { + sqlite3_free(knn_data); + } + + return rc; +} + +int vec0Filter_fullscan(vec0_vtab *p, vec0_cursor *pCur) { + int rc; + char *zSql; + struct vec0_query_fullscan_data *fullscan_data; + + fullscan_data = sqlite3_malloc(sizeof(*fullscan_data)); + if (!fullscan_data) { + return SQLITE_NOMEM; + } + memset(fullscan_data, 0, sizeof(*fullscan_data)); + + zSql = sqlite3_mprintf(" SELECT rowid " + " FROM " VEC0_SHADOW_ROWIDS_NAME + " ORDER by chunk_id, chunk_offset ", + p->schemaName, p->tableName); + if (!zSql) { + rc = SQLITE_NOMEM; + goto error; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &fullscan_data->rowids_stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) { + // IMP: V09901_26739 + vtab_set_error(&p->base, "Error preparing rowid scan: %s", + sqlite3_errmsg(p->db)); + goto error; + } + + rc = sqlite3_step(fullscan_data->rowids_stmt); + + // DONE when there's no rowids, ROW when there are, both "success" + if (!(rc == SQLITE_ROW || rc == SQLITE_DONE)) { + goto error; + } + + fullscan_data->done = rc == SQLITE_DONE; + pCur->query_plan = VEC0_QUERY_PLAN_FULLSCAN; + pCur->fullscan_data = fullscan_data; + return SQLITE_OK; + +error: + vec0_query_fullscan_data_clear(fullscan_data); + sqlite3_free(fullscan_data); + return rc; +} + +int vec0Filter_point(vec0_cursor *pCur, vec0_vtab *p, int argc, + sqlite3_value **argv) { + int rc; + assert(argc == 1); + i64 rowid; + struct vec0_query_point_data *point_data = NULL; + + point_data = sqlite3_malloc(sizeof(*point_data)); + if (!point_data) { + rc = SQLITE_NOMEM; + goto error; + } + memset(point_data, 0, sizeof(*point_data)); + + if (p->pkIsText) { + rc = vec0_rowid_from_id(p, argv[0], &rowid); + if (rc == SQLITE_EMPTY) { + goto eof; + } + if (rc != SQLITE_OK) { + goto error; + } + } else { + rowid = sqlite3_value_int64(argv[0]); + } + + for (int i = 0; i < p->numVectorColumns; i++) { + rc = vec0_get_vector_data(p, rowid, i, &point_data->vectors[i], NULL); + if (rc == SQLITE_EMPTY) { + goto eof; + } + if (rc != SQLITE_OK) { + goto error; + } + } + + point_data->rowid = rowid; + point_data->done = 0; + pCur->point_data = point_data; + pCur->query_plan = VEC0_QUERY_PLAN_POINT; + return SQLITE_OK; + +eof: + point_data->rowid = rowid; + point_data->done = 1; + pCur->point_data = point_data; + pCur->query_plan = VEC0_QUERY_PLAN_POINT; + return SQLITE_OK; + +error: + vec0_query_point_data_clear(point_data); + sqlite3_free(point_data); + return rc; +} + +static int vec0Filter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, + const char *idxStr, int argc, sqlite3_value **argv) { + vec0_vtab *p = (vec0_vtab *)pVtabCursor->pVtab; + vec0_cursor *pCur = (vec0_cursor *)pVtabCursor; + vec0_cursor_clear(pCur); + + int idxStrLength = strlen(idxStr); + if(idxStrLength <= 0) { + return SQLITE_ERROR; + } + if((idxStrLength-1) % 4 != 0) { + return SQLITE_ERROR; + } + int numValueEntries = (idxStrLength-1) / 4; + if(numValueEntries != argc) { + return SQLITE_ERROR; + } + + char query_plan = idxStr[0]; + switch(query_plan) { + case VEC0_QUERY_PLAN_FULLSCAN: + return vec0Filter_fullscan(p, pCur); + case VEC0_QUERY_PLAN_KNN: + return vec0Filter_knn(pCur, p, idxNum, idxStr, argc, argv); + case VEC0_QUERY_PLAN_POINT: + return vec0Filter_point(pCur, p, argc, argv); + default: + vtab_set_error(pVtabCursor->pVtab, "unknown idxStr '%s'", idxStr); + return SQLITE_ERROR; + } +} + +static int vec0Rowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { + vec0_cursor *pCur = (vec0_cursor *)cur; + switch (pCur->query_plan) { + case VEC0_QUERY_PLAN_FULLSCAN: { + *pRowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0); + return SQLITE_OK; + } + case VEC0_QUERY_PLAN_POINT: { + *pRowid = pCur->point_data->rowid; + return SQLITE_OK; + } + case VEC0_QUERY_PLAN_KNN: { + vtab_set_error(cur->pVtab, + "Internal sqlite-vec error: expected point query plan in " + "vec0Rowid, found %d", + pCur->query_plan); + return SQLITE_ERROR; + } + } + return SQLITE_ERROR; +} + +static int vec0Next(sqlite3_vtab_cursor *cur) { + vec0_cursor *pCur = (vec0_cursor *)cur; + switch (pCur->query_plan) { + case VEC0_QUERY_PLAN_FULLSCAN: { + if (!pCur->fullscan_data) { + return SQLITE_ERROR; + } + int rc = sqlite3_step(pCur->fullscan_data->rowids_stmt); + if (rc == SQLITE_DONE) { + pCur->fullscan_data->done = 1; + return SQLITE_OK; + } + if (rc == SQLITE_ROW) { + return SQLITE_OK; + } + return SQLITE_ERROR; + } + case VEC0_QUERY_PLAN_KNN: { + if (!pCur->knn_data) { + return SQLITE_ERROR; + } + + pCur->knn_data->current_idx++; + return SQLITE_OK; + } + case VEC0_QUERY_PLAN_POINT: { + if (!pCur->point_data) { + return SQLITE_ERROR; + } + pCur->point_data->done = 1; + return SQLITE_OK; + } + } + return SQLITE_ERROR; +} + +static int vec0Eof(sqlite3_vtab_cursor *cur) { + vec0_cursor *pCur = (vec0_cursor *)cur; + switch (pCur->query_plan) { + case VEC0_QUERY_PLAN_FULLSCAN: { + if (!pCur->fullscan_data) { + return 1; + } + return pCur->fullscan_data->done; + } + case VEC0_QUERY_PLAN_KNN: { + if (!pCur->knn_data) { + return 1; + } + // return (pCur->knn_data->current_idx >= pCur->knn_data->k) || + // (pCur->knn_data->distances[pCur->knn_data->current_idx] == FLT_MAX); + return (pCur->knn_data->current_idx >= pCur->knn_data->k_used); + } + case VEC0_QUERY_PLAN_POINT: { + if (!pCur->point_data) { + return 1; + } + return pCur->point_data->done; + } + } + return 1; +} + +static int vec0Column_fullscan(vec0_vtab *pVtab, vec0_cursor *pCur, + sqlite3_context *context, int i) { + if (!pCur->fullscan_data) { + sqlite3_result_error( + context, "Internal sqlite-vec error: fullscan_data is NULL.", -1); + return SQLITE_ERROR; + } + i64 rowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0); + if (i == VEC0_COLUMN_ID) { + return vec0_result_id(pVtab, context, rowid); + } + else if (vec0_column_idx_is_vector(pVtab, i)) { + void *v; + int sz; + int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i); + int rc = vec0_get_vector_data(pVtab, rowid, vector_idx, &v, &sz); + if (rc != SQLITE_OK) { + return rc; + } + sqlite3_result_blob(context, v, sz, sqlite3_free); + sqlite3_result_subtype(context, + pVtab->vector_columns[vector_idx].element_type); + + } + else if (i == vec0_column_distance_idx(pVtab)) { + sqlite3_result_null(context); + } + else if(vec0_column_idx_is_partition(pVtab, i)) { + int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i); + sqlite3_value * v; + int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + else if(vec0_column_idx_is_auxiliary(pVtab, i)) { + int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i); + sqlite3_value * v; + int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + + else if(vec0_column_idx_is_metadata(pVtab, i)) { + if(sqlite3_vtab_nochange(context)) { + return SQLITE_OK; + } + int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i); + int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context); + if(rc != SQLITE_OK) { + // IMP: V15466_32305 + const char * zErr = sqlite3_mprintf( + "Could not extract metadata value for column %.*s at rowid %lld", + pVtab->metadata_columns[metadata_idx].name_length, + pVtab->metadata_columns[metadata_idx].name, rowid + ); + if(zErr) { + sqlite3_result_error(context, zErr, -1); + sqlite3_free((void *) zErr); + }else { + sqlite3_result_error_nomem(context); + } + } + } + + return SQLITE_OK; +} + +static int vec0Column_point(vec0_vtab *pVtab, vec0_cursor *pCur, + sqlite3_context *context, int i) { + if (!pCur->point_data) { + sqlite3_result_error(context, + "Internal sqlite-vec error: point_data is NULL.", -1); + return SQLITE_ERROR; + } + if (i == VEC0_COLUMN_ID) { + return vec0_result_id(pVtab, context, pCur->point_data->rowid); + } + else if (i == vec0_column_distance_idx(pVtab)) { + sqlite3_result_null(context); + return SQLITE_OK; + } + else if (vec0_column_idx_is_vector(pVtab, i)) { + if (sqlite3_vtab_nochange(context)) { + sqlite3_result_null(context); + return SQLITE_OK; + } + int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i); + sqlite3_result_blob( + context, pCur->point_data->vectors[vector_idx], + vector_column_byte_size(pVtab->vector_columns[vector_idx]), + SQLITE_TRANSIENT); + sqlite3_result_subtype(context, + pVtab->vector_columns[vector_idx].element_type); + return SQLITE_OK; + } + else if(vec0_column_idx_is_partition(pVtab, i)) { + if(sqlite3_vtab_nochange(context)) { + return SQLITE_OK; + } + int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i); + i64 rowid = pCur->point_data->rowid; + sqlite3_value * v; + int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + else if(vec0_column_idx_is_auxiliary(pVtab, i)) { + if(sqlite3_vtab_nochange(context)) { + return SQLITE_OK; + } + i64 rowid = pCur->point_data->rowid; + int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i); + sqlite3_value * v; + int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + + else if(vec0_column_idx_is_metadata(pVtab, i)) { + if(sqlite3_vtab_nochange(context)) { + return SQLITE_OK; + } + i64 rowid = pCur->point_data->rowid; + int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i); + int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context); + if(rc != SQLITE_OK) { + const char * zErr = sqlite3_mprintf( + "Could not extract metadata value for column %.*s at rowid %lld", + pVtab->metadata_columns[metadata_idx].name_length, + pVtab->metadata_columns[metadata_idx].name, rowid + ); + if(zErr) { + sqlite3_result_error(context, zErr, -1); + sqlite3_free((void *) zErr); + }else { + sqlite3_result_error_nomem(context); + } + } + } + + return SQLITE_OK; +} + +static int vec0Column_knn(vec0_vtab *pVtab, vec0_cursor *pCur, + sqlite3_context *context, int i) { + if (!pCur->knn_data) { + sqlite3_result_error(context, + "Internal sqlite-vec error: knn_data is NULL.", -1); + return SQLITE_ERROR; + } + if (i == VEC0_COLUMN_ID) { + i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx]; + return vec0_result_id(pVtab, context, rowid); + } + else if (i == vec0_column_distance_idx(pVtab)) { + sqlite3_result_double( + context, pCur->knn_data->distances[pCur->knn_data->current_idx]); + return SQLITE_OK; + } + else if (vec0_column_idx_is_vector(pVtab, i)) { + void *out; + int sz; + int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i); + int rc = vec0_get_vector_data( + pVtab, pCur->knn_data->rowids[pCur->knn_data->current_idx], vector_idx, + &out, &sz); + if (rc != SQLITE_OK) { + return rc; + } + sqlite3_result_blob(context, out, sz, sqlite3_free); + sqlite3_result_subtype(context, + pVtab->vector_columns[vector_idx].element_type); + return SQLITE_OK; + } + else if(vec0_column_idx_is_partition(pVtab, i)) { + int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i); + i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx]; + sqlite3_value * v; + int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + else if(vec0_column_idx_is_auxiliary(pVtab, i)) { + int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i); + i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx]; + sqlite3_value * v; + int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v); + if(rc == SQLITE_OK) { + sqlite3_result_value(context, v); + sqlite3_value_free(v); + }else { + sqlite3_result_error_code(context, rc); + } + } + + else if(vec0_column_idx_is_metadata(pVtab, i)) { + int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i); + i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx]; + int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context); + if(rc != SQLITE_OK) { + const char * zErr = sqlite3_mprintf( + "Could not extract metadata value for column %.*s at rowid %lld", + pVtab->metadata_columns[metadata_idx].name_length, + pVtab->metadata_columns[metadata_idx].name, rowid + ); + if(zErr) { + sqlite3_result_error(context, zErr, -1); + sqlite3_free((void *) zErr); + }else { + sqlite3_result_error_nomem(context); + } + } + } + + return SQLITE_OK; +} + +static int vec0Column(sqlite3_vtab_cursor *cur, sqlite3_context *context, + int i) { + vec0_cursor *pCur = (vec0_cursor *)cur; + vec0_vtab *pVtab = (vec0_vtab *)cur->pVtab; + switch (pCur->query_plan) { + case VEC0_QUERY_PLAN_FULLSCAN: { + return vec0Column_fullscan(pVtab, pCur, context, i); + } + case VEC0_QUERY_PLAN_KNN: { + return vec0Column_knn(pVtab, pCur, context, i); + } + case VEC0_QUERY_PLAN_POINT: { + return vec0Column_point(pVtab, pCur, context, i); + } + } + return SQLITE_OK; +} + +/** + * @brief Handles the "insert rowid" step of a row insert operation of a vec0 + * table. + * + * This function will insert a new row into the _rowids vec0 shadow table. + * + * @param p: virtual table + * @param idValue: Value containing the inserted rowid/id value. + * @param rowid: Output rowid, will point to the "real" i64 rowid + * value that was inserted + * @return int SQLITE_OK on success, error code on failure + */ +int vec0Update_InsertRowidStep(vec0_vtab *p, sqlite3_value *idValue, + i64 *rowid) { + + /** + * An insert into a vec0 table can happen a few different ways: + * 1) With default INTEGER primary key: With a supplied i64 rowid + * 2) With default INTEGER primary key: WITHOUT a supplied rowid + * 3) With TEXT primary key: supplied text rowid + */ + + int rc; + + // Option 3: vtab has a user-defined TEXT primary key, so ensure a text value + // is provided. + if (p->pkIsText) { + if (sqlite3_value_type(idValue) != SQLITE_TEXT) { + // IMP: V04200_21039 + vtab_set_error(&p->base, + "The %s virtual table was declared with a TEXT primary " + "key, but a non-TEXT value was provided in an INSERT.", + p->tableName); + return SQLITE_ERROR; + } + + return vec0_rowids_insert_id(p, idValue, rowid); + } + + // Option 1: User supplied a i64 rowid + if (sqlite3_value_type(idValue) == SQLITE_INTEGER) { + i64 suppliedRowid = sqlite3_value_int64(idValue); + rc = vec0_rowids_insert_rowid(p, suppliedRowid); + if (rc == SQLITE_OK) { + *rowid = suppliedRowid; + } + return rc; + } + + // Option 2: User did not suppled a rowid + + if (sqlite3_value_type(idValue) != SQLITE_NULL) { + // IMP: V30855_14925 + vtab_set_error(&p->base, + "Only integers are allows for primary key values on %s", + p->tableName); + return SQLITE_ERROR; + } + // NULL to get next auto-incremented value + return vec0_rowids_insert_id(p, NULL, rowid); +} + +/** + * @brief Determines the "next available" chunk position for a newly inserted + * vec0 row. + * + * This operation may insert a new "blank" chunk the _chunks table, if there is + * no more space in previous chunks. + * + * @param p: virtual table + * @param partitionKeyValues: array of partition key column values, to constrain + * against any partition key columns. + * @param chunk_rowid: Output rowid of the chunk in the _chunks virtual table + * that has the avialabiity. + * @param chunk_offset: Output the index of the available space insert the + * chunk, based on the index of the first available validity bit. + * @param pBlobValidity: Output blob of the validity column of the available + * chunk. Will be opened with read/write permissions. + * @param pValidity: Output buffer of the original chunk's validity column. + * Needs to be cleaned up with sqlite3_free(). + * @return int SQLITE_OK on success, error code on failure + */ +int vec0Update_InsertNextAvailableStep( + vec0_vtab *p, + sqlite3_value ** partitionKeyValues, + i64 *chunk_rowid, i64 *chunk_offset, + sqlite3_blob **blobChunksValidity, + const unsigned char **bufferChunksValidity) { + + int rc; + i64 validitySize; + *chunk_offset = -1; + + rc = vec0_get_latest_chunk_rowid(p, chunk_rowid, partitionKeyValues); + if(rc == SQLITE_EMPTY) { + goto done; + } + if (rc != SQLITE_OK) { + goto cleanup; + } + + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity", + *chunk_rowid, 1, blobChunksValidity); + if (rc != SQLITE_OK) { + // IMP: V22053_06123 + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "could not open validity blob on %s.%s.%lld", + p->schemaName, p->shadowChunksName, *chunk_rowid); + goto cleanup; + } + + validitySize = sqlite3_blob_bytes(*blobChunksValidity); + if (validitySize != p->chunk_size / CHAR_BIT) { + // IMP: V29362_13432 + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "validity blob size mismatch on " + "%s.%s.%lld, expected %lld but received %lld.", + p->schemaName, p->shadowChunksName, *chunk_rowid, + (i64)(p->chunk_size / CHAR_BIT), validitySize); + rc = SQLITE_ERROR; + goto cleanup; + } + + *bufferChunksValidity = sqlite3_malloc(validitySize); + if (!(*bufferChunksValidity)) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "Could not allocate memory for validity bitmap"); + rc = SQLITE_NOMEM; + goto cleanup; + } + + rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity, + validitySize, 0); + + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "Could not read validity bitmap for %s.%s.%lld", + p->schemaName, p->shadowChunksName, *chunk_rowid); + goto cleanup; + } + + // find the next available offset, ie first `0` in the bitmap. + for (int i = 0; i < validitySize; i++) { + if ((*bufferChunksValidity)[i] == 0b11111111) + continue; + for (int j = 0; j < CHAR_BIT; j++) { + if (((((*bufferChunksValidity)[i] >> j) & 1) == 0)) { + *chunk_offset = (i * CHAR_BIT) + j; + goto done; + } + } + } + +done: + // latest chunk was full, so need to create a new one + if (*chunk_offset == -1) { + rc = vec0_new_chunk(p, partitionKeyValues, chunk_rowid); + if (rc != SQLITE_OK) { + // IMP: V08441_25279 + vtab_set_error(&p->base, + VEC_INTERAL_ERROR "Could not insert a new vector chunk"); + rc = SQLITE_ERROR; // otherwise raises a DatabaseError and not operational + // error? + goto cleanup; + } + *chunk_offset = 0; + + // blobChunksValidity and pValidity are stale, pointing to the previous + // (full) chunk. to re-assign them + rc = sqlite3_blob_close(*blobChunksValidity); + sqlite3_free((void *)*bufferChunksValidity); + *blobChunksValidity = NULL; + *bufferChunksValidity = NULL; + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR + "unknown error, blobChunksValidity could not be closed, " + "please file an issue."); + rc = SQLITE_ERROR; + goto cleanup; + } + + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, + "validity", *chunk_rowid, 1, blobChunksValidity); + if (rc != SQLITE_OK) { + vtab_set_error( + &p->base, + VEC_INTERAL_ERROR + "Could not open validity blob for newly created chunk %s.%s.%lld", + p->schemaName, p->shadowChunksName, *chunk_rowid); + goto cleanup; + } + validitySize = sqlite3_blob_bytes(*blobChunksValidity); + if (validitySize != p->chunk_size / CHAR_BIT) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "validity blob size mismatch for newly created chunk " + "%s.%s.%lld. Exepcted %lld, got %lld", + p->schemaName, p->shadowChunksName, *chunk_rowid, + p->chunk_size / CHAR_BIT, validitySize); + goto cleanup; + } + *bufferChunksValidity = sqlite3_malloc(validitySize); + rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity, + validitySize, 0); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "could not read validity blob newly created chunk " + "%s.%s.%lld", + p->schemaName, p->shadowChunksName, *chunk_rowid); + goto cleanup; + } + } + + rc = SQLITE_OK; + +cleanup: + return rc; +} + +/** + * @brief Write the vector data into the provided vector blob at the given + * offset + * + * @param blobVectors SQLite BLOB to write to + * @param chunk_offset the "offset" (ie validity bitmap position) to write the + * vector to + * @param bVector pointer to the vector containing data + * @param dimensions how many dimensions the vector has + * @param element_type the vector type + * @return result of sqlite3_blob_write, SQLITE_OK on success, otherwise failure + */ +static int +vec0_write_vector_to_vector_blob(sqlite3_blob *blobVectors, i64 chunk_offset, + const void *bVector, size_t dimensions, + enum VectorElementType element_type) { + int n; + int offset; + + switch (element_type) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: + n = dimensions * sizeof(f32); + offset = chunk_offset * dimensions * sizeof(f32); + break; + case SQLITE_VEC_ELEMENT_TYPE_INT8: + n = dimensions * sizeof(i8); + offset = chunk_offset * dimensions * sizeof(i8); + break; + case SQLITE_VEC_ELEMENT_TYPE_BIT: + n = dimensions / CHAR_BIT; + offset = chunk_offset * dimensions / CHAR_BIT; + break; + } + + return sqlite3_blob_write(blobVectors, bVector, n, offset); +} + +/** + * @brief + * + * @param p vec0 virtual table + * @param chunk_rowid: which chunk to write to + * @param chunk_offset: the offset inside the chunk to write the vector to. + * @param rowid: the rowid of the inserting row + * @param vectorDatas: array of the vector data to insert + * @param blobValidity: writeable validity blob of the row's assigned chunk. + * @param validity: snapshot buffer of the valdity column from the row's + * assigned chunk. + * @return int SQLITE_OK on success, error code on failure + */ +int vec0Update_InsertWriteFinalStep(vec0_vtab *p, i64 chunk_rowid, + i64 chunk_offset, i64 rowid, + void *vectorDatas[], + sqlite3_blob *blobChunksValidity, + const unsigned char *bufferChunksValidity) { + int rc, brc; + sqlite3_blob *blobChunksRowids = NULL; + + // mark the validity bit for this row in the chunk's validity bitmap + // Get the byte offset of the bitmap + char unsigned bx = bufferChunksValidity[chunk_offset / CHAR_BIT]; + // set the bit at the chunk_offset position inside that byte + bx = bx | (1 << (chunk_offset % CHAR_BIT)); + // write that 1 byte + rc = sqlite3_blob_write(blobChunksValidity, &bx, 1, chunk_offset / CHAR_BIT); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, VEC_INTERAL_ERROR "could not mark validity bit "); + return rc; + } + + // Go insert the vector data into the vector chunk shadow tables + for (int i = 0; i < p->numVectorColumns; i++) { + sqlite3_blob *blobVectors; + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i], + "vectors", chunk_rowid, 1, &blobVectors); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "Error opening vector blob at %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid); + goto cleanup; + } + + i64 expected = + p->chunk_size * vector_column_byte_size(p->vector_columns[i]); + i64 actual = sqlite3_blob_bytes(blobVectors); + + if (actual != expected) { + // IMP: V16386_00456 + vtab_set_error( + &p->base, + VEC_INTERAL_ERROR + "vector blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid, expected, + actual); + rc = SQLITE_ERROR; + // already error, can ignore result code + sqlite3_blob_close(blobVectors); + goto cleanup; + }; + + rc = vec0_write_vector_to_vector_blob( + blobVectors, chunk_offset, vectorDatas[i], + p->vector_columns[i].dimensions, p->vector_columns[i].element_type); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "could not write vector blob on %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid); + rc = SQLITE_ERROR; + // already error, can ignore result code + sqlite3_blob_close(blobVectors); + goto cleanup; + } + rc = sqlite3_blob_close(blobVectors); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR + "could not close vector blob on %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid); + rc = SQLITE_ERROR; + goto cleanup; + } + } + + // write the new rowid to the rowids column of the _chunks table + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids", + chunk_rowid, 1, &blobChunksRowids); + if (rc != SQLITE_OK) { + // IMP: V09221_26060 + vtab_set_error(&p->base, + VEC_INTERAL_ERROR "could not open rowids blob on %s.%s.%lld", + p->schemaName, p->shadowChunksName, chunk_rowid); + goto cleanup; + } + i64 expected = p->chunk_size * sizeof(i64); + i64 actual = sqlite3_blob_bytes(blobChunksRowids); + if (expected != actual) { + // IMP: V12779_29618 + vtab_set_error( + &p->base, + VEC_INTERAL_ERROR + "rowids blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld", + p->schemaName, p->shadowChunksName, chunk_rowid, expected, actual); + rc = SQLITE_ERROR; + goto cleanup; + } + rc = sqlite3_blob_write(blobChunksRowids, &rowid, sizeof(i64), + chunk_offset * sizeof(i64)); + if (rc != SQLITE_OK) { + vtab_set_error( + &p->base, VEC_INTERAL_ERROR "could not write rowids blob on %s.%s.%lld", + p->schemaName, p->shadowChunksName, chunk_rowid); + rc = SQLITE_ERROR; + goto cleanup; + } + + // Now with all the vectors inserted, go back and update the _rowids table + // with the new chunk_rowid/chunk_offset values + rc = vec0_rowids_update_position(p, rowid, chunk_rowid, chunk_offset); + +cleanup: + brc = sqlite3_blob_close(blobChunksRowids); + if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) { + vtab_set_error( + &p->base, VEC_INTERAL_ERROR "could not close rowids blob on %s.%s.%lld", + p->schemaName, p->shadowChunksName, chunk_rowid); + return brc; + } + return rc; +} + +int vec0_write_metadata_value(vec0_vtab *p, int metadata_column_idx, i64 rowid, i64 chunk_id, i64 chunk_offset, sqlite3_value * v, int isupdate) { + int rc; + struct Vec0MetadataColumnDefinition * metadata_column = &p->metadata_columns[metadata_column_idx]; + vec0_metadata_column_kind kind = metadata_column->kind; + + // verify input value matches column type + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + if(sqlite3_value_type(v) != SQLITE_INTEGER || ((sqlite3_value_int(v) != 0) && (sqlite3_value_int(v) != 1))) { + rc = SQLITE_ERROR; + vtab_set_error(&p->base, "Expected 0 or 1 for BOOLEAN metadata column %.*s", metadata_column->name_length, metadata_column->name); + goto done; + } + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + if(sqlite3_value_type(v) != SQLITE_INTEGER) { + rc = SQLITE_ERROR; + vtab_set_error(&p->base, "Expected integer for INTEGER metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v))); + goto done; + } + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + if(sqlite3_value_type(v) != SQLITE_FLOAT) { + rc = SQLITE_ERROR; + vtab_set_error(&p->base, "Expected float for FLOAT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v))); + goto done; + } + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + if(sqlite3_value_type(v) != SQLITE_TEXT) { + rc = SQLITE_ERROR; + vtab_set_error(&p->base, "Expected text for TEXT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v))); + goto done; + } + break; + } + } + + sqlite3_blob * blobValue = NULL; + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_column_idx], "data", chunk_id, 1, &blobValue); + if(rc != SQLITE_OK) { + goto done; + } + + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + u8 block; + int value = sqlite3_value_int(v); + rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT)); + if(rc != SQLITE_OK) { + goto done; + } + + if (value) { + block |= 1 << (chunk_offset % CHAR_BIT); + } else { + block &= ~(1 << (chunk_offset % CHAR_BIT)); + } + + rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT); + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + i64 value = sqlite3_value_int64(v); + rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64)); + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + double value = sqlite3_value_double(v); + rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(double)); + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + int prev_n; + rc = sqlite3_blob_read(blobValue, &prev_n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + if(rc != SQLITE_OK) { + goto done; + } + + const char * s = (const char *) sqlite3_value_text(v); + int n = sqlite3_value_bytes(v); + u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + memcpy(view, &n, sizeof(int)); + memcpy(view+4, s, min(n, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH-4)); + + rc = sqlite3_blob_write(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + const char * zSql; + + if(isupdate && (prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)) { + zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " SET data = ?2 WHERE rowid = ?1", p->schemaName, p->tableName, metadata_column_idx); + }else { + zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " (rowid, data) VALUES (?1, ?2)", p->schemaName, p->tableName, metadata_column_idx); + } + if(!zSql) { + rc = SQLITE_NOMEM; + goto done; + } + sqlite3_stmt * stmt; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_bind_int64(stmt, 1, rowid); + sqlite3_bind_text(stmt, 2, s, n, SQLITE_STATIC); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + + if(rc != SQLITE_DONE) { + rc = SQLITE_ERROR; + goto done; + } + } + else if(prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_column_idx); + if(!zSql) { + rc = SQLITE_NOMEM; + goto done; + } + sqlite3_stmt * stmt; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + + if(rc != SQLITE_DONE) { + rc = SQLITE_ERROR; + goto done; + } + } + break; + } + } + + if(rc != SQLITE_OK) { + + } + rc = sqlite3_blob_close(blobValue); + if(rc != SQLITE_OK) { + goto done; + } + + done: + return rc; +} + + +/** + * @brief Handles INSERT INTO operations on a vec0 table. + * + * @return int SQLITE_OK on success, otherwise error code on failure + */ +int vec0Update_Insert(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv, + sqlite_int64 *pRowid) { + UNUSED_PARAMETER(argc); + vec0_vtab *p = (vec0_vtab *)pVTab; + int rc; + // Rowid for the inserted row, deterimined by the inserted ID + _rowids shadow + // table + i64 rowid; + + // Array to hold the vector data of the inserted row. Individual elements will + // have a lifetime bound to the argv[..] values. + void *vectorDatas[VEC0_MAX_VECTOR_COLUMNS]; + // Array to hold cleanup functions for vectorDatas[] + vector_cleanup cleanups[VEC0_MAX_VECTOR_COLUMNS]; + + sqlite3_value * partitionKeyValues[VEC0_MAX_PARTITION_COLUMNS]; + + // Rowid of the chunk in the _chunks shadow table that the row will be a part + // of. + i64 chunk_rowid; + // offset within the chunk where the rowid belongs + i64 chunk_offset; + + // a write-able blob of the validity column for the given chunk. Used to mark + // validity bit + sqlite3_blob *blobChunksValidity = NULL; + // buffer for the valididty column for the given chunk. Maybe not needed here? + const unsigned char *bufferChunksValidity = NULL; + int numReadVectors = 0; + + // Read all provided partition key values into partitionKeyValues + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) { + continue; + } + int partition_key_idx = p->user_column_idxs[i]; + partitionKeyValues[partition_key_idx] = argv[2+VEC0_COLUMN_USERN_START + i]; + + int new_value_type = sqlite3_value_type(partitionKeyValues[partition_key_idx]); + if((new_value_type != SQLITE_NULL) && (new_value_type != p->paritition_columns[partition_key_idx].type)) { + // IMP: V11454_28292 + vtab_set_error( + pVTab, + "Parition key type mismatch: The partition key column %.*s has type %s, but %s was provided.", + p->paritition_columns[partition_key_idx].name_length, + p->paritition_columns[partition_key_idx].name, + type_name(p->paritition_columns[partition_key_idx].type), + type_name(new_value_type) + ); + rc = SQLITE_ERROR; + goto cleanup; + } + } + + // read all the inserted vectors into vectorDatas, validate their lengths. + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) { + continue; + } + int vector_column_idx = p->user_column_idxs[i]; + sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i]; + size_t dimensions; + + char *pzError; + enum VectorElementType elementType; + rc = vector_from_value(valueVector, &vectorDatas[vector_column_idx], &dimensions, + &elementType, &cleanups[vector_column_idx], &pzError); + if (rc != SQLITE_OK) { + // IMP: V06519_23358 + vtab_set_error( + pVTab, "Inserted vector for the \"%.*s\" column is invalid: %z", + p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name, pzError); + rc = SQLITE_ERROR; + goto cleanup; + } + + numReadVectors++; + if (elementType != p->vector_columns[vector_column_idx].element_type) { + // IMP: V08221_25059 + vtab_set_error( + pVTab, + "Inserted vector for the \"%.*s\" column is expected to be of type " + "%s, but a %s vector was provided.", + p->vector_columns[i].name_length, p->vector_columns[i].name, + vector_subtype_name(p->vector_columns[i].element_type), + vector_subtype_name(elementType)); + rc = SQLITE_ERROR; + goto cleanup; + } + + if (dimensions != p->vector_columns[vector_column_idx].dimensions) { + // IMP: V01145_17984 + vtab_set_error( + pVTab, + "Dimension mismatch for inserted vector for the \"%.*s\" column. " + "Expected %d dimensions but received %d.", + p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name, + p->vector_columns[vector_column_idx].dimensions, dimensions); + rc = SQLITE_ERROR; + goto cleanup; + } + } + + // Cannot insert a value in the hidden "distance" column + if (sqlite3_value_type(argv[2 + vec0_column_distance_idx(p)]) != + SQLITE_NULL) { + // IMP: V24228_08298 + vtab_set_error(pVTab, + "A value was provided for the hidden \"distance\" column."); + rc = SQLITE_ERROR; + goto cleanup; + } + // Cannot insert a value in the hidden "k" column + if (sqlite3_value_type(argv[2 + vec0_column_k_idx(p)]) != SQLITE_NULL) { + // IMP: V11875_28713 + vtab_set_error(pVTab, "A value was provided for the hidden \"k\" column."); + rc = SQLITE_ERROR; + goto cleanup; + } + + // Step #1: Insert/get a rowid for this row, from the _rowids table. + rc = vec0Update_InsertRowidStep(p, argv[2 + VEC0_COLUMN_ID], &rowid); + if (rc != SQLITE_OK) { + goto cleanup; + } + + // Step #2: Find the next "available" position in the _chunks table for this + // row. + rc = vec0Update_InsertNextAvailableStep(p, partitionKeyValues, + &chunk_rowid, &chunk_offset, + &blobChunksValidity, + &bufferChunksValidity); + if (rc != SQLITE_OK) { + goto cleanup; + } + + // Step #3: With the next available chunk position, write out all the vectors + // to their specified location. + rc = vec0Update_InsertWriteFinalStep(p, chunk_rowid, chunk_offset, rowid, + vectorDatas, blobChunksValidity, + bufferChunksValidity); + if (rc != SQLITE_OK) { + goto cleanup; + } + + if(p->numAuxiliaryColumns > 0) { + sqlite3_stmt *stmt; + sqlite3_str * s = sqlite3_str_new(NULL); + sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_AUXILIARY_NAME "(rowid ", p->schemaName, p->tableName); + for(int i = 0; i < p->numAuxiliaryColumns; i++) { + sqlite3_str_appendf(s, ", value%02d", i); + } + sqlite3_str_appendall(s, ") VALUES (? "); + for(int i = 0; i < p->numAuxiliaryColumns; i++) { + sqlite3_str_appendall(s, ", ?"); + } + sqlite3_str_appendall(s, ")"); + char * zSql = sqlite3_str_finish(s); + // TODO double check error handling ehre + if(!zSql) { + rc = SQLITE_NOMEM; + goto cleanup; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + if(rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_bind_int64(stmt, 1, rowid); + + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) { + continue; + } + int auxiliary_key_idx = p->user_column_idxs[i]; + sqlite3_value * v = argv[2+VEC0_COLUMN_USERN_START + i]; + int v_type = sqlite3_value_type(v); + if(v_type != SQLITE_NULL && (v_type != p->auxiliary_columns[auxiliary_key_idx].type)) { + sqlite3_finalize(stmt); + rc = SQLITE_CONSTRAINT; + vtab_set_error( + pVTab, + "Auxiliary column type mismatch: The auxiliary column %.*s has type %s, but %s was provided.", + p->auxiliary_columns[auxiliary_key_idx].name_length, + p->auxiliary_columns[auxiliary_key_idx].name, + type_name(p->auxiliary_columns[auxiliary_key_idx].type), + type_name(v_type) + ); + goto cleanup; + } + // first 1 is for 1-based indexing on sqlite3_bind_*, second 1 is to account for initial rowid parameter + sqlite3_bind_value(stmt, 1 + 1 + auxiliary_key_idx, v); + } + + rc = sqlite3_step(stmt); + if(rc != SQLITE_DONE) { + sqlite3_finalize(stmt); + rc = SQLITE_ERROR; + goto cleanup; + } + sqlite3_finalize(stmt); + } + + + for(int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) { + continue; + } + int metadata_idx = p->user_column_idxs[i]; + sqlite3_value *v = argv[2 + VEC0_COLUMN_USERN_START + i]; + rc = vec0_write_metadata_value(p, metadata_idx, rowid, chunk_rowid, chunk_offset, v, 0); + if(rc != SQLITE_OK) { + goto cleanup; + } + } + + *pRowid = rowid; + rc = SQLITE_OK; + +cleanup: + for (int i = 0; i < numReadVectors; i++) { + cleanups[i](vectorDatas[i]); + } + sqlite3_free((void *)bufferChunksValidity); + int brc = sqlite3_blob_close(blobChunksValidity); + if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) { + vtab_set_error(&p->base, + VEC_INTERAL_ERROR "unknown error, blobChunksValidity could " + "not be closed, please file an issue"); + return brc; + } + return rc; +} + +int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id, + u64 chunk_offset) { + int rc, brc; + sqlite3_blob *blobChunksValidity = NULL; + char unsigned bx; + int validityOffset = chunk_offset / CHAR_BIT; + + // 2. ensure chunks.validity bit is 1, then set to 0 + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity", + chunk_id, 1, &blobChunksValidity); + if (rc != SQLITE_OK) { + // IMP: V26002_10073 + vtab_set_error(&p->base, "could not open validity blob for %s.%s.%lld", + p->schemaName, p->shadowChunksName, chunk_id); + return SQLITE_ERROR; + } + // will skip the sqlite3_blob_bytes(blobChunksValidity) check for now, + // the read below would catch it + + rc = sqlite3_blob_read(blobChunksValidity, &bx, sizeof(bx), validityOffset); + if (rc != SQLITE_OK) { + // IMP: V21193_05263 + vtab_set_error( + &p->base, "could not read validity blob for %s.%s.%lld at %d", + p->schemaName, p->shadowChunksName, chunk_id, validityOffset); + goto cleanup; + } + if (!(bx >> (chunk_offset % CHAR_BIT))) { + // IMP: V21193_05263 + rc = SQLITE_ERROR; + vtab_set_error( + &p->base, + "vec0 deletion error: validity bit is not set for %s.%s.%lld at %d", + p->schemaName, p->shadowChunksName, chunk_id, validityOffset); + goto cleanup; + } + char unsigned mask = ~(1 << (chunk_offset % CHAR_BIT)); + char result = bx & mask; + rc = sqlite3_blob_write(blobChunksValidity, &result, sizeof(bx), + validityOffset); + if (rc != SQLITE_OK) { + vtab_set_error( + &p->base, "could not write to validity blob for %s.%s.%lld at %d", + p->schemaName, p->shadowChunksName, chunk_id, validityOffset); + goto cleanup; + } + +cleanup: + + brc = sqlite3_blob_close(blobChunksValidity); + if (rc != SQLITE_OK) + return rc; + if (brc != SQLITE_OK) { + vtab_set_error(&p->base, + "vec0 deletion error: Error commiting validity blob " + "transaction on %s.%s.%lld at %d", + p->schemaName, p->shadowChunksName, chunk_id, + validityOffset); + return brc; + } + return SQLITE_OK; +} + +int vec0Update_Delete_ClearRowid(vec0_vtab *p, i64 chunk_id, + u64 chunk_offset) { + int rc, brc; + sqlite3_blob *blobChunksRowids = NULL; + i64 zero = 0; + + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids", + chunk_id, 1, &blobChunksRowids); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "could not open rowids blob for %s.%s.%lld", + p->schemaName, p->shadowChunksName, chunk_id); + return SQLITE_ERROR; + } + + rc = sqlite3_blob_write(blobChunksRowids, &zero, sizeof(zero), + chunk_offset * sizeof(i64)); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + "could not write to rowids blob for %s.%s.%lld at %llu", + p->schemaName, p->shadowChunksName, chunk_id, chunk_offset); + } + + brc = sqlite3_blob_close(blobChunksRowids); + if (rc != SQLITE_OK) + return rc; + if (brc != SQLITE_OK) { + vtab_set_error(&p->base, + "vec0 deletion error: Error commiting rowids blob " + "transaction on %s.%s.%lld at %llu", + p->schemaName, p->shadowChunksName, chunk_id, chunk_offset); + return brc; + } + return SQLITE_OK; +} + +int vec0Update_Delete_ClearVectors(vec0_vtab *p, i64 chunk_id, + u64 chunk_offset) { + int rc, brc; + for (int i = 0; i < p->numVectorColumns; i++) { + sqlite3_blob *blobVectors = NULL; + size_t n = vector_column_byte_size(p->vector_columns[i]); + + rc = sqlite3_blob_open(p->db, p->schemaName, + p->shadowVectorChunksNames[i], "vectors", + chunk_id, 1, &blobVectors); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + "could not open vector blob for %s.%s.%lld column %d", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id, i); + return SQLITE_ERROR; + } + + void *zeroBuf = sqlite3_malloc(n); + if (!zeroBuf) { + sqlite3_blob_close(blobVectors); + return SQLITE_NOMEM; + } + memset(zeroBuf, 0, n); + + rc = sqlite3_blob_write(blobVectors, zeroBuf, n, chunk_offset * n); + sqlite3_free(zeroBuf); + if (rc != SQLITE_OK) { + vtab_set_error( + &p->base, + "could not write to vector blob for %s.%s.%lld at %llu column %d", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id, + chunk_offset, i); + } + + brc = sqlite3_blob_close(blobVectors); + if (rc != SQLITE_OK) + return rc; + if (brc != SQLITE_OK) { + vtab_set_error(&p->base, + "vec0 deletion error: Error commiting vector blob " + "transaction on %s.%s.%lld column %d", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id, i); + return brc; + } + } + return SQLITE_OK; +} + +int vec0Update_Delete_DeleteChunkIfEmpty(vec0_vtab *p, i64 chunk_id, + int *deleted) { + int rc, brc; + sqlite3_blob *blobValidity = NULL; + *deleted = 0; + + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity", + chunk_id, 0, &blobValidity); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, + "could not open validity blob for chunk %lld", chunk_id); + return SQLITE_ERROR; + } + + int validitySize = sqlite3_blob_bytes(blobValidity); + unsigned char *validityBuf = sqlite3_malloc(validitySize); + if (!validityBuf) { + sqlite3_blob_close(blobValidity); + return SQLITE_NOMEM; + } + + rc = sqlite3_blob_read(blobValidity, validityBuf, validitySize, 0); + brc = sqlite3_blob_close(blobValidity); + if (rc != SQLITE_OK) { + sqlite3_free(validityBuf); + return rc; + } + if (brc != SQLITE_OK) { + sqlite3_free(validityBuf); + return brc; + } + + int allZero = 1; + for (int i = 0; i < validitySize; i++) { + if (validityBuf[i] != 0) { + allZero = 0; + break; + } + } + sqlite3_free(validityBuf); + + if (!allZero) { + return SQLITE_OK; + } + + // All validity bits are zero — delete this chunk and its associated data + char *zSql; + sqlite3_stmt *stmt; + + // Delete from _chunks + zSql = sqlite3_mprintf( + "DELETE FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE rowid = ?", + p->schemaName, p->tableName); + if (!zSql) + return SQLITE_NOMEM; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) + return rc; + sqlite3_bind_int64(stmt, 1, chunk_id); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if (rc != SQLITE_DONE) + return SQLITE_ERROR; + + // Delete from each _vector_chunksNN + for (int i = 0; i < p->numVectorColumns; i++) { + zSql = sqlite3_mprintf( + "DELETE FROM " VEC0_SHADOW_VECTOR_N_NAME " WHERE rowid = ?", + p->schemaName, p->tableName, i); + if (!zSql) + return SQLITE_NOMEM; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) + return rc; + sqlite3_bind_int64(stmt, 1, chunk_id); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if (rc != SQLITE_DONE) + return SQLITE_ERROR; + } + + // Delete from each _metadatachunksNN + for (int i = 0; i < p->numMetadataColumns; i++) { + zSql = sqlite3_mprintf( + "DELETE FROM " VEC0_SHADOW_METADATA_N_NAME " WHERE rowid = ?", + p->schemaName, p->tableName, i); + if (!zSql) + return SQLITE_NOMEM; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) + return rc; + sqlite3_bind_int64(stmt, 1, chunk_id); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if (rc != SQLITE_DONE) + return SQLITE_ERROR; + } + + // Invalidate cached stmtLatestChunk so it gets re-prepared on next insert + if (p->stmtLatestChunk) { + sqlite3_finalize(p->stmtLatestChunk); + p->stmtLatestChunk = NULL; + } + + *deleted = 1; + return SQLITE_OK; +} + +int vec0Update_Delete_DeleteRowids(vec0_vtab *p, i64 rowid) { + int rc; + sqlite3_stmt *stmt = NULL; + + char *zSql = + sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?", + p->schemaName, p->tableName); + if (!zSql) { + return SQLITE_NOMEM; + } + + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + if (rc != SQLITE_DONE) { + goto cleanup; + } + rc = SQLITE_OK; + +cleanup: + sqlite3_finalize(stmt); + return rc; +} + +int vec0Update_Delete_DeleteAux(vec0_vtab *p, i64 rowid) { + int rc; + sqlite3_stmt *stmt = NULL; + + char *zSql = + sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?", + p->schemaName, p->tableName); + if (!zSql) { + return SQLITE_NOMEM; + } + + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + sqlite3_free(zSql); + if (rc != SQLITE_OK) { + goto cleanup; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + if (rc != SQLITE_DONE) { + goto cleanup; + } + rc = SQLITE_OK; + +cleanup: + sqlite3_finalize(stmt); + return rc; +} + +int vec0Update_Delete_ClearMetadata(vec0_vtab *p, int metadata_idx, i64 rowid, i64 chunk_id, + u64 chunk_offset) { + int rc; + sqlite3_blob * blobValue; + vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind; + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 1, &blobValue); + if(rc != SQLITE_OK) { + return rc; + } + + switch(kind) { + case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { + u8 block; + rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT)); + if(rc != SQLITE_OK) { + goto done; + } + + block &= ~(1 << (chunk_offset % CHAR_BIT)); + rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT); + break; + } + case VEC0_METADATA_COLUMN_KIND_INTEGER: { + i64 v = 0; + rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(i64)); + break; + } + case VEC0_METADATA_COLUMN_KIND_FLOAT: { + double v = 0; + rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(double)); + break; + } + case VEC0_METADATA_COLUMN_KIND_TEXT: { + int n; + rc = sqlite3_blob_read(blobValue, &n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + if(rc != SQLITE_OK) { + goto done; + } + + u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH]; + memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + rc = sqlite3_blob_write(blobValue, &view, sizeof(view), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH); + if(rc != SQLITE_OK) { + goto done; + } + + if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) { + const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx); + if(!zSql) { + rc = SQLITE_NOMEM; + goto done; + } + sqlite3_stmt * stmt; + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + if(rc != SQLITE_OK) { + goto done; + } + sqlite3_bind_int64(stmt, 1, rowid); + rc = sqlite3_step(stmt); + sqlite3_finalize(stmt); + if(rc != SQLITE_DONE) { + rc = SQLITE_ERROR; + goto done; + } + // Fix for https://github.com/asg017/sqlite-vec/issues/274 + // sqlite3_step returns SQLITE_DONE (101) on DML success, but the + // `done:` epilogue treats anything other than SQLITE_OK as an error. + // Without this, SQLITE_DONE propagates up to vec0Update_Delete, + // which aborts the DELETE scan and silently drops remaining rows. + rc = SQLITE_OK; + } + break; + } + } + int rc2; + done: + rc2 = sqlite3_blob_close(blobValue); + if(rc == SQLITE_OK) { + return rc2; + } + return rc; +} + +int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) { + vec0_vtab *p = (vec0_vtab *)pVTab; + int rc; + i64 rowid; + i64 chunk_id; + i64 chunk_offset; + + if (p->pkIsText) { + rc = vec0_rowid_from_id(p, idValue, &rowid); + if (rc != SQLITE_OK) { + return rc; + } + } else { + rowid = sqlite3_value_int64(idValue); + } + + // 1. Find chunk position for given rowid + // 2. Ensure that validity bit for position is 1, then set to 0 + // 3. Zero out rowid in chunks.rowid + // 4. Zero out vector data in all vector column chunks + // 5. Delete value in _rowids table + + // 1. get chunk_id and chunk_offset from _rowids + rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + + // 2. clear validity bit + rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + + // 3. zero out rowid in chunks.rowids + rc = vec0Update_Delete_ClearRowid(p, chunk_id, chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + + // 4. zero out any data in vector chunks tables + rc = vec0Update_Delete_ClearVectors(p, chunk_id, chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + + // 5. delete from _rowids table + rc = vec0Update_Delete_DeleteRowids(p, rowid); + if (rc != SQLITE_OK) { + return rc; + } + + // 6. delete any auxiliary rows + if(p->numAuxiliaryColumns > 0) { + rc = vec0Update_Delete_DeleteAux(p, rowid); + if (rc != SQLITE_OK) { + return rc; + } + } + + // 7. delete metadata + for(int i = 0; i < p->numMetadataColumns; i++) { + rc = vec0Update_Delete_ClearMetadata(p, i, rowid, chunk_id, chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + } + + // 8. reclaim chunk if fully empty + { + int chunkDeleted; + rc = vec0Update_Delete_DeleteChunkIfEmpty(p, chunk_id, &chunkDeleted); + if (rc != SQLITE_OK) { + return rc; + } + } + + return SQLITE_OK; +} + +int vec0Update_UpdateAuxColumn(vec0_vtab *p, int auxiliary_column_idx, sqlite3_value * value, i64 rowid) { + int rc; + sqlite3_stmt *stmt; + const char * zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_AUXILIARY_NAME " SET value%02d = ? WHERE rowid = ?", p->schemaName, p->tableName, auxiliary_column_idx); + if(!zSql) { + return SQLITE_NOMEM; + } + rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); + if(rc != SQLITE_OK) { + return rc; + } + sqlite3_bind_value(stmt, 1, value); + sqlite3_bind_int64(stmt, 2, rowid); + rc = sqlite3_step(stmt); + if(rc != SQLITE_DONE) { + sqlite3_finalize(stmt); + return SQLITE_ERROR; + } + sqlite3_finalize(stmt); + return SQLITE_OK; +} + +int vec0Update_UpdateVectorColumn(vec0_vtab *p, i64 chunk_id, i64 chunk_offset, + int i, sqlite3_value *valueVector) { + int rc; + + sqlite3_blob *blobVectors = NULL; + + char *pzError; + size_t dimensions; + enum VectorElementType elementType; + void *vector; + vector_cleanup cleanup = vector_cleanup_noop; + // https://github.com/asg017/sqlite-vec/issues/53 + rc = vector_from_value(valueVector, &vector, &dimensions, &elementType, + &cleanup, &pzError); + if (rc != SQLITE_OK) { + // IMP: V15203_32042 + vtab_set_error( + &p->base, "Updated vector for the \"%.*s\" column is invalid: %z", + p->vector_columns[i].name_length, p->vector_columns[i].name, pzError); + rc = SQLITE_ERROR; + goto cleanup; + } + if (elementType != p->vector_columns[i].element_type) { + // IMP: V03643_20481 + vtab_set_error( + &p->base, + "Updated vector for the \"%.*s\" column is expected to be of type " + "%s, but a %s vector was provided.", + p->vector_columns[i].name_length, p->vector_columns[i].name, + vector_subtype_name(p->vector_columns[i].element_type), + vector_subtype_name(elementType)); + rc = SQLITE_ERROR; + goto cleanup; + } + if (dimensions != p->vector_columns[i].dimensions) { + // IMP: V25739_09810 + vtab_set_error( + &p->base, + "Dimension mismatch for new updated vector for the \"%.*s\" column. " + "Expected %d dimensions but received %d.", + p->vector_columns[i].name_length, p->vector_columns[i].name, + p->vector_columns[i].dimensions, dimensions); + rc = SQLITE_ERROR; + goto cleanup; + } + + rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i], + "vectors", chunk_id, 1, &blobVectors); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "Could not open vectors blob for %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id); + goto cleanup; + } + rc = vec0_write_vector_to_vector_blob(blobVectors, chunk_offset, vector, + p->vector_columns[i].dimensions, + p->vector_columns[i].element_type); + if (rc != SQLITE_OK) { + vtab_set_error(&p->base, "Could not write to vectors blob for %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id); + goto cleanup; + } + +cleanup: + cleanup(vector); + int brc = sqlite3_blob_close(blobVectors); + if (rc != SQLITE_OK) { + return rc; + } + if (brc != SQLITE_OK) { + vtab_set_error( + &p->base, + "Could not commit blob transaction for vectors blob for %s.%s.%lld", + p->schemaName, p->shadowVectorChunksNames[i], chunk_id); + return brc; + } + return SQLITE_OK; +} + +int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) { + UNUSED_PARAMETER(argc); + vec0_vtab *p = (vec0_vtab *)pVTab; + int rc; + i64 chunk_id; + i64 chunk_offset; + + i64 rowid; + if (p->pkIsText) { + const char *a = (const char *)sqlite3_value_text(argv[0]); + const char *b = (const char *)sqlite3_value_text(argv[1]); + // IMP: V08886_25725 + if ((sqlite3_value_bytes(argv[0]) != sqlite3_value_bytes(argv[1])) || + strncmp(a, b, sqlite3_value_bytes(argv[0])) != 0) { + vtab_set_error(pVTab, + "UPDATEs on vec0 primary key values are not allowed."); + return SQLITE_ERROR; + } + rc = vec0_rowid_from_id(p, argv[0], &rowid); + if (rc != SQLITE_OK) { + return rc; + } + } else { + rowid = sqlite3_value_int64(argv[0]); + } + + // 1) get chunk_id and chunk_offset from _rowids + rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset); + if (rc != SQLITE_OK) { + return rc; + } + + // 2) update any partition key values + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) { + continue; + } + sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i]; + if(sqlite3_value_nochange(value)) { + continue; + } + vtab_set_error(pVTab, "UPDATE on partition key columns are not supported yet. "); + return SQLITE_ERROR; + } + + // 3) handle auxiliary column updates + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) { + continue; + } + int auxiliary_column_idx = p->user_column_idxs[i]; + sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i]; + if(sqlite3_value_nochange(value)) { + continue; + } + rc = vec0Update_UpdateAuxColumn(p, auxiliary_column_idx, value, rowid); + if(rc != SQLITE_OK) { + return SQLITE_ERROR; + } + } + + // 4) handle metadata column updates + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) { + continue; + } + int metadata_column_idx = p->user_column_idxs[i]; + sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i]; + if(sqlite3_value_nochange(value)) { + continue; + } + rc = vec0_write_metadata_value(p, metadata_column_idx, rowid, chunk_id, chunk_offset, value, 1); + if(rc != SQLITE_OK) { + return rc; + } + } + + // 5) iterate over all new vectors, update the vectors + for (int i = 0; i < vec0_num_defined_user_columns(p); i++) { + if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) { + continue; + } + int vector_idx = p->user_column_idxs[i]; + sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i]; + // in vec0Column, we check sqlite3_vtab_nochange() on vector columns. + // If the vector column isn't being changed, we return NULL; + // That's not great, that means vector columns can never be NULLABLE + // (bc we cant distinguish if an updated vector is truly NULL or nochange). + // Also it means that if someone tries to run `UPDATE v SET X = NULL`, + // we can't effectively detect and raise an error. + // A better solution would be to use a custom result_type for "empty", + // but subtypes don't appear to survive xColumn -> xUpdate, it's always 0. + // So for now, we'll just use NULL and warn people to not SET X = NULL + // in the docs. + if (sqlite3_value_type(valueVector) == SQLITE_NULL) { + continue; + } + + rc = vec0Update_UpdateVectorColumn(p, chunk_id, chunk_offset, vector_idx, + valueVector); + if (rc != SQLITE_OK) { + return SQLITE_ERROR; + } + } + + return SQLITE_OK; +} + +static int vec0Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv, + sqlite_int64 *pRowid) { + // DELETE operation + if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) { + return vec0Update_Delete(pVTab, argv[0]); + } + // INSERT operation + else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) { + return vec0Update_Insert(pVTab, argc, argv, pRowid); + } + // UPDATE operation + else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) { + return vec0Update_Update(pVTab, argc, argv); + } else { + vtab_set_error(pVTab, "Unrecognized xUpdate operation provided for vec0."); + return SQLITE_ERROR; + } +} + +static int vec0ShadowName(const char *zName) { + static const char *azName[] = { + "rowids", "chunks", "auxiliary", "info", + + // Up to VEC0_MAX_METADATA_COLUMNS + // TODO be smarter about this man + "metadatachunks00", + "metadatachunks01", + "metadatachunks02", + "metadatachunks03", + "metadatachunks04", + "metadatachunks05", + "metadatachunks06", + "metadatachunks07", + "metadatachunks08", + "metadatachunks09", + "metadatachunks10", + "metadatachunks11", + "metadatachunks12", + "metadatachunks13", + "metadatachunks14", + "metadatachunks15", + + // Up to + "metadatatext00", + "metadatatext01", + "metadatatext02", + "metadatatext03", + "metadatatext04", + "metadatatext05", + "metadatatext06", + "metadatatext07", + "metadatatext08", + "metadatatext09", + "metadatatext10", + "metadatatext11", + "metadatatext12", + "metadatatext13", + "metadatatext14", + "metadatatext15", + }; + + for (size_t i = 0; i < sizeof(azName) / sizeof(azName[0]); i++) { + if (sqlite3_stricmp(zName, azName[i]) == 0) + return 1; + } + //for(size_t i = 0; i < )"vector_chunks", "metadatachunks" + return 0; +} + +static int vec0Begin(sqlite3_vtab *pVTab) { + UNUSED_PARAMETER(pVTab); + return SQLITE_OK; +} +static int vec0Sync(sqlite3_vtab *pVTab) { + UNUSED_PARAMETER(pVTab); + vec0_vtab *p = (vec0_vtab *)pVTab; + if (p->stmtLatestChunk) { + sqlite3_finalize(p->stmtLatestChunk); + p->stmtLatestChunk = NULL; + } + if (p->stmtRowidsInsertRowid) { + sqlite3_finalize(p->stmtRowidsInsertRowid); + p->stmtRowidsInsertRowid = NULL; + } + if (p->stmtRowidsInsertId) { + sqlite3_finalize(p->stmtRowidsInsertId); + p->stmtRowidsInsertId = NULL; + } + if (p->stmtRowidsUpdatePosition) { + sqlite3_finalize(p->stmtRowidsUpdatePosition); + p->stmtRowidsUpdatePosition = NULL; + } + if (p->stmtRowidsGetChunkPosition) { + sqlite3_finalize(p->stmtRowidsGetChunkPosition); + p->stmtRowidsGetChunkPosition = NULL; + } + return SQLITE_OK; +} +static int vec0Commit(sqlite3_vtab *pVTab) { + UNUSED_PARAMETER(pVTab); + return SQLITE_OK; +} +static int vec0Rollback(sqlite3_vtab *pVTab) { + UNUSED_PARAMETER(pVTab); + return SQLITE_OK; +} + +static sqlite3_module vec0Module = { + /* iVersion */ 3, + /* xCreate */ vec0Create, + /* xConnect */ vec0Connect, + /* xBestIndex */ vec0BestIndex, + /* xDisconnect */ vec0Disconnect, + /* xDestroy */ vec0Destroy, + /* xOpen */ vec0Open, + /* xClose */ vec0Close, + /* xFilter */ vec0Filter, + /* xNext */ vec0Next, + /* xEof */ vec0Eof, + /* xColumn */ vec0Column, + /* xRowid */ vec0Rowid, + /* xUpdate */ vec0Update, + /* xBegin */ vec0Begin, + /* xSync */ vec0Sync, + /* xCommit */ vec0Commit, + /* xRollback */ vec0Rollback, + /* xFindFunction */ 0, + /* xRename */ 0, // https://github.com/asg017/sqlite-vec/issues/43 + /* xSavepoint */ 0, + /* xRelease */ 0, + /* xRollbackTo */ 0, + /* xShadowName */ vec0ShadowName, +#if SQLITE_VERSION_NUMBER >= 3044000 + /* xIntegrity */ 0, // https://github.com/asg017/sqlite-vec/issues/44 +#endif +}; +#pragma endregion + +static char *POINTER_NAME_STATIC_BLOB_DEF = "vec0-static_blob_def"; +struct static_blob_definition { + void *p; + size_t dimensions; + size_t nvectors; + enum VectorElementType element_type; +}; +static void vec_static_blob_from_raw(sqlite3_context *context, int argc, + sqlite3_value **argv) { + + assert(argc == 4); + struct static_blob_definition *p; + p = sqlite3_malloc(sizeof(*p)); + if (!p) { + sqlite3_result_error_nomem(context); + return; + } + memset(p, 0, sizeof(*p)); + p->p = (void *)sqlite3_value_int64(argv[0]); + p->element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32; + p->dimensions = sqlite3_value_int64(argv[2]); + p->nvectors = sqlite3_value_int64(argv[3]); + sqlite3_result_pointer(context, p, POINTER_NAME_STATIC_BLOB_DEF, + sqlite3_free); +} +#pragma region vec_static_blobs() table function + +#define MAX_STATIC_BLOBS 16 + +typedef struct static_blob static_blob; +struct static_blob { + char *name; + void *p; + size_t dimensions; + size_t nvectors; + enum VectorElementType element_type; +}; + +typedef struct vec_static_blob_data vec_static_blob_data; +struct vec_static_blob_data { + static_blob static_blobs[MAX_STATIC_BLOBS]; +}; + +typedef struct vec_static_blobs_vtab vec_static_blobs_vtab; +struct vec_static_blobs_vtab { + sqlite3_vtab base; + vec_static_blob_data *data; +}; + +typedef struct vec_static_blobs_cursor vec_static_blobs_cursor; +struct vec_static_blobs_cursor { + sqlite3_vtab_cursor base; + sqlite3_int64 iRowid; +}; + +static int vec_static_blobsConnect(sqlite3 *db, void *pAux, int argc, + const char *const *argv, + sqlite3_vtab **ppVtab, char **pzErr) { + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + UNUSED_PARAMETER(pzErr); + + vec_static_blobs_vtab *pNew; +#define VEC_STATIC_BLOBS_NAME 0 +#define VEC_STATIC_BLOBS_DATA 1 +#define VEC_STATIC_BLOBS_DIMENSIONS 2 +#define VEC_STATIC_BLOBS_COUNT 3 + int rc = sqlite3_declare_vtab( + db, "CREATE TABLE x(name, data, dimensions hidden, count hidden)"); + if (rc == SQLITE_OK) { + pNew = sqlite3_malloc(sizeof(*pNew)); + *ppVtab = (sqlite3_vtab *)pNew; + if (pNew == 0) + return SQLITE_NOMEM; + memset(pNew, 0, sizeof(*pNew)); + pNew->data = pAux; + } + return rc; +} + +static int vec_static_blobsDisconnect(sqlite3_vtab *pVtab) { + vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVtab; + sqlite3_free(p); + return SQLITE_OK; +} + +static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc, + sqlite3_value **argv, sqlite_int64 *pRowid) { + UNUSED_PARAMETER(pRowid); + vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVTab; + // DELETE operation + if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) { + return SQLITE_ERROR; + } + // INSERT operation + else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) { + const char *key = + (const char *)sqlite3_value_text(argv[2 + VEC_STATIC_BLOBS_NAME]); + int idx = -1; + for (int i = 0; i < MAX_STATIC_BLOBS; i++) { + if (!p->data->static_blobs[i].name) { + p->data->static_blobs[i].name = sqlite3_mprintf("%s", key); + idx = i; + break; + } + } + if (idx < 0) + abort(); + struct static_blob_definition *def = sqlite3_value_pointer( + argv[2 + VEC_STATIC_BLOBS_DATA], POINTER_NAME_STATIC_BLOB_DEF); + p->data->static_blobs[idx].p = def->p; + p->data->static_blobs[idx].dimensions = def->dimensions; + p->data->static_blobs[idx].nvectors = def->nvectors; + p->data->static_blobs[idx].element_type = def->element_type; + + return SQLITE_OK; + } + // UPDATE operation + else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) { + return SQLITE_ERROR; + } + return SQLITE_ERROR; +} + +static int vec_static_blobsOpen(sqlite3_vtab *p, + sqlite3_vtab_cursor **ppCursor) { + UNUSED_PARAMETER(p); + vec_static_blobs_cursor *pCur; + pCur = sqlite3_malloc(sizeof(*pCur)); + if (pCur == 0) + return SQLITE_NOMEM; + memset(pCur, 0, sizeof(*pCur)); + *ppCursor = &pCur->base; + return SQLITE_OK; +} + +static int vec_static_blobsClose(sqlite3_vtab_cursor *cur) { + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; + sqlite3_free(pCur); + return SQLITE_OK; +} + +static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab, + sqlite3_index_info *pIdxInfo) { + UNUSED_PARAMETER(pVTab); + pIdxInfo->idxNum = 1; + pIdxInfo->estimatedCost = (double)10; + pIdxInfo->estimatedRows = 10; + return SQLITE_OK; +} + +static int vec_static_blobsNext(sqlite3_vtab_cursor *cur); +static int vec_static_blobsFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, + const char *idxStr, int argc, + sqlite3_value **argv) { + UNUSED_PARAMETER(idxNum); + UNUSED_PARAMETER(idxStr); + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)pVtabCursor; + pCur->iRowid = -1; + vec_static_blobsNext(pVtabCursor); + return SQLITE_OK; +} + +static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur, + sqlite_int64 *pRowid) { + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; + *pRowid = pCur->iRowid; + return SQLITE_OK; +} + +static int vec_static_blobsNext(sqlite3_vtab_cursor *cur) { + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; + vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pCur->base.pVtab; + pCur->iRowid++; + while (pCur->iRowid < MAX_STATIC_BLOBS) { + if (p->data->static_blobs[pCur->iRowid].name) { + return SQLITE_OK; + } + pCur->iRowid++; + } + return SQLITE_OK; +} + +static int vec_static_blobsEof(sqlite3_vtab_cursor *cur) { + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; + return pCur->iRowid >= MAX_STATIC_BLOBS; +} + +static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur, + sqlite3_context *context, int i) { + vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; + vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)cur->pVtab; + switch (i) { + case VEC_STATIC_BLOBS_NAME: + sqlite3_result_text(context, p->data->static_blobs[pCur->iRowid].name, -1, + SQLITE_TRANSIENT); + break; + case VEC_STATIC_BLOBS_DATA: + sqlite3_result_null(context); + break; + case VEC_STATIC_BLOBS_DIMENSIONS: + sqlite3_result_int64(context, + p->data->static_blobs[pCur->iRowid].dimensions); + break; + case VEC_STATIC_BLOBS_COUNT: + sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].nvectors); + break; + } + return SQLITE_OK; +} + +static sqlite3_module vec_static_blobsModule = { + /* iVersion */ 3, + /* xCreate */ 0, + /* xConnect */ vec_static_blobsConnect, + /* xBestIndex */ vec_static_blobsBestIndex, + /* xDisconnect */ vec_static_blobsDisconnect, + /* xDestroy */ 0, + /* xOpen */ vec_static_blobsOpen, + /* xClose */ vec_static_blobsClose, + /* xFilter */ vec_static_blobsFilter, + /* xNext */ vec_static_blobsNext, + /* xEof */ vec_static_blobsEof, + /* xColumn */ vec_static_blobsColumn, + /* xRowid */ vec_static_blobsRowid, + /* xUpdate */ vec_static_blobsUpdate, + /* xBegin */ 0, + /* xSync */ 0, + /* xCommit */ 0, + /* xRollback */ 0, + /* xFindMethod */ 0, + /* xRename */ 0, + /* xSavepoint */ 0, + /* xRelease */ 0, + /* xRollbackTo */ 0, + /* xShadowName */ 0, +#if SQLITE_VERSION_NUMBER >= 3044000 + /* xIntegrity */ 0 +#endif +}; +#pragma endregion + +#pragma region vec_static_blob_entries() table function + +typedef struct vec_static_blob_entries_vtab vec_static_blob_entries_vtab; +struct vec_static_blob_entries_vtab { + sqlite3_vtab base; + static_blob *blob; +}; +typedef enum { + VEC_SBE__QUERYPLAN_FULLSCAN = 1, + VEC_SBE__QUERYPLAN_KNN = 2 +} vec_sbe_query_plan; + +struct sbe_query_knn_data { + i64 k; + i64 k_used; + // Array of rowids of size k. Must be freed with sqlite3_free(). + i32 *rowids; + // Array of distances of size k. Must be freed with sqlite3_free(). + f32 *distances; + i64 current_idx; +}; +void sbe_query_knn_data_clear(struct sbe_query_knn_data *knn_data) { + if (!knn_data) + return; + + if (knn_data->rowids) { + sqlite3_free(knn_data->rowids); + knn_data->rowids = NULL; + } + if (knn_data->distances) { + sqlite3_free(knn_data->distances); + knn_data->distances = NULL; + } +} + +typedef struct vec_static_blob_entries_cursor vec_static_blob_entries_cursor; +struct vec_static_blob_entries_cursor { + sqlite3_vtab_cursor base; + sqlite3_int64 iRowid; + vec_sbe_query_plan query_plan; + struct sbe_query_knn_data *knn_data; +}; + +static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc, + const char *const *argv, + sqlite3_vtab **ppVtab, char **pzErr) { + UNUSED_PARAMETER(argc); + UNUSED_PARAMETER(argv); + UNUSED_PARAMETER(pzErr); + vec_static_blob_data *blob_data = pAux; + int idx = -1; + for (int i = 0; i < MAX_STATIC_BLOBS; i++) { + if (!blob_data->static_blobs[i].name) + continue; + if (strncmp(blob_data->static_blobs[i].name, argv[3], + strlen(blob_data->static_blobs[i].name)) == 0) { + idx = i; + break; + } + } + if (idx < 0) + abort(); + vec_static_blob_entries_vtab *pNew; +#define VEC_STATIC_BLOB_ENTRIES_VECTOR 0 +#define VEC_STATIC_BLOB_ENTRIES_DISTANCE 1 +#define VEC_STATIC_BLOB_ENTRIES_K 2 + int rc = sqlite3_declare_vtab( + db, "CREATE TABLE x(vector, distance hidden, k hidden)"); + if (rc == SQLITE_OK) { + pNew = sqlite3_malloc(sizeof(*pNew)); + *ppVtab = (sqlite3_vtab *)pNew; + if (pNew == 0) + return SQLITE_NOMEM; + memset(pNew, 0, sizeof(*pNew)); + pNew->blob = &blob_data->static_blobs[idx]; + } + return rc; +} + +static int vec_static_blob_entriesCreate(sqlite3 *db, void *pAux, int argc, + const char *const *argv, + sqlite3_vtab **ppVtab, char **pzErr) { + return vec_static_blob_entriesConnect(db, pAux, argc, argv, ppVtab, pzErr); +} + +static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) { + vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVtab; + sqlite3_free(p); + return SQLITE_OK; +} + +static int vec_static_blob_entriesOpen(sqlite3_vtab *p, + sqlite3_vtab_cursor **ppCursor) { + UNUSED_PARAMETER(p); + vec_static_blob_entries_cursor *pCur; + pCur = sqlite3_malloc(sizeof(*pCur)); + if (pCur == 0) + return SQLITE_NOMEM; + memset(pCur, 0, sizeof(*pCur)); + *ppCursor = &pCur->base; + return SQLITE_OK; +} + +static int vec_static_blob_entriesClose(sqlite3_vtab_cursor *cur) { + vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; + sqlite3_free(pCur->knn_data); + sqlite3_free(pCur); + return SQLITE_OK; +} + +static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, + sqlite3_index_info *pIdxInfo) { + vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVTab; + int iMatchTerm = -1; + int iLimitTerm = -1; + // int iRowidTerm = -1; // https://github.com/asg017/sqlite-vec/issues/47 + int iKTerm = -1; + + for (int i = 0; i < pIdxInfo->nConstraint; i++) { + if (!pIdxInfo->aConstraint[i].usable) + continue; + + int iColumn = pIdxInfo->aConstraint[i].iColumn; + int op = pIdxInfo->aConstraint[i].op; + if (op == SQLITE_INDEX_CONSTRAINT_MATCH && + iColumn == VEC_STATIC_BLOB_ENTRIES_VECTOR) { + if (iMatchTerm > -1) { + // https://github.com/asg017/sqlite-vec/issues/51 + return SQLITE_ERROR; + } + iMatchTerm = i; + } + if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) { + iLimitTerm = i; + } + if (op == SQLITE_INDEX_CONSTRAINT_EQ && + iColumn == VEC_STATIC_BLOB_ENTRIES_K) { + iKTerm = i; + } + } + if (iMatchTerm >= 0) { + if (iLimitTerm < 0 && iKTerm < 0) { + // https://github.com/asg017/sqlite-vec/issues/51 + return SQLITE_ERROR; + } + if (iLimitTerm >= 0 && iKTerm >= 0) { + return SQLITE_ERROR; // limit or k, not both + } + if (pIdxInfo->nOrderBy < 1) { + vtab_set_error(pVTab, "ORDER BY distance required"); + return SQLITE_CONSTRAINT; + } + if (pIdxInfo->nOrderBy > 1) { + // https://github.com/asg017/sqlite-vec/issues/51 + vtab_set_error(pVTab, "more than 1 ORDER BY clause provided"); + return SQLITE_CONSTRAINT; + } + if (pIdxInfo->aOrderBy[0].iColumn != VEC_STATIC_BLOB_ENTRIES_DISTANCE) { + vtab_set_error(pVTab, "ORDER BY must be on the distance column"); + return SQLITE_CONSTRAINT; + } + if (pIdxInfo->aOrderBy[0].desc) { + vtab_set_error(pVTab, + "Only ascending in ORDER BY distance clause is supported, " + "DESC is not supported yet."); + return SQLITE_CONSTRAINT; + } + + pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN; + pIdxInfo->estimatedCost = (double)10; + pIdxInfo->estimatedRows = 10; + + pIdxInfo->orderByConsumed = 1; + pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1; + pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1; + if (iLimitTerm >= 0) { + pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = 2; + pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1; + } else { + pIdxInfo->aConstraintUsage[iKTerm].argvIndex = 2; + pIdxInfo->aConstraintUsage[iKTerm].omit = 1; + } + + } else { + pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN; + pIdxInfo->estimatedCost = (double)p->blob->nvectors; + pIdxInfo->estimatedRows = p->blob->nvectors; + } + return SQLITE_OK; +} + +static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, + int idxNum, const char *idxStr, + int argc, sqlite3_value **argv) { + UNUSED_PARAMETER(idxStr); + assert(argc >= 0 && argc <= 3); + vec_static_blob_entries_cursor *pCur = + (vec_static_blob_entries_cursor *)pVtabCursor; + vec_static_blob_entries_vtab *p = + (vec_static_blob_entries_vtab *)pCur->base.pVtab; + + if (idxNum == VEC_SBE__QUERYPLAN_KNN) { + assert(argc == 2); + pCur->query_plan = VEC_SBE__QUERYPLAN_KNN; + struct sbe_query_knn_data *knn_data; + knn_data = sqlite3_malloc(sizeof(*knn_data)); + if (!knn_data) { + return SQLITE_NOMEM; + } + memset(knn_data, 0, sizeof(*knn_data)); + + void *queryVector; + size_t dimensions; + enum VectorElementType elementType; + vector_cleanup cleanup; + char *err; + int rc = vector_from_value(argv[0], &queryVector, &dimensions, &elementType, + &cleanup, &err); + if (rc != SQLITE_OK) { + return SQLITE_ERROR; + } + if (elementType != p->blob->element_type) { + return SQLITE_ERROR; + } + if (dimensions != p->blob->dimensions) { + return SQLITE_ERROR; + } + + i64 k = min(sqlite3_value_int64(argv[1]), (i64)p->blob->nvectors); + if (k < 0) { + // HANDLE https://github.com/asg017/sqlite-vec/issues/55 + return SQLITE_ERROR; + } + if (k == 0) { + knn_data->k = 0; + pCur->knn_data = knn_data; + return SQLITE_OK; + } + + size_t bsize = (p->blob->nvectors + 7) & ~7; + + i32 *topk_rowids = sqlite3_malloc(k * sizeof(i32)); + if (!topk_rowids) { + // HANDLE https://github.com/asg017/sqlite-vec/issues/55 + return SQLITE_ERROR; + } + f32 *distances = sqlite3_malloc(bsize * sizeof(f32)); + if (!distances) { + // HANDLE https://github.com/asg017/sqlite-vec/issues/55 + return SQLITE_ERROR; + } + + for (size_t i = 0; i < p->blob->nvectors; i++) { + // https://github.com/asg017/sqlite-vec/issues/52 + float *v = ((float *)p->blob->p) + (i * p->blob->dimensions); + distances[i] = + distance_l2_sqr_float(v, (float *)queryVector, &p->blob->dimensions); + } + u8 *candidates = bitmap_new(bsize); + assert(candidates); + + u8 *taken = bitmap_new(bsize); + assert(taken); + + bitmap_fill(candidates, bsize); + for (size_t i = bsize; i >= p->blob->nvectors; i--) { + bitmap_set(candidates, i, 0); + } + i32 k_used = 0; + min_idx(distances, bsize, candidates, topk_rowids, k, taken, &k_used); + knn_data->current_idx = 0; + knn_data->distances = distances; + knn_data->k = k; + knn_data->rowids = topk_rowids; + + pCur->knn_data = knn_data; + } else { + pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN; + pCur->iRowid = 0; + } + + return SQLITE_OK; +} + +static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur, + sqlite_int64 *pRowid) { + vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; + switch (pCur->query_plan) { + case VEC_SBE__QUERYPLAN_FULLSCAN: { + *pRowid = pCur->iRowid; + return SQLITE_OK; + } + case VEC_SBE__QUERYPLAN_KNN: { + i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx]; + *pRowid = (sqlite3_int64)rowid; + return SQLITE_OK; + } + } + return SQLITE_ERROR; +} + +static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) { + vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; + switch (pCur->query_plan) { + case VEC_SBE__QUERYPLAN_FULLSCAN: { + pCur->iRowid++; + return SQLITE_OK; + } + case VEC_SBE__QUERYPLAN_KNN: { + pCur->knn_data->current_idx++; + return SQLITE_OK; + } + } + return SQLITE_ERROR; +} + +static int vec_static_blob_entriesEof(sqlite3_vtab_cursor *cur) { + vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; + vec_static_blob_entries_vtab *p = + (vec_static_blob_entries_vtab *)pCur->base.pVtab; + switch (pCur->query_plan) { + case VEC_SBE__QUERYPLAN_FULLSCAN: { + return (size_t)pCur->iRowid >= p->blob->nvectors; + } + case VEC_SBE__QUERYPLAN_KNN: { + return pCur->knn_data->current_idx >= pCur->knn_data->k; + } + } + return SQLITE_ERROR; +} + +static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur, + sqlite3_context *context, int i) { + vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; + vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)cur->pVtab; + + switch (pCur->query_plan) { + case VEC_SBE__QUERYPLAN_FULLSCAN: { + switch (i) { + case VEC_STATIC_BLOB_ENTRIES_VECTOR: + + sqlite3_result_blob( + context, + ((unsigned char *)p->blob->p) + + (pCur->iRowid * p->blob->dimensions * sizeof(float)), + p->blob->dimensions * sizeof(float), SQLITE_TRANSIENT); + sqlite3_result_subtype(context, p->blob->element_type); + break; + } + return SQLITE_OK; + } + case VEC_SBE__QUERYPLAN_KNN: { + switch (i) { + case VEC_STATIC_BLOB_ENTRIES_VECTOR: { + i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx]; + sqlite3_result_blob(context, + ((unsigned char *)p->blob->p) + + (rowid * p->blob->dimensions * sizeof(float)), + p->blob->dimensions * sizeof(float), + SQLITE_TRANSIENT); + sqlite3_result_subtype(context, p->blob->element_type); + break; + } + } + return SQLITE_OK; + } + } + return SQLITE_ERROR; +} + +static sqlite3_module vec_static_blob_entriesModule = { + /* iVersion */ 3, + /* xCreate */ + vec_static_blob_entriesCreate, // handle rm? + // https://github.com/asg017/sqlite-vec/issues/55 + /* xConnect */ vec_static_blob_entriesConnect, + /* xBestIndex */ vec_static_blob_entriesBestIndex, + /* xDisconnect */ vec_static_blob_entriesDisconnect, + /* xDestroy */ vec_static_blob_entriesDisconnect, + /* xOpen */ vec_static_blob_entriesOpen, + /* xClose */ vec_static_blob_entriesClose, + /* xFilter */ vec_static_blob_entriesFilter, + /* xNext */ vec_static_blob_entriesNext, + /* xEof */ vec_static_blob_entriesEof, + /* xColumn */ vec_static_blob_entriesColumn, + /* xRowid */ vec_static_blob_entriesRowid, + /* xUpdate */ 0, + /* xBegin */ 0, + /* xSync */ 0, + /* xCommit */ 0, + /* xRollback */ 0, + /* xFindMethod */ 0, + /* xRename */ 0, + /* xSavepoint */ 0, + /* xRelease */ 0, + /* xRollbackTo */ 0, + /* xShadowName */ 0, +#if SQLITE_VERSION_NUMBER >= 3044000 + /* xIntegrity */ 0 +#endif +}; +#pragma endregion + +#ifdef SQLITE_VEC_ENABLE_AVX +#define SQLITE_VEC_DEBUG_BUILD_AVX "avx" +#else +#define SQLITE_VEC_DEBUG_BUILD_AVX "" +#endif +#ifdef SQLITE_VEC_ENABLE_NEON +#define SQLITE_VEC_DEBUG_BUILD_NEON "neon" +#else +#define SQLITE_VEC_DEBUG_BUILD_NEON "" +#endif + +#define SQLITE_VEC_DEBUG_BUILD \ + SQLITE_VEC_DEBUG_BUILD_AVX " " SQLITE_VEC_DEBUG_BUILD_NEON + +#define SQLITE_VEC_DEBUG_STRING \ + "Version: " SQLITE_VEC_VERSION "\n" \ + "Date: " SQLITE_VEC_DATE "\n" \ + "Commit: " SQLITE_VEC_SOURCE "\n" \ + "Build flags: " SQLITE_VEC_DEBUG_BUILD + +SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg, + const sqlite3_api_routines *pApi) { +#ifndef SQLITE_CORE + SQLITE_EXTENSION_INIT2(pApi); +#endif + int rc = SQLITE_OK; + +#define DEFAULT_FLAGS (SQLITE_UTF8 | SQLITE_INNOCUOUS | SQLITE_DETERMINISTIC) + + rc = sqlite3_create_function_v2(db, "vec_version", 0, DEFAULT_FLAGS, + SQLITE_VEC_VERSION, _static_text_func, NULL, + NULL, NULL); + if (rc != SQLITE_OK) { + return rc; + } + rc = sqlite3_create_function_v2(db, "vec_debug", 0, DEFAULT_FLAGS, + SQLITE_VEC_DEBUG_STRING, _static_text_func, + NULL, NULL, NULL); + if (rc != SQLITE_OK) { + return rc; + } + static struct { + const char *zFName; + void (*xFunc)(sqlite3_context *, int, sqlite3_value **); + int nArg; + int flags; + } aFunc[] = { + // clang-format off + //{"vec_version", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_VERSION }, + //{"vec_debug", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_DEBUG_STRING }, + {"vec_distance_l2", vec_distance_l2, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, }, + {"vec_distance_l1", vec_distance_l1, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, }, + {"vec_distance_hamming",vec_distance_hamming, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, }, + {"vec_distance_cosine", vec_distance_cosine, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, }, + {"vec_length", vec_length, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE, }, + {"vec_type", vec_type, 1, DEFAULT_FLAGS, }, + {"vec_to_json", vec_to_json, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_add", vec_add, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_sub", vec_sub, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_slice", vec_slice, 3, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_normalize", vec_normalize, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_f32", vec_f32, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_bit", vec_bit, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_int8", vec_int8, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_quantize_int8", vec_quantize_int8, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + {"vec_quantize_binary", vec_quantize_binary, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, }, + // clang-format on + }; + + static struct { + char *name; + const sqlite3_module *module; + void *p; + void (*xDestroy)(void *); + } aMod[] = { + // clang-format off + {"vec0", &vec0Module, NULL, NULL}, + {"vec_each", &vec_eachModule, NULL, NULL}, + // clang-format on + }; + + for (unsigned long i = 0; i < countof(aFunc) && rc == SQLITE_OK; i++) { + rc = sqlite3_create_function_v2(db, aFunc[i].zFName, aFunc[i].nArg, + aFunc[i].flags, NULL, aFunc[i].xFunc, NULL, + NULL, NULL); + if (rc != SQLITE_OK) { + *pzErrMsg = sqlite3_mprintf("Error creating function %s: %s", + aFunc[i].zFName, sqlite3_errmsg(db)); + return rc; + } + } + + for (unsigned long i = 0; i < countof(aMod) && rc == SQLITE_OK; i++) { + rc = sqlite3_create_module_v2(db, aMod[i].name, aMod[i].module, NULL, NULL); + if (rc != SQLITE_OK) { + *pzErrMsg = sqlite3_mprintf("Error creating module %s: %s", aMod[i].name, + sqlite3_errmsg(db)); + return rc; + } + } + + return SQLITE_OK; +} + +#ifndef SQLITE_VEC_OMIT_FS +SQLITE_VEC_API int sqlite3_vec_numpy_init(sqlite3 *db, char **pzErrMsg, + const sqlite3_api_routines *pApi) { + UNUSED_PARAMETER(pzErrMsg); +#ifndef SQLITE_CORE + SQLITE_EXTENSION_INIT2(pApi); +#endif + int rc = SQLITE_OK; + rc = sqlite3_create_function_v2(db, "vec_npy_file", 1, SQLITE_RESULT_SUBTYPE, + NULL, vec_npy_file, NULL, NULL, NULL); + if(rc != SQLITE_OK) { + return rc; + } + rc = sqlite3_create_module_v2(db, "vec_npy_each", &vec_npy_eachModule, NULL, NULL); + return rc; +} +#endif + +SQLITE_VEC_API int +sqlite3_vec_static_blobs_init(sqlite3 *db, char **pzErrMsg, + const sqlite3_api_routines *pApi) { + UNUSED_PARAMETER(pzErrMsg); +#ifndef SQLITE_CORE + SQLITE_EXTENSION_INIT2(pApi); +#endif + + int rc = SQLITE_OK; + vec_static_blob_data *static_blob_data; + static_blob_data = sqlite3_malloc(sizeof(*static_blob_data)); + if (!static_blob_data) { + return SQLITE_NOMEM; + } + memset(static_blob_data, 0, sizeof(*static_blob_data)); + + rc = sqlite3_create_function_v2( + db, "vec_static_blob_from_raw", 4, + DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, NULL, + vec_static_blob_from_raw, NULL, NULL, NULL); + if (rc != SQLITE_OK) + return rc; + + rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule, + static_blob_data, sqlite3_free); + if (rc != SQLITE_OK) + return rc; + rc = sqlite3_create_module_v2(db, "vec_static_blob_entries", + &vec_static_blob_entriesModule, + static_blob_data, NULL); + if (rc != SQLITE_OK) + return rc; + return rc; +} diff --git a/internal/sqlitevec/sqlite-vec.h b/internal/sqlitevec/sqlite-vec.h new file mode 100644 index 00000000..daed5fcc --- /dev/null +++ b/internal/sqlitevec/sqlite-vec.h @@ -0,0 +1,38 @@ +#ifndef SQLITE_VEC_H +#define SQLITE_VEC_H + +#ifndef SQLITE_CORE +#include "sqlite3ext.h" +#else +#include "sqlite3.h" +#endif + +#ifdef SQLITE_VEC_STATIC +#define SQLITE_VEC_API +#else +#ifdef _WIN32 +#define SQLITE_VEC_API __declspec(dllexport) +#else +#define SQLITE_VEC_API +#endif +#endif + +#define SQLITE_VEC_VERSION "v0.1.9" +#define SQLITE_VEC_DATE "2026-03-31T07:59:06Z" +#define SQLITE_VEC_SOURCE "e9f598abfa0c06b328d8fe5da9c3760cce74be10" +#define SQLITE_VEC_VERSION_MAJOR 0 +#define SQLITE_VEC_VERSION_MINOR 1 +#define SQLITE_VEC_VERSION_PATCH 9 + +#ifdef __cplusplus +extern "C" { +#endif + +SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg, + const sqlite3_api_routines *pApi); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/internal/store/hybrid_cte_test.go b/internal/store/hybrid_cte_test.go index 0a474682..e5b54348 100644 --- a/internal/store/hybrid_cte_test.go +++ b/internal/store/hybrid_cte_test.go @@ -20,7 +20,7 @@ import ( "database/sql" "testing" - sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo" + sqlite_vec "github.com/ory/lumen/internal/sqlitevec" ) // TestHybridCTE_VecAndFTS5InCTE verifies that vec0 and FTS5 MATCH clauses diff --git a/internal/store/shared.go b/internal/store/shared.go new file mode 100644 index 00000000..a323fa73 --- /dev/null +++ b/internal/store/shared.go @@ -0,0 +1,894 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 + +package store + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "fmt" + "math" + "os" + "path/filepath" + "strings" + "time" + + "github.com/ory/lumen/internal/chunker" + sqlite_vec "github.com/ory/lumen/internal/sqlitevec" +) + +const sharedSchemaVersion = "1" + +// CollectionStats describes both project-local references and collection-wide +// physical storage. Chunk references may exceed UniqueVectors because vectors +// are content-addressed across file revisions and projects. +type CollectionStats struct { + ProjectFiles int + ProjectChunks int + UniqueVectors int + ChunkReferences int + SharedReferences int + DatabaseBytes int64 + ReclaimableBytes int64 + VectorStorage string +} + +// CleanupStats reports a shared collection garbage-collection pass. +type CleanupStats struct { + ProjectsRemoved int + VectorsRemoved int + BytesReclaimed int64 + ProjectsLeft int +} + +func openCollection(dsn string, dimensions int, vectorStorage string) (*Store, error) { + db, err := sql.Open("sqlite3", dsn) + if err != nil { + return nil, fmt.Errorf("open db: %w", err) + } + db.SetMaxOpenConns(1) + for _, pragma := range []string{ + "PRAGMA auto_vacuum=INCREMENTAL", + "PRAGMA journal_mode=WAL", + "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", + "PRAGMA cache_size=-64000", + "PRAGMA temp_store=MEMORY", + "PRAGMA busy_timeout=120000", + } { + if _, err := db.Exec(pragma); err != nil { + _ = db.Close() + return nil, fmt.Errorf("exec %q: %w", pragma, err) + } + } + // Legacy per-worktree databases remain readable during the lazy migration + // window. New profile paths always create the shared schema; an existing + // legacy path is upgraded by normal indexing without making it unreadable. + if legacy, _ := checkTableExists(db, "files"); legacy { + if shared, _ := checkTableExists(db, "collection_meta"); !shared { + _ = db.Close() + return openStore(dsn, dimensions) + } + } + if err := createCollectionSchema(db, dimensions, vectorStorage); err != nil { + _ = db.Close() + return nil, fmt.Errorf("create shared schema: %w", err) + } + + s := &Store{ + db: db, + dimensions: dimensions, + shared: true, + vectorStorage: vectorStorage, + dsn: dsn, + } + if dsn == ":memory:" { + return s, nil + } + readDB, err := sql.Open("sqlite3", dsn) + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("open read db: %w", err) + } + readDB.SetMaxOpenConns(1) + for _, pragma := range []string{ + "PRAGMA query_only=ON", + "PRAGMA foreign_keys=ON", + "PRAGMA cache_size=-64000", + "PRAGMA temp_store=MEMORY", + "PRAGMA busy_timeout=120000", + } { + if _, err := readDB.Exec(pragma); err != nil { + _ = readDB.Close() + _ = db.Close() + return nil, fmt.Errorf("read db %q: %w", pragma, err) + } + } + s.readDB = readDB + return s, nil +} + +func createCollectionSchema(db *sql.DB, dimensions int, vectorStorage string) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS collection_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + )`, + `CREATE TABLE IF NOT EXISTS projects ( + id INTEGER PRIMARY KEY, + path TEXT NOT NULL UNIQUE, + created_at TEXT NOT NULL, + last_accessed_at TEXT NOT NULL + )`, + `CREATE TABLE IF NOT EXISTS project_meta ( + project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + key TEXT NOT NULL, + value TEXT NOT NULL, + PRIMARY KEY(project_id, key) + ) WITHOUT ROWID`, + `CREATE TABLE IF NOT EXISTS file_revisions ( + id INTEGER PRIMARY KEY, + relative_path TEXT NOT NULL, + content_hash BLOB NOT NULL, + complete INTEGER NOT NULL DEFAULT 0, + UNIQUE(relative_path, content_hash) + )`, + `CREATE TABLE IF NOT EXISTS project_files ( + project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + relative_path TEXT NOT NULL, + file_revision_id INTEGER NOT NULL REFERENCES file_revisions(id), + PRIMARY KEY(project_id, relative_path) + ) WITHOUT ROWID`, + `CREATE INDEX IF NOT EXISTS idx_project_files_revision ON project_files(file_revision_id)`, + `CREATE TABLE IF NOT EXISTS vector_keys ( + id INTEGER PRIMARY KEY, + input_hash BLOB NOT NULL UNIQUE + )`, + `CREATE TABLE IF NOT EXISTS chunk_defs ( + id INTEGER PRIMARY KEY, + file_revision_id INTEGER NOT NULL REFERENCES file_revisions(id) ON DELETE CASCADE, + ordinal INTEGER NOT NULL, + chunk_key TEXT NOT NULL, + vector_id INTEGER NOT NULL REFERENCES vector_keys(id), + symbol TEXT NOT NULL, + kind TEXT NOT NULL, + start_line INTEGER NOT NULL, + end_line INTEGER NOT NULL, + UNIQUE(file_revision_id, ordinal) + )`, + `CREATE INDEX IF NOT EXISTS idx_chunk_defs_revision ON chunk_defs(file_revision_id)`, + `CREATE INDEX IF NOT EXISTS idx_chunk_defs_vector ON chunk_defs(vector_id)`, + } + for _, stmt := range stmts { + if _, err := db.Exec(stmt); err != nil { + return fmt.Errorf("exec %q: %w", stmt, err) + } + } + + want := map[string]string{ + "schema_version": sharedSchemaVersion, + "vec_dimensions": fmt.Sprintf("%d", dimensions), + "vector_storage": vectorStorage, + } + for key, value := range want { + var existing string + err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = ?`, key).Scan(&existing) + switch { + case err == sql.ErrNoRows: + if _, err := db.Exec(`INSERT INTO collection_meta(key, value) VALUES (?, ?)`, key, value); err != nil { + return err + } + case err != nil: + return err + case existing != value: + return fmt.Errorf("collection profile mismatch for %s: stored %q, requested %q", key, existing, value) + } + } + + exists, err := checkTableExists(db, "vec_vectors") + if err != nil { + return err + } + if !exists { + elementType := "int8" + if vectorStorage == "float32" { + elementType = "float" + } + stmt := fmt.Sprintf(`CREATE VIRTUAL TABLE vec_vectors USING vec0( + vector_id INTEGER PRIMARY KEY, + embedding %s[%d] distance_metric=cosine + )`, elementType, dimensions) + if _, err := db.Exec(stmt); err != nil { + return fmt.Errorf("create vec_vectors: %w", err) + } + } + return nil +} + +// IsShared reports whether this store uses the repository-scoped schema. +func (s *Store) IsShared() bool { return s.shared } + +// UseProject selects (and, if necessary, creates) a project membership in the +// shared collection. Calling it repeatedly for the same path is cheap. +func (s *Store) UseProject(projectPath string) error { + if !s.shared { + return nil + } + if projectPath == "" { + projectPath = ":default" + } else if abs, err := filepath.Abs(projectPath); err == nil { + projectPath = filepath.Clean(abs) + } + if s.projectID != 0 && s.projectPath == projectPath { + s.stampSharedAccess() + return nil + } + var existingID int64 + if err := s.reader().QueryRow(`SELECT id FROM projects WHERE path = ?`, projectPath).Scan(&existingID); err == nil { + s.projectID = existingID + s.projectPath = projectPath + s.stampSharedAccess() + return nil + } else if err != sql.ErrNoRows { + return fmt.Errorf("select project: %w", err) + } + now := nowRFC3339() + if _, err := s.db.Exec(`INSERT INTO projects(path, created_at, last_accessed_at) VALUES (?, ?, ?) + ON CONFLICT(path) DO UPDATE SET last_accessed_at = excluded.last_accessed_at`, projectPath, now, now); err != nil { + return fmt.Errorf("register project: %w", err) + } + if err := s.db.QueryRow(`SELECT id FROM projects WHERE path = ?`, projectPath).Scan(&s.projectID); err != nil { + return fmt.Errorf("select project: %w", err) + } + s.projectPath = projectPath + return nil +} + +// stampSharedAccess is best-effort and deliberately uses a short-lived +// connection with the same bounded timeout as legacy access stamping. Search +// must not wait behind a concurrent indexer's write transaction merely to +// refresh lifecycle metadata. +func (s *Store) stampSharedAccess() { + if s.dsn == "" || s.dsn == ":memory:" { + _, _ = s.db.Exec(`UPDATE projects SET last_accessed_at = ? WHERE id = ?`, nowRFC3339(), s.projectID) + return + } + db, err := sql.Open("sqlite3", s.dsn) + if err != nil { + return + } + defer func() { _ = db.Close() }() + db.SetMaxOpenConns(1) + if _, err := db.Exec(fmt.Sprintf("PRAGMA busy_timeout=%d", accessStampBusyTimeoutMS)); err != nil { + return + } + _, _ = db.Exec(`UPDATE projects SET last_accessed_at = ? WHERE id = ?`, nowRFC3339(), s.projectID) +} + +func nowRFC3339() string { return time.Now().UTC().Format(time.RFC3339) } + +func (s *Store) setProjectMeta(key, value string) error { + _, err := s.db.Exec(`INSERT INTO project_meta(project_id, key, value) VALUES (?, ?, ?) + ON CONFLICT(project_id, key) DO UPDATE SET value = excluded.value`, s.projectID, key, value) + return err +} + +func (s *Store) getProjectMeta(key string) (string, error) { + var value string + err := s.reader().QueryRow(`SELECT value FROM project_meta WHERE project_id = ? AND key = ?`, s.projectID, key).Scan(&value) + return value, err +} + +func (s *Store) getProjectMetaBatch(keys []string) (map[string]string, error) { + result := make(map[string]string, len(keys)) + if len(keys) == 0 { + return result, nil + } + marks := make([]string, len(keys)) + args := make([]any, 0, len(keys)+1) + args = append(args, s.projectID) + for i, key := range keys { + marks[i] = "?" + args = append(args, key) + } + rows, err := s.reader().Query(`SELECT key, value FROM project_meta WHERE project_id = ? AND key IN (`+strings.Join(marks, ",")+`)`, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + for rows.Next() { + var key, value string + if err := rows.Scan(&key, &value); err != nil { + return nil, err + } + result[key] = value + } + return result, rows.Err() +} + +func hashBlob(hash string) []byte { + if decoded, err := hex.DecodeString(hash); err == nil { + return decoded + } + return []byte(hash) +} + +// EmbeddingInput returns the exact, filepath-aware input whose hash identifies +// a vector in the shared collection. +func EmbeddingInput(c chunker.Chunk) string { + return "// " + c.FilePath + "\n" + c.Content +} + +func embeddingInputHash(c chunker.Chunk) [sha256.Size]byte { + return sha256.Sum256([]byte(EmbeddingInput(c))) +} + +// AttachExistingFileRevision links the selected project to an already indexed +// complete revision. It returns false when the revision has not been seen in +// this collection and must be chunked. +func (s *Store) AttachExistingFileRevision(relativePath, contentHash string) (bool, error) { + tx, err := s.db.Begin() + if err != nil { + return false, err + } + defer func() { _ = tx.Rollback() }() + var revisionID int64 + err = tx.QueryRow(`SELECT id FROM file_revisions WHERE relative_path = ? AND content_hash = ? AND complete = 1`, relativePath, hashBlob(contentHash)).Scan(&revisionID) + if err == sql.ErrNoRows { + return false, nil + } + if err != nil { + return false, err + } + if err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID); err != nil { + return false, err + } + if err := gcUnreferencedTx(tx); err != nil { + return false, err + } + return true, tx.Commit() +} + +// MissingChunkInputs returns the chunk positions whose exact embedding inputs +// are absent from the collection. Callers only need to embed these positions. +func (s *Store) MissingChunkInputs(chunks []chunker.Chunk) ([]int, error) { + missing := make([]int, 0, len(chunks)) + for i, c := range chunks { + h := embeddingInputHash(c) + var exists bool + if err := s.reader().QueryRow(`SELECT EXISTS(SELECT 1 FROM vector_keys WHERE input_hash = ?)`, h[:]).Scan(&exists); err != nil { + return nil, err + } + if !exists { + missing = append(missing, i) + } + } + return missing, nil +} + +// StoreFileRevision atomically installs a complete revision and updates the +// selected project's membership. vectors is keyed by chunk position and only +// needs entries returned by MissingChunkInputs; concurrent winners are reused. +func (s *Store) StoreFileRevision(relativePath, contentHash string, chunks []chunker.Chunk, vectors map[int][]float32) (bool, error) { + tx, err := s.db.Begin() + if err != nil { + return false, err + } + defer func() { _ = tx.Rollback() }() + + res, err := tx.Exec(`INSERT OR IGNORE INTO file_revisions(relative_path, content_hash, complete) VALUES (?, ?, 0)`, relativePath, hashBlob(contentHash)) + if err != nil { + return false, err + } + inserted, _ := res.RowsAffected() + var revisionID int64 + if err := tx.QueryRow(`SELECT id FROM file_revisions WHERE relative_path = ? AND content_hash = ?`, relativePath, hashBlob(contentHash)).Scan(&revisionID); err != nil { + return false, err + } + + if inserted > 0 { + for i, c := range chunks { + h := embeddingInputHash(c) + var vectorID int64 + err := tx.QueryRow(`SELECT id FROM vector_keys WHERE input_hash = ?`, h[:]).Scan(&vectorID) + if err == sql.ErrNoRows { + vec, ok := vectors[i] + if !ok { + return false, fmt.Errorf("missing embedding for chunk %d (%s)", i, c.ID) + } + result, err := tx.Exec(`INSERT INTO vector_keys(input_hash) VALUES (?)`, h[:]) + if err != nil { + return false, err + } + vectorID, err = result.LastInsertId() + if err != nil { + return false, err + } + blob, err := s.serializeVector(vec) + if err != nil { + return false, err + } + insertSQL := `INSERT INTO vec_vectors(vector_id, embedding) VALUES (?, ?)` + if s.vectorStorage == "int8" { + insertSQL = `INSERT INTO vec_vectors(vector_id, embedding) VALUES (?, vec_int8(?))` + } + if _, err := tx.Exec(insertSQL, vectorID, blob); err != nil { + return false, fmt.Errorf("insert shared vector: %w", err) + } + } else if err != nil { + return false, err + } + if _, err := tx.Exec(`INSERT INTO chunk_defs(file_revision_id, ordinal, chunk_key, vector_id, symbol, kind, start_line, end_line) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, revisionID, i, c.ID, vectorID, c.Symbol, c.Kind, c.StartLine, c.EndLine); err != nil { + return false, err + } + } + if _, err := tx.Exec(`UPDATE file_revisions SET complete = 1 WHERE id = ?`, revisionID); err != nil { + return false, err + } + } else { + var complete bool + if err := tx.QueryRow(`SELECT complete FROM file_revisions WHERE id = ?`, revisionID).Scan(&complete); err != nil { + return false, err + } + if !complete { + return false, fmt.Errorf("file revision %s is incomplete", relativePath) + } + // A force reindex supplies vectors even for an existing revision. Refresh + // those physical rows in place while preserving their stable numeric IDs + // and all sharing relationships. + for position, vec := range vectors { + if position < 0 || position >= len(chunks) { + return false, fmt.Errorf("vector position %d out of range", position) + } + h := embeddingInputHash(chunks[position]) + var vectorID int64 + if err := tx.QueryRow(`SELECT id FROM vector_keys WHERE input_hash = ?`, h[:]).Scan(&vectorID); err != nil { + return false, err + } + blob, err := s.serializeVector(vec) + if err != nil { + return false, err + } + if _, err := tx.Exec(`DELETE FROM vec_vectors WHERE vector_id = ?`, vectorID); err != nil { + return false, fmt.Errorf("remove shared vector for refresh: %w", err) + } + insertSQL := `INSERT INTO vec_vectors(vector_id, embedding) VALUES (?, ?)` + if s.vectorStorage == "int8" { + insertSQL = `INSERT INTO vec_vectors(vector_id, embedding) VALUES (?, vec_int8(?))` + } + if _, err := tx.Exec(insertSQL, vectorID, blob); err != nil { + return false, fmt.Errorf("refresh shared vector: %w", err) + } + } + } + + if err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID); err != nil { + return false, err + } + if err := gcUnreferencedTx(tx); err != nil { + return false, err + } + return inserted > 0, tx.Commit() +} + +func replaceProjectFileTx(tx *sql.Tx, projectID int64, path string, revisionID int64) error { + _, err := tx.Exec(`INSERT INTO project_files(project_id, relative_path, file_revision_id) VALUES (?, ?, ?) + ON CONFLICT(project_id, relative_path) DO UPDATE SET file_revision_id = excluded.file_revision_id`, projectID, path, revisionID) + return err +} + +func (s *Store) serializeVector(vector []float32) ([]byte, error) { + if len(vector) != s.dimensions { + return nil, fmt.Errorf("vector dimensions mismatch: got %d, want %d", len(vector), s.dimensions) + } + if s.vectorStorage == "float32" { + return sqlite_vec.SerializeFloat32(vector) + } + return quantizeInt8(vector), nil +} + +func quantizeInt8(vector []float32) []byte { + maxAbs := float32(0) + for _, value := range vector { + abs := float32(math.Abs(float64(value))) + if abs > maxAbs { + maxAbs = abs + } + } + result := make([]byte, len(vector)) + if maxAbs == 0 { + return result + } + for i, value := range vector { + q := int(math.Round(float64(value / maxAbs * 127))) + q = max(-127, min(127, q)) + result[i] = byte(int8(q)) + } + return result +} + +func (s *Store) upsertSharedFile(path, hash string) error { + if hash != "" { + if attached, err := s.AttachExistingFileRevision(path, hash); err != nil || attached { + return err + } + } + _, err := s.db.Exec(`INSERT OR IGNORE INTO file_revisions(relative_path, content_hash, complete) VALUES (?, ?, 0)`, path, hashBlob(hash)) + if err != nil { + return err + } + var revisionID int64 + if err := s.db.QueryRow(`SELECT id FROM file_revisions WHERE relative_path = ? AND content_hash = ?`, path, hashBlob(hash)).Scan(&revisionID); err != nil { + return err + } + return replaceProjectFileDB(s.db, s.projectID, path, revisionID) +} + +func replaceProjectFileDB(db *sql.DB, projectID int64, path string, revisionID int64) error { + _, err := db.Exec(`INSERT INTO project_files(project_id, relative_path, file_revision_id) VALUES (?, ?, ?) + ON CONFLICT(project_id, relative_path) DO UPDATE SET file_revision_id = excluded.file_revision_id`, projectID, path, revisionID) + return err +} + +func (s *Store) insertSharedChunks(chunks []chunker.Chunk, vectors [][]float32) error { + if len(chunks) != len(vectors) { + return fmt.Errorf("chunks and vectors length mismatch: %d vs %d", len(chunks), len(vectors)) + } + byFile := make(map[string][]int) + for i, c := range chunks { + byFile[c.FilePath] = append(byFile[c.FilePath], i) + } + for path, positions := range byFile { + var hash []byte + if err := s.db.QueryRow(`SELECT fr.content_hash FROM project_files pf JOIN file_revisions fr ON fr.id = pf.file_revision_id + WHERE pf.project_id = ? AND pf.relative_path = ?`, s.projectID, path).Scan(&hash); err != nil { + return err + } + fileChunks := make([]chunker.Chunk, len(positions)) + fileVectors := make(map[int][]float32, len(positions)) + for i, pos := range positions { + fileChunks[i] = chunks[pos] + fileVectors[i] = vectors[pos] + } + if _, err := s.StoreFileRevision(path, hex.EncodeToString(hash), fileChunks, fileVectors); err != nil { + return err + } + } + return nil +} + +func (s *Store) removeProjectFile(path string) error { + tx, err := s.db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + if _, err := tx.Exec(`DELETE FROM project_files WHERE project_id = ? AND relative_path = ?`, s.projectID, path); err != nil { + return err + } + if err := gcUnreferencedTx(tx); err != nil { + return err + } + return tx.Commit() +} + +func gcUnreferencedTx(tx *sql.Tx) error { + if _, err := tx.Exec(`DELETE FROM file_revisions WHERE NOT EXISTS ( + SELECT 1 FROM project_files WHERE file_revision_id = file_revisions.id)`); err != nil { + return err + } + rows, err := tx.Query(`SELECT id FROM vector_keys WHERE NOT EXISTS ( + SELECT 1 FROM chunk_defs WHERE vector_id = vector_keys.id)`) + if err != nil { + return err + } + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + _ = rows.Close() + return err + } + ids = append(ids, id) + } + if err := rows.Close(); err != nil { + return err + } + for _, id := range ids { + if _, err := tx.Exec(`DELETE FROM vec_vectors WHERE vector_id = ?`, id); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM vector_keys WHERE id = ?`, id); err != nil { + return err + } + } + return nil +} + +func (s *Store) projectFileHashes() (map[string]string, error) { + rows, err := s.reader().Query(`SELECT pf.relative_path, fr.content_hash + FROM project_files pf JOIN file_revisions fr ON fr.id = pf.file_revision_id + WHERE pf.project_id = ?`, s.projectID) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + result := make(map[string]string) + for rows.Next() { + var path string + var hash []byte + if err := rows.Scan(&path, &hash); err != nil { + return nil, err + } + result[path] = hex.EncodeToString(hash) + } + return result, rows.Err() +} + +func (s *Store) searchShared(ctx context.Context, queryVec []float32, limit int, maxDistance float64, pathPrefix string) ([]SearchResult, error) { + if limit <= 0 { + return nil, nil + } + blob, err := s.serializeVector(queryVec) + if err != nil { + return nil, fmt.Errorf("serialize query: %w", err) + } + var totalVectors int + if err := s.reader().QueryRow(`SELECT count(*) FROM vector_keys`).Scan(&totalVectors); err != nil { + return nil, err + } + if totalVectors == 0 { + return nil, nil + } + candidates := min(totalVectors, max(limit, 32)) + for { + results, boundaryDistance, err := s.searchSharedCandidates(ctx, blob, limit, candidates, maxDistance, pathPrefix) + if err != nil { + return nil, err + } + if candidates >= totalVectors || (len(results) >= limit && results[len(results)-1].Distance < boundaryDistance) { + return results, nil + } + candidates = min(totalVectors, candidates*2) + } +} + +func (s *Store) searchSharedCandidates(ctx context.Context, blob []byte, limit, candidates int, maxDistance float64, pathPrefix string) ([]SearchResult, float64, error) { + where := []string{"pf.project_id = ?"} + args := []any{blob, candidates, s.projectID} + if maxDistance > 0 { + where = append(where, "knn.distance < ?") + args = append(args, maxDistance) + } + if pathPrefix != "" { + where = append(where, "(pf.relative_path = ? OR pf.relative_path LIKE ? || '/%')") + args = append(args, pathPrefix, pathPrefix) + } + args = append(args, limit) + matchExpression := "?" + if s.vectorStorage == "int8" { + matchExpression = "vec_int8(?)" + } + query := `WITH knn AS ( + SELECT vector_id, distance FROM vec_vectors + WHERE embedding MATCH ` + matchExpression + ` AND k = ? + ) + SELECT pf.relative_path, cd.symbol, cd.kind, cd.start_line, cd.end_line, knn.distance, + max(knn.distance) OVER () + FROM knn + JOIN chunk_defs cd ON cd.vector_id = knn.vector_id + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE ` + strings.Join(where, " AND ") + ` + ORDER BY knn.distance, pf.relative_path, cd.start_line, cd.end_line, cd.symbol, cd.id + LIMIT ?` + rows, err := s.reader().QueryContext(ctx, query, args...) + if err != nil { + return nil, 0, fmt.Errorf("shared search query: %w", err) + } + defer func() { _ = rows.Close() }() + var results []SearchResult + var boundaryDistance float64 + for rows.Next() { + var result SearchResult + if err := rows.Scan(&result.FilePath, &result.Symbol, &result.Kind, &result.StartLine, &result.EndLine, &result.Distance, &boundaryDistance); err != nil { + return nil, 0, err + } + results = append(results, result) + } + if err := rows.Err(); err != nil { + return nil, 0, err + } + if err := rows.Close(); err != nil { + return nil, 0, err + } + return results, boundaryDistance, nil +} + +func (s *Store) sharedStats() (StoreStats, error) { + stats, err := s.CollectionStats() + if err != nil { + return StoreStats{}, err + } + return StoreStats{ + TotalFiles: stats.ProjectFiles, + TotalChunks: stats.ProjectChunks, + UniqueVectors: stats.UniqueVectors, + SharedReferences: stats.SharedReferences, + DatabaseBytes: stats.DatabaseBytes, + ReclaimableBytes: stats.ReclaimableBytes, + VectorStorage: stats.VectorStorage, + }, nil +} + +// CollectionStats returns project-local counts and physical collection usage. +func (s *Store) CollectionStats() (CollectionStats, error) { + var stats CollectionStats + err := s.reader().QueryRow(`SELECT + (SELECT count(*) FROM project_files WHERE project_id = ?), + (SELECT count(*) FROM chunk_defs cd JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id WHERE pf.project_id = ?), + (SELECT count(*) FROM vector_keys), + (SELECT count(*) FROM chunk_defs cd JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id)`, s.projectID, s.projectID).Scan( + &stats.ProjectFiles, &stats.ProjectChunks, &stats.UniqueVectors, &stats.ChunkReferences) + if err != nil { + return stats, err + } + stats.SharedReferences = stats.ChunkReferences - stats.UniqueVectors + stats.VectorStorage = s.vectorStorage + var pageCount, pageSize, freelist int64 + if err := s.reader().QueryRow(`PRAGMA page_count`).Scan(&pageCount); err != nil { + return stats, err + } + if err := s.reader().QueryRow(`PRAGMA page_size`).Scan(&pageSize); err != nil { + return stats, err + } + if err := s.reader().QueryRow(`PRAGMA freelist_count`).Scan(&freelist); err != nil { + return stats, err + } + stats.DatabaseBytes = pageCount * pageSize + stats.ReclaimableBytes = freelist * pageSize + return stats, nil +} + +func (s *Store) sharedTopSymbols(n int) ([]string, error) { + rows, err := s.reader().Query(`SELECT cd.symbol FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? GROUP BY cd.symbol ORDER BY count(*) DESC LIMIT ?`, s.projectID, n) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var symbols []string + for rows.Next() { + var symbol string + if err := rows.Scan(&symbol); err != nil { + return nil, err + } + symbols = append(symbols, symbol) + } + return symbols, rows.Err() +} + +func (s *Store) sharedHasSentinelFiles() (bool, error) { + var exists bool + err := s.reader().QueryRow(`SELECT EXISTS( + SELECT 1 FROM project_files pf JOIN file_revisions fr ON fr.id = pf.file_revision_id + WHERE pf.project_id = ? AND fr.complete = 0)`, s.projectID).Scan(&exists) + return exists, err +} + +// CleanupStaleProjects removes memberships not accessed since cutoff, garbage +// collects data no remaining project references, and incrementally reclaims +// free pages. Missing project directories are also considered stale. +func (s *Store) CleanupStaleProjects(cutoff time.Time) (CleanupStats, error) { + before, err := s.CollectionStats() + if err != nil { + return CleanupStats{}, err + } + rows, err := s.db.Query(`SELECT p.id, p.path, p.last_accessed_at, + NOT EXISTS(SELECT 1 FROM project_files pf WHERE pf.project_id = p.id) + AND NOT EXISTS(SELECT 1 FROM project_meta pm WHERE pm.project_id = p.id) + FROM projects p`) + if err != nil { + return CleanupStats{}, err + } + type project struct { + id int64 + path string + last string + unused bool + } + var stale []project + for rows.Next() { + var p project + if err := rows.Scan(&p.id, &p.path, &p.last, &p.unused); err != nil { + _ = rows.Close() + return CleanupStats{}, err + } + accessed, parseErr := time.Parse(time.RFC3339, p.last) + _, statErr := os.Stat(p.path) + defaultUnused := p.path == ":default" && p.unused + if defaultUnused || (p.path != ":default" && (os.IsNotExist(statErr) || parseErr != nil || accessed.Before(cutoff))) { + stale = append(stale, p) + } + } + if err := rows.Close(); err != nil { + return CleanupStats{}, err + } + tx, err := s.db.Begin() + if err != nil { + return CleanupStats{}, err + } + defer func() { _ = tx.Rollback() }() + for _, p := range stale { + if _, err := tx.Exec(`DELETE FROM projects WHERE id = ?`, p.id); err != nil { + return CleanupStats{}, err + } + } + if err := gcUnreferencedTx(tx); err != nil { + return CleanupStats{}, err + } + if err := tx.Commit(); err != nil { + return CleanupStats{}, err + } + after, err := s.CollectionStats() + if err != nil { + return CleanupStats{}, err + } + _, _ = s.db.Exec(`PRAGMA incremental_vacuum`) + if s.dsn != ":memory:" { + _, _ = s.db.Exec(`PRAGMA wal_checkpoint(PASSIVE)`) + } + final, _ := s.CollectionStats() + var projectsLeft int + _ = s.reader().QueryRow(`SELECT count(*) FROM projects`).Scan(&projectsLeft) + return CleanupStats{ + ProjectsRemoved: len(stale), + VectorsRemoved: before.UniqueVectors - after.UniqueVectors, + BytesReclaimed: max(int64(0), before.DatabaseBytes-final.DatabaseBytes), + ProjectsLeft: projectsLeft, + }, nil +} + +// CleanupCollectionAt opens dbPath only when it is a shared collection and +// performs project-aware cleanup. The boolean is false for legacy indexes. +func CleanupCollectionAt(dbPath string, cutoff time.Time) (CleanupStats, bool, error) { + db, err := sql.Open("sqlite3", dbPath) + if err != nil { + return CleanupStats{}, false, err + } + var shared bool + if err := db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'collection_meta')`).Scan(&shared); err != nil { + _ = db.Close() + return CleanupStats{}, false, err + } + if !shared { + _ = db.Close() + return CleanupStats{}, false, nil + } + var dimensions int + var storage string + if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vec_dimensions'`).Scan(&dimensions); err != nil { + _ = db.Close() + return CleanupStats{}, true, err + } + if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vector_storage'`).Scan(&storage); err != nil { + _ = db.Close() + return CleanupStats{}, true, err + } + _ = db.Close() + s, err := openCollection(dbPath, dimensions, storage) + if err != nil { + return CleanupStats{}, true, err + } + defer func() { _ = s.Close() }() + stats, err := s.CleanupStaleProjects(cutoff) + return stats, true, err +} diff --git a/internal/store/shared_test.go b/internal/store/shared_test.go new file mode 100644 index 00000000..d98a8373 --- /dev/null +++ b/internal/store/shared_test.go @@ -0,0 +1,379 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +package store + +import ( + "context" + "math/rand" + "os" + "path/filepath" + "strconv" + "sync" + "testing" + "time" + + "github.com/ory/lumen/internal/chunker" +) + +func TestSharedCollectionReusesRevisionsAndVectors(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "index.db") + projectA := filepath.Join(t.TempDir(), "worktree-a") + projectB := filepath.Join(t.TempDir(), "worktree-b") + s, err := NewCollection(dbPath, 4, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + + chunks := []chunker.Chunk{ + {ID: "a", FilePath: "main.go", Symbol: "A", Kind: "function", StartLine: 1, EndLine: 2, Content: "func A() {}"}, + {ID: "b", FilePath: "main.go", Symbol: "B", Kind: "function", StartLine: 3, EndLine: 4, Content: "func B() {}"}, + } + vectors := map[int][]float32{ + 0: {1, 0, 0, 0}, + 1: {0, 1, 0, 0}, + } + created, err := s.StoreFileRevision("main.go", "abcd", chunks, vectors) + if err != nil { + t.Fatal(err) + } + if !created { + t.Fatal("first revision should be created") + } + + if err := s.UseProject(projectB); err != nil { + t.Fatal(err) + } + reused, err := s.AttachExistingFileRevision("main.go", "abcd") + if err != nil { + t.Fatal(err) + } + if !reused { + t.Fatal("second worktree should reuse the complete revision") + } + stats, err := s.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 2 || stats.ChunkReferences != 4 || stats.SharedReferences != 2 { + t.Fatalf("unexpected dedup stats: %+v", stats) + } + + results, err := s.Search(context.Background(), []float32{1, 0, 0, 0}, 2, 0, "") + if err != nil { + t.Fatal(err) + } + if len(results) != 2 || results[0].Symbol != "A" { + t.Fatalf("unexpected project-local search results: %+v", results) + } + + if err := s.UseProject(projectA); err != nil { + t.Fatal(err) + } + if err := s.DeleteFileChunks("main.go"); err != nil { + t.Fatal(err) + } + stats, err = s.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 2 { + t.Fatalf("shared vectors removed while another project referenced them: %+v", stats) + } + + if err := s.UseProject(projectB); err != nil { + t.Fatal(err) + } + if err := s.DeleteFileChunks("main.go"); err != nil { + t.Fatal(err) + } + stats, err = s.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 0 || stats.ChunkReferences != 0 { + t.Fatalf("last-reference GC left vectors behind: %+v", stats) + } +} + +func TestSharedCollectionFloat32Override(t *testing.T) { + s, err := NewCollection(":memory:", 3, "float32", t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + chunk := chunker.Chunk{ID: "c", FilePath: "x.go", Symbol: "X", Kind: "function", StartLine: 1, EndLine: 1, Content: "func X() {}"} + if _, err := s.StoreFileRevision("x.go", "01", []chunker.Chunk{chunk}, map[int][]float32{0: {0.1, 0.2, 0.3}}); err != nil { + t.Fatal(err) + } + results, err := s.Search(context.Background(), []float32{0.1, 0.2, 0.3}, 1, 0, "") + if err != nil { + t.Fatal(err) + } + if len(results) != 1 || results[0].Symbol != "X" { + t.Fatalf("unexpected float32 results: %+v", results) + } +} + +func TestSharedSearchAdaptivelyExpandsSparseProjectCandidates(t *testing.T) { + s, err := NewCollection(":memory:", 4, "int8", "/project-a") + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + + chunks := make([]chunker.Chunk, 80) + vectors := make(map[int][]float32, len(chunks)) + for i := range chunks { + suffix := strconv.Itoa(i + 1) + chunks[i] = chunker.Chunk{ + ID: "a" + suffix, FilePath: "crowded.go", Symbol: "Crowded", Kind: "function", + StartLine: i + 1, EndLine: i + 1, Content: "crowded input " + suffix, + } + vectors[i] = []float32{1, float32(i+1) / 10000, 0, 0} + } + if _, err := s.StoreFileRevision("crowded.go", "aa", chunks, vectors); err != nil { + t.Fatal(err) + } + + if err := s.UseProject("/project-b"); err != nil { + t.Fatal(err) + } + target := chunker.Chunk{ID: "target", FilePath: "nested/target.go", Symbol: "Target", Kind: "function", StartLine: 1, EndLine: 1, Content: "target input"} + if _, err := s.StoreFileRevision("nested/target.go", "bb", []chunker.Chunk{target}, map[int][]float32{0: {0, 1, 0, 0}}); err != nil { + t.Fatal(err) + } + + // The nearest 80 global vectors belong to project A. Project B's result + // is only found after candidate doubling exhausts that dense prefix. + results, err := s.Search(context.Background(), []float32{1, 0, 0, 0}, 1, 0, "nested") + if err != nil { + t.Fatal(err) + } + if len(results) != 1 || results[0].Symbol != "Target" { + t.Fatalf("adaptive sparse-project search failed: %+v", results) + } +} + +func TestSharedCollectionConcurrentRevisionInsertionIsIdempotent(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "index.db") + projectA, projectB := t.TempDir(), t.TempDir() + a, err := NewCollection(dbPath, 4, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = a.Close() }() + b, err := NewCollection(dbPath, 4, "int8", projectB) + if err != nil { + t.Fatal(err) + } + defer func() { _ = b.Close() }() + chunk := chunker.Chunk{ID: "same", FilePath: "same.go", Symbol: "Same", Kind: "function", StartLine: 1, EndLine: 1, Content: "func Same() {}"} + + start := make(chan struct{}) + errs := make(chan error, 2) + var wg sync.WaitGroup + for _, collection := range []*Store{a, b} { + wg.Add(1) + go func(s *Store) { + defer wg.Done() + <-start + _, err := s.StoreFileRevision("same.go", "cc", []chunker.Chunk{chunk}, map[int][]float32{0: {1, 0, 0, 0}}) + errs <- err + }(collection) + } + close(start) + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + stats, err := a.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 1 || stats.ChunkReferences != 2 { + t.Fatalf("concurrent insertion was not idempotent: %+v", stats) + } +} + +func TestSharedCleanupRemovesOnlyStaleMemberships(t *testing.T) { + projectA, projectB := t.TempDir(), t.TempDir() + s, err := NewCollection(":memory:", 4, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + chunk := chunker.Chunk{ID: "shared", FilePath: "shared.go", Symbol: "Shared", Kind: "function", StartLine: 1, EndLine: 1, Content: "func Shared() {}"} + if _, err := s.StoreFileRevision("shared.go", "dd", []chunker.Chunk{chunk}, map[int][]float32{0: {1, 0, 0, 0}}); err != nil { + t.Fatal(err) + } + if err := s.UseProject(projectB); err != nil { + t.Fatal(err) + } + if attached, err := s.AttachExistingFileRevision("shared.go", "dd"); err != nil || !attached { + t.Fatalf("attach second project: attached=%v err=%v", attached, err) + } + old := time.Now().Add(-60 * 24 * time.Hour).UTC().Format(time.RFC3339) + if _, err := s.db.Exec(`UPDATE projects SET last_accessed_at = ? WHERE path = ?`, old, projectA); err != nil { + t.Fatal(err) + } + cleanup, err := s.CleanupStaleProjects(time.Now().Add(-30 * 24 * time.Hour)) + if err != nil { + t.Fatal(err) + } + if cleanup.ProjectsRemoved != 1 || cleanup.VectorsRemoved != 0 { + t.Fatalf("unexpected cleanup stats: %+v", cleanup) + } + stats, err := s.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 1 || stats.ChunkReferences != 1 { + t.Fatalf("active project's shared data was not preserved: %+v", stats) + } +} + +func TestInt8RecallAt8AgainstFloat32(t *testing.T) { + const ( + dimensions = 64 + vectorCount = 200 + queryCount = 20 + limit = 8 + ) + project := t.TempDir() + floatStore, err := NewCollection(":memory:", dimensions, "float32", project) + if err != nil { + t.Fatal(err) + } + defer func() { _ = floatStore.Close() }() + int8Store, err := NewCollection(":memory:", dimensions, "int8", project) + if err != nil { + t.Fatal(err) + } + defer func() { _ = int8Store.Close() }() + + rng := rand.New(rand.NewSource(42)) //nolint:gosec // deterministic recall fixture + chunks := make([]chunker.Chunk, vectorCount) + vectors := make(map[int][]float32, vectorCount) + for i := range vectorCount { + suffix := strconv.Itoa(i) + chunks[i] = chunker.Chunk{ID: suffix, FilePath: "vectors.go", Symbol: "V" + suffix, Kind: "function", StartLine: i + 1, EndLine: i + 1, Content: "vector " + suffix} + vector := make([]float32, dimensions) + for j := range vector { + vector[j] = float32(rng.NormFloat64()) + } + vectors[i] = vector + } + for _, s := range []*Store{floatStore, int8Store} { + if _, err := s.StoreFileRevision("vectors.go", "ee", chunks, vectors); err != nil { + t.Fatal(err) + } + } + + matches := 0 + for range queryCount { + query := make([]float32, dimensions) + for j := range query { + query[j] = float32(rng.NormFloat64()) + } + want, err := floatStore.Search(context.Background(), query, limit, 0, "") + if err != nil { + t.Fatal(err) + } + got, err := int8Store.Search(context.Background(), query, limit, 0, "") + if err != nil { + t.Fatal(err) + } + wantSymbols := make(map[string]struct{}, len(want)) + for _, result := range want { + wantSymbols[result.Symbol] = struct{}{} + } + for _, result := range got { + if _, ok := wantSymbols[result.Symbol]; ok { + matches++ + } + } + } + recall := float64(matches) / float64(queryCount*limit) + if recall < 0.95 { + t.Fatalf("int8 recall@8 = %.3f, want >= 0.95", recall) + } +} + +func TestSharedInt8StorageAtMostTwentyPercentOfSeparateFloat32(t *testing.T) { + const ( + dimensions = 768 + chunkCount = 1000 + ) + base := t.TempDir() + projectA, projectB := t.TempDir(), t.TempDir() + chunks := make([]chunker.Chunk, chunkCount) + vectorMap := make(map[int][]float32, chunkCount) + vectorSlice := make([][]float32, chunkCount) + rng := rand.New(rand.NewSource(7)) //nolint:gosec // deterministic size fixture + for i := range chunkCount { + suffix := strconv.Itoa(i) + chunks[i] = chunker.Chunk{ID: suffix, FilePath: "bulk.go", Symbol: "Bulk" + suffix, Kind: "function", StartLine: i + 1, EndLine: i + 1, Content: "bulk " + suffix} + vector := make([]float32, dimensions) + for j := range vector { + vector[j] = float32(rng.NormFloat64()) + } + vectorMap[i] = vector + vectorSlice[i] = vector + } + + sharedPath := filepath.Join(base, "shared.db") + shared, err := NewCollection(sharedPath, dimensions, "int8", projectA) + if err != nil { + t.Fatal(err) + } + if _, err := shared.StoreFileRevision("bulk.go", "ff", chunks, vectorMap); err != nil { + t.Fatal(err) + } + if err := shared.UseProject(projectB); err != nil { + t.Fatal(err) + } + if attached, err := shared.AttachExistingFileRevision("bulk.go", "ff"); err != nil || !attached { + t.Fatalf("attach shared revision: attached=%v err=%v", attached, err) + } + if err := shared.Close(); err != nil { + t.Fatal(err) + } + + var separateBytes int64 + for i := range 2 { + path := filepath.Join(base, "float-"+strconv.Itoa(i)+".db") + legacy, err := New(path, dimensions) + if err != nil { + t.Fatal(err) + } + if err := legacy.UpsertFile("bulk.go", "ff"); err != nil { + t.Fatal(err) + } + if err := legacy.InsertChunks(chunks, vectorSlice); err != nil { + t.Fatal(err) + } + if err := legacy.Close(); err != nil { + t.Fatal(err) + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + separateBytes += info.Size() + } + sharedInfo, err := os.Stat(sharedPath) + if err != nil { + t.Fatal(err) + } + ratio := float64(sharedInfo.Size()) / float64(separateBytes) + if ratio > 0.20 { + t.Fatalf("shared int8 size ratio = %.3f (%d/%d), want <= 0.20", ratio, sharedInfo.Size(), separateBytes) + } +} diff --git a/internal/store/store.go b/internal/store/store.go index 5e07114b..796f047d 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -23,10 +23,10 @@ import ( "strings" "time" - sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo" _ "github.com/mattn/go-sqlite3" // register sqlite3 driver "github.com/ory/lumen/internal/chunker" + sqlite_vec "github.com/ory/lumen/internal/sqlitevec" ) // MetaLastAccessedAt is the project_meta key holding the RFC3339 UTC timestamp @@ -74,16 +74,26 @@ type SearchResult struct { // StoreStats holds aggregate statistics about the store contents. type StoreStats struct { //nolint:revive // StoreStats is intentionally named to avoid ambiguity at call sites - TotalFiles int - TotalChunks int + TotalFiles int + TotalChunks int + UniqueVectors int + SharedReferences int + DatabaseBytes int64 + ReclaimableBytes int64 + VectorStorage string } // Store manages SQLite + sqlite-vec storage for code chunks and their // embedding vectors. type Store struct { - db *sql.DB - readDB *sql.DB // separate read-only connection; nil for :memory: databases - dimensions int + db *sql.DB + readDB *sql.DB // separate read-only connection; nil for :memory: databases + dimensions int + shared bool + vectorStorage string + projectPath string + projectID int64 + dsn string } // New opens (or creates) a SQLite database at dsn, enables WAL mode and @@ -103,7 +113,40 @@ func New(dsn string, dimensions int) (*Store, error) { if err != nil { return s, err } - s.stampAccess(dsn) + if s.shared && s.projectID == 0 { + if err := s.UseProject(""); err != nil { + _ = s.Close() + return nil, err + } + } + if s.shared { + s.stampSharedAccess() + } else { + s.stampAccess(dsn) + } + return s, nil +} + +// NewCollection opens a repository-scoped shared collection and selects the +// project membership identified by projectPath. vectorStorage must be int8 or +// float32. Multiple Store instances may safely select different worktrees in +// the same database. +func NewCollection(dsn string, dimensions int, vectorStorage, projectPath string) (*Store, error) { + if vectorStorage != "int8" && vectorStorage != "float32" { + return nil, fmt.Errorf("unsupported vector storage %q", vectorStorage) + } + s, err := openCollection(dsn, dimensions, vectorStorage) + if err != nil && IsCorruptionErr(err) && dsn != ":memory:" { + deleteDBFiles(dsn) + s, err = openCollection(dsn, dimensions, vectorStorage) + } + if err != nil { + return nil, err + } + if err := s.UseProject(projectPath); err != nil { + _ = s.Close() + return nil, err + } return s, nil } @@ -165,6 +208,18 @@ func openStore(dsn string, dimensions int) (*Store, error) { return nil, fmt.Errorf("exec %q: %w", p, err) } } + // A low-level caller may open a database already upgraded to the shared + // schema (for example metadata tooling during a rolling upgrade). Return a + // shared view instead of attempting to overlay the legacy tables. + if shared, _ := checkTableExists(db, "collection_meta"); shared { + var storage string + if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vector_storage'`).Scan(&storage); err != nil { + _ = db.Close() + return nil, err + } + _ = db.Close() + return openCollection(dsn, dimensions, storage) + } if err := createSchema(db, dimensions); err != nil { _ = db.Close() @@ -372,6 +427,9 @@ func ReadMetaAt(dbPath string, keys ...string) (map[string]string, error) { // SetMeta upserts a key-value pair in the project_meta table. func (s *Store) SetMeta(key, value string) error { + if s.shared { + return s.setProjectMeta(key, value) + } _, err := s.db.Exec( `INSERT INTO project_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, @@ -383,6 +441,9 @@ func (s *Store) SetMeta(key, value string) error { // GetMeta retrieves a value from the project_meta table by key. // It uses the read-only connection when available for concurrency with writes. func (s *Store) GetMeta(key string) (string, error) { + if s.shared { + return s.getProjectMeta(key) + } var val string err := s.reader().QueryRow("SELECT value FROM project_meta WHERE key = ?", key).Scan(&val) if err != nil { @@ -395,6 +456,9 @@ func (s *Store) GetMeta(key string) (string, error) { // Missing keys are absent from the returned map. Uses the read-only connection // when available for concurrency with writes. func (s *Store) GetMetaBatch(keys []string) (map[string]string, error) { + if s.shared { + return s.getProjectMetaBatch(keys) + } return queryMeta(s.reader(), keys) } @@ -437,6 +501,9 @@ func queryMeta(db *sql.DB, keys []string) (map[string]string, error) { // UpsertFile inserts or updates a file path and its content hash. func (s *Store) UpsertFile(path, hash string) error { + if s.shared { + return s.upsertSharedFile(path, hash) + } _, err := s.db.Exec( `INSERT INTO files (path, hash) VALUES (?, ?) ON CONFLICT(path) DO UPDATE SET hash = excluded.hash`, @@ -453,6 +520,9 @@ func (s *Store) UpsertFile(path, hash string) error { // would cause an error. The deduplication loop below handles within-batch // duplicates only. func (s *Store) InsertChunks(chunks []chunker.Chunk, vectors [][]float32) error { + if s.shared { + return s.insertSharedChunks(chunks, vectors) + } if len(chunks) != len(vectors) { return fmt.Errorf("chunks and vectors length mismatch: %d vs %d", len(chunks), len(vectors)) } @@ -527,6 +597,9 @@ func insertChunkAndVector(chunkStmt, vecStmt interface { // DeleteFileChunks removes all chunks (and their vectors) associated with the // given file path, then removes the file record itself. func (s *Store) DeleteFileChunks(filePath string) error { + if s.shared { + return s.removeProjectFile(filePath) + } tx, err := s.db.Begin() if err != nil { return fmt.Errorf("begin tx: %w", err) @@ -568,6 +641,9 @@ func (s *Store) DeleteFileChunks(filePath string) error { // connection (e.g. during indexing). The provided context is used for // query cancellation. func (s *Store) Search(ctx context.Context, queryVec []float32, limit int, maxDistance float64, pathPrefix string) ([]SearchResult, error) { + if s.shared { + return s.searchShared(ctx, queryVec, limit, maxDistance, pathPrefix) + } blob, err := sqlite_vec.SerializeFloat32(queryVec) if err != nil { return nil, fmt.Errorf("serialize query: %w", err) @@ -626,6 +702,9 @@ func (s *Store) Search(ctx context.Context, queryVec []float32, limit int, maxDi // GetFileHashes returns a map of file path to content hash for all tracked files. func (s *Store) GetFileHashes() (map[string]string, error) { + if s.shared { + return s.projectFileHashes() + } rows, err := s.db.Query("SELECT path, hash FROM files") if err != nil { return nil, fmt.Errorf("query files: %w", err) @@ -646,6 +725,9 @@ func (s *Store) GetFileHashes() (map[string]string, error) { // Stats returns aggregate statistics about the store contents in one query. // Uses the read-only connection when available for concurrency with writes. func (s *Store) Stats() (StoreStats, error) { + if s.shared { + return s.sharedStats() + } var stats StoreStats err := s.reader().QueryRow( `SELECT (SELECT count(*) FROM files), (SELECT count(*) FROM chunks)`, @@ -658,6 +740,9 @@ func (s *Store) Stats() (StoreStats, error) { // TopSymbols returns the n most frequently occurring symbol names in the store. func (s *Store) TopSymbols(n int) ([]string, error) { + if s.shared { + return s.sharedTopSymbols(n) + } rows, err := s.reader().Query( "SELECT symbol FROM chunks GROUP BY symbol ORDER BY count(*) DESC LIMIT ?", n, ) @@ -682,6 +767,9 @@ func (s *Store) TopSymbols(n int) ([]string, error) { // Uses the read-only connection when available for consistency with other // read methods (GetMeta, Stats, Search). func (s *Store) HasSentinelFiles() (bool, error) { + if s.shared { + return s.sharedHasSentinelFiles() + } var exists bool err := s.reader().QueryRow("SELECT EXISTS(SELECT 1 FROM files WHERE hash = '')").Scan(&exists) return exists, err diff --git a/skills/doctor/SKILL.md b/skills/doctor/SKILL.md index 91cee630..5453ca42 100644 --- a/skills/doctor/SKILL.md +++ b/skills/doctor/SKILL.md @@ -17,8 +17,13 @@ project. 3. Report a concise summary: - Embedding service status, backend, host, and model - Index totals: files, chunks, last indexed time, stale or fresh + - Storage totals: unique vectors, shared references, deduplication ratio, + vector precision, database bytes, and reclaimable bytes - Any MCP or plugin setup issue that blocks the tools 4. If no index exists yet, explain that the Lumen `semantic_search` tool seeds the index on first use. 5. If the user wants eager indexing instead of waiting for the next search, suggest running `lumen index .` in the repository root. +6. Interpret storage totals as collection-wide values shared by compatible Git + worktrees; file, chunk, freshness, and indexing values are project-local. + A high deduplication ratio is expected and is not an index-health problem. diff --git a/skills/reindex/SKILL.md b/skills/reindex/SKILL.md index e9c534c6..2f9988e1 100644 --- a/skills/reindex/SKILL.md +++ b/skills/reindex/SKILL.md @@ -17,12 +17,13 @@ Refresh or rebuild the bundled Lumen index for the current project. missing indexes automatically. 3. If the user explicitly asks for a clean rebuild, explain the options and run one via the shell: - - `lumen index --force .` — rebuilds only the current project's index from - scratch. Prefer this. - - `lumen clean --days 0 && lumen index .` — deletes every eligible cached - index except those protected by active locks, then rebuilds. Use only when - the user asks for a full wipe. + - `lumen index --force .` — reprocesses every file for the current project + without wiping other worktrees or their shared vectors. Prefer this. + - `lumen clean --days 0 && lumen index .` — deletes every cached index on the + host before rebuilding. Use only when the user asks for a full wipe. - `lumen clean` — removes indexes for projects that no longer exist or have not been used in 30 days. Use to reclaim disk space, not to rebuild the current project. -4. After the refresh or rebuild, report the new index status. +4. After the refresh or rebuild, report the new index status, including vector + precision, unique vectors, shared references, deduplication ratio, database + bytes, and reclaimable bytes when available. From 7a0ba24a8ded68c07ceaa3ea63e07007a76d6911 Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:29:04 +0200 Subject: [PATCH 2/9] fix(index): harden shared index lifecycle --- cmd/clean_test.go | 50 +++++++++++++ cmd/index.go | 8 +- cmd/index_test.go | 28 +++++++ cmd/stdio_test.go | 24 +++--- internal/index/index.go | 94 +++++++++++++++++++----- internal/index/index_concurrency_test.go | 68 +++++++++++++++++ 6 files changed, 235 insertions(+), 37 deletions(-) diff --git a/cmd/clean_test.go b/cmd/clean_test.go index 904ad9b6..b0aaee1b 100644 --- a/cmd/clean_test.go +++ b/cmd/clean_test.go @@ -445,6 +445,56 @@ func TestClean_HoldsLockDuringRemovalAndReleasesItAfterFailure(t *testing.T) { lock.Release() } +func TestClean_HoldsCollectionLockThroughSharedCleanupAndRemoval(t *testing.T) { + tmp := resolvedTempDir(t) + t.Setenv("XDG_DATA_HOME", tmp) + + project := projectDir(t, "stale-shared") + dbPath := config.DBPathForProjectProfile(project, embedder.DefaultModel, 4, "int8", 512) + require.NoError(t, os.MkdirAll(filepath.Dir(dbPath), 0o755)) + s, err := store.NewCollection(dbPath, 4, "int8", project) + require.NoError(t, err) + require.NoError(t, s.Close()) + db, err := sql.Open("sqlite3", dbPath) + require.NoError(t, err) + _, err = db.Exec(`UPDATE projects SET last_accessed_at = ? WHERE path = ?`, daysAgo(45), project) + require.NoError(t, err) + require.NoError(t, db.Close()) + + hashDir := filepath.Dir(dbPath) + lockPath := indexlock.LockPathForDB(dbPath) + removeErr := errors.New("injected removal failure") + cleanupLockHeld := false + removeLockHeld := false + originalCleanup := cleanupCollectionAt + originalRemove := removeIndexDir + cleanupCollectionAt = func(path string, cutoff time.Time) (store.CleanupStats, bool, error) { + assert.Equal(t, dbPath, path) + cleanupLockHeld = indexlock.IsAnyHeld(lockPath) + return originalCleanup(path, cutoff) + } + removeIndexDir = func(path string) error { + assert.Equal(t, hashDir, path) + removeLockHeld = indexlock.IsAnyHeld(lockPath) + return removeErr + } + t.Cleanup(func() { + cleanupCollectionAt = originalCleanup + removeIndexDir = originalRemove + }) + + _, _, err = runCleanIndexes(t, tmp, 30) + require.ErrorIs(t, err, removeErr) + assert.True(t, cleanupLockHeld, "cleanup must hold the collection lock while mutating the database") + assert.True(t, removeLockHeld, "cleanup must retain the collection lock while removing the directory") + assert.DirExists(t, hashDir, "the injected removal failure must leave the collection directory") + + lock, lockErr := indexlock.TryAcquire(lockPath) + require.NoError(t, lockErr) + require.NotNil(t, lock, "cleanup must release the collection lock after a failure") + lock.Release() +} + func TestClean_RejectsPositionalArgs(t *testing.T) { require.Error(t, cleanCmd.Args(cleanCmd, []string{"/some/project"}), "clean takes no positional arguments") diff --git a/cmd/index.go b/cmd/index.go index 2da800bc..8241dcc2 100644 --- a/cmd/index.go +++ b/cmd/index.go @@ -225,13 +225,7 @@ func setupIndexer(cfg *config.ConfigService, emb *embedder.FailoverEmbedder, dbP } func configuredDBPath(cfg *config.ConfigService, projectPath, model string) string { - dimensions, known := config.ModelDimensions(model) - if !known { - servers := cfg.Servers() - if len(servers) > 0 && servers[0].Model == model { - dimensions = cfg.ServerDims(0) - } - } + dimensions := cfg.ServerDims(0) return config.DBPathForProjectProfile(projectPath, model, dimensions, cfg.VectorStorage(), cfg.MaxChunkTokens()) } diff --git a/cmd/index_test.go b/cmd/index_test.go index cd18c9c0..56610b68 100644 --- a/cmd/index_test.go +++ b/cmd/index_test.go @@ -19,6 +19,8 @@ import ( "path/filepath" "strings" "testing" + + "github.com/ory/lumen/internal/config" ) func TestRunIndex_RefusesUnindexableRoot(t *testing.T) { @@ -35,3 +37,29 @@ func TestRunIndex_RefusesUnindexableRoot(t *testing.T) { t.Fatalf("expected error to mention the .lumenignore catch-all reason, got %q", err.Error()) } } + +func TestConfiguredDBPathUsesExplicitDimensionsForKnownModel(t *testing.T) { + t.Setenv("XDG_DATA_HOME", t.TempDir()) + const model = "ordis/jina-embeddings-v2-base-code" + cmd := flagsTestHarness(t, ` +servers: + - backend: ollama + host: http://localhost:11434 + model: ordis/jina-embeddings-v2-base-code + dims: 1024 +`) + cfg, err := loadConfigWithFlags(cmd) + if err != nil { + t.Fatal(err) + } + project := t.TempDir() + got := configuredDBPath(cfg, project, model) + want := config.DBPathForProjectProfile(project, model, 1024, cfg.VectorStorage(), cfg.MaxChunkTokens()) + registryPath := config.DBPathForProjectProfile(project, model, 768, cfg.VectorStorage(), cfg.MaxChunkTokens()) + if got != want { + t.Fatalf("configuredDBPath = %q, want explicit-dimension path %q", got, want) + } + if got == registryPath { + t.Fatal("explicit dimensions must not reuse the built-in model profile") + } +} diff --git a/cmd/stdio_test.go b/cmd/stdio_test.go index ea9104be..e7469e67 100644 --- a/cmd/stdio_test.go +++ b/cmd/stdio_test.go @@ -399,10 +399,10 @@ func TestIndexerCache_GetOrCreate_ModelChangeCreatesSeparateIndexer(t *testing.T t.Fatalf("expected same effective root, got %q vs %q", rootA, rootB) } - if _, err := os.Stat(config.DBPathForProject(rootA, "model-a")); err != nil { + if _, err := os.Stat(ic.dbPath(rootA, "model-a")); err != nil { t.Fatalf("expected model-a DB to exist: %v", err) } - if _, err := os.Stat(config.DBPathForProject(rootB, "model-b")); err != nil { + if _, err := os.Stat(ic.dbPath(rootB, "model-b")); err != nil { t.Fatalf("expected model-b DB to exist: %v", err) } @@ -556,7 +556,7 @@ func TestIndexerCache_GetOrCreate_PreferredRoot(t *testing.T) { cfg: newTestConfigService(t, 512), } // Pre-create the DB file at parentDir so the preferred root is adopted. - dbPath := config.DBPathForProject(parentDir, "stub") + dbPath := ic.dbPath(parentDir, "stub") if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil { t.Fatal(err) } @@ -1110,7 +1110,7 @@ func TestEnsureIndexed_SkipsWhenLockHeld(t *testing.T) { t.Fatalf("getOrCreate: %v", err) } - dbPath := config.DBPathForProject(effectiveRoot, ic.embedder.ModelName()) + dbPath := ic.dbPath(effectiveRoot, ic.embedder.ModelName()) lockPath := indexlock.LockPathForDB(dbPath) // Ensure the lock file's parent directory exists (getOrCreate creates the DB @@ -1223,8 +1223,10 @@ func TestGetOrCreate_PrePopulatesTTLFromRecentIndex(t *testing.T) { t.Fatal(err) } + t.Setenv("LUMEN_FRESHNESS_TTL", "30s") + cfg := newTestConfigService(t, 512) // Build a real DB at the expected path and stamp it with a recent timestamp. - dbPath := config.DBPathForProject(projectDir, "stub") + dbPath := configuredDBPath(cfg, projectDir, "stub") if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil { t.Fatal(err) } @@ -1232,10 +1234,9 @@ func TestGetOrCreate_PrePopulatesTTLFromRecentIndex(t *testing.T) { t.Fatal(err) } - t.Setenv("LUMEN_FRESHNESS_TTL", "30s") ic := &indexerCache{ embedder: &stubEmbedder{}, - cfg: newTestConfigService(t, 512), + cfg: cfg, } idx, _, _, err := ic.getOrCreate(projectDir, "") if err != nil { @@ -1256,7 +1257,9 @@ func TestGetOrCreate_DoesNotPrePopulateTTLFromOldIndex(t *testing.T) { t.Fatal(err) } - dbPath := config.DBPathForProject(projectDir, "stub") + t.Setenv("LUMEN_FRESHNESS_TTL", "30s") + cfg := newTestConfigService(t, 512) + dbPath := configuredDBPath(cfg, projectDir, "stub") if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil { t.Fatal(err) } @@ -1265,10 +1268,9 @@ func TestGetOrCreate_DoesNotPrePopulateTTLFromOldIndex(t *testing.T) { t.Fatal(err) } - t.Setenv("LUMEN_FRESHNESS_TTL", "30s") ic := &indexerCache{ embedder: &stubEmbedder{}, - cfg: newTestConfigService(t, 512), + cfg: cfg, } idx, _, _, err := ic.getOrCreate(projectDir, "") if err != nil { @@ -1443,7 +1445,7 @@ func TestEnsureIndexed_FreshnessTTL(t *testing.T) { Limit: 8, } - dbPath := config.DBPathForProject(effectiveRoot, ic.embedder.ModelName()) + dbPath := ic.dbPath(effectiveRoot, ic.embedder.ModelName()) // First call: no TTL entry yet — runs EnsureFresh and records lastCheckedAt. _, err = ic.ensureIndexed(idx, input, effectiveRoot, dbPath, nil) diff --git a/internal/index/index.go b/internal/index/index.go index 715dfe84..fcb1bcfe 100644 --- a/internal/index/index.go +++ b/internal/index/index.go @@ -89,6 +89,7 @@ type StatusInfo struct { // Indexer orchestrates chunking, embedding, and storage for a code index. type Indexer struct { mu sync.Mutex + projectMu sync.RWMutex store *store.Store emb embedder.Embedder chunker chunker.Chunker @@ -138,25 +139,34 @@ func NewIndexerForProject(dsn string, emb embedder.Embedder, maxChunkTokens int, } // rebuildStore closes the current store, deletes the database files, and -// opens a fresh store. Must be called while holding idx.mu.Lock() or before -// the Indexer is shared with other goroutines. -func (idx *Indexer) rebuildStore() error { +// opens a fresh store. Callers must hold idx.mu but must not hold a project +// lease: the exclusive project lock drains concurrent readers before the +// underlying connections are replaced. +func (idx *Indexer) rebuildStore(projectPath string) error { + idx.projectMu.Lock() + defer idx.projectMu.Unlock() + _ = idx.store.Close() if idx.dsn != "" && idx.dsn != ":memory:" { for _, suffix := range []string{"", "-wal", "-shm"} { _ = os.Remove(idx.dsn + suffix) } } - s, err := store.NewCollection(idx.dsn, idx.emb.Dimensions(), idx.vectorStorage, idx.projectPath) + s, err := store.NewCollection(idx.dsn, idx.emb.Dimensions(), idx.vectorStorage, projectPath) if err != nil { return fmt.Errorf("open fresh store: %w", err) } idx.store = s + idx.projectPath = projectPath return nil } // Close closes the underlying store. func (idx *Indexer) Close() error { + idx.mu.Lock() + defer idx.mu.Unlock() + idx.projectMu.Lock() + defer idx.projectMu.Unlock() return idx.store.Close() } @@ -185,9 +195,11 @@ func (idx *Indexer) Index(ctx context.Context, projectDir string, force bool, pr idx.mu.Lock() defer idx.mu.Unlock() - if err := idx.selectProject(projectDir); err != nil { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { return Stats{}, err } + defer func() { releaseProject() }() storedHash, err := idx.store.GetMeta("root_hash") if err != nil && err != sql.ErrNoRows { @@ -216,9 +228,15 @@ func (idx *Indexer) Index(ctx context.Context, projectDir string, force bool, pr idx.logger.Error("corrupted database detected during index, rebuilding", "project", projectDir, "err", indexErr) } - if rebuildErr := idx.rebuildStore(); rebuildErr != nil { + releaseProject() + releaseProject = func() {} + if rebuildErr := idx.rebuildStore(projectDir); rebuildErr != nil { return Stats{}, fmt.Errorf("rebuild corrupted db: %w", rebuildErr) } + releaseProject, err = idx.lockProject(projectDir) + if err != nil { + return Stats{}, err + } // Retry with force=true so the fresh DB gets a full index pass. stats, indexErr = idx.indexWithTree(ctx, projectDir, "", true, curTree, progress) if indexErr != nil { @@ -253,9 +271,11 @@ func (idx *Indexer) EnsureFresh(ctx context.Context, projectDir string, progress idx.mu.Lock() defer idx.mu.Unlock() - if err := idx.selectProject(projectDir); err != nil { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { return false, Stats{}, err } + defer func() { releaseProject() }() storedHash, err := idx.store.GetMeta("root_hash") if err != nil && err != sql.ErrNoRows { @@ -286,9 +306,15 @@ func (idx *Indexer) EnsureFresh(ctx context.Context, projectDir string, progress idx.logger.Error("corrupted database detected during reindex, rebuilding", "project", projectDir, "err", err) } - if rebuildErr := idx.rebuildStore(); rebuildErr != nil { + releaseProject() + releaseProject = func() {} + if rebuildErr := idx.rebuildStore(projectDir); rebuildErr != nil { return false, Stats{}, fmt.Errorf("rebuild corrupted db: %w", rebuildErr) } + releaseProject, err = idx.lockProject(projectDir) + if err != nil { + return false, Stats{}, err + } // Retry with empty storedHash so the fresh DB gets a full index pass. stats, err = idx.indexWithTree(ctx, projectDir, "", false, curTree, progress) if err != nil { @@ -552,6 +578,8 @@ func (idx *Indexer) indexWithTree(ctx context.Context, projectDir, oldRootHash s // stored in the last_indexed_at metadata field. Returns (zero, false) if the // field is absent or unparseable (e.g. the index has never been run). func (idx *Indexer) LastIndexedAt() (time.Time, bool) { + idx.projectMu.RLock() + defer idx.projectMu.RUnlock() val, err := idx.store.GetMeta("last_indexed_at") if err != nil || val == "" { return time.Time{}, false @@ -570,9 +598,11 @@ func (idx *Indexer) LastIndexedAt() (time.Time, bool) { // IsFresh does not acquire the indexer mutex; it reads through the store's // read-only connection (SQLite WAL isolation). func (idx *Indexer) IsFresh(projectDir string) (bool, error) { - if err := idx.selectProject(projectDir); err != nil { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { return false, err } + defer releaseProject() curTree, err := merkle.BuildTree(projectDir, makeSkip(projectDir)) if err != nil { return false, fmt.Errorf("build merkle tree: %w", err) @@ -597,9 +627,11 @@ func (idx *Indexer) IsFresh(projectDir string) (bool, error) { // acquire the indexer mutex, relying on SQLite WAL mode for isolation. func (idx *Indexer) Search(ctx context.Context, projectDir string, queryVec []float32, limit int, maxDistance float64, pathPrefix string) ([]store.SearchResult, error) { - if err := idx.selectProject(projectDir); err != nil { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { return nil, err } + defer releaseProject() return idx.store.Search(ctx, queryVec, limit, maxDistance, pathPrefix) } @@ -611,9 +643,11 @@ func (idx *Indexer) Search(ctx context.Context, projectDir string, queryVec []fl func (idx *Indexer) Status(projectDir string) (StatusInfo, error) { var info StatusInfo info.ProjectPath = projectDir - if err := idx.selectProject(projectDir); err != nil { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { return info, err } + defer releaseProject() storeStats, err := idx.store.Stats() if err != nil { @@ -643,15 +677,37 @@ func (idx *Indexer) Status(projectDir string) (StatusInfo, error) { return info, nil } -func (idx *Indexer) selectProject(projectDir string) error { - if projectDir == "" { - projectDir = idx.projectPath - } - if err := idx.store.UseProject(projectDir); err != nil { - return fmt.Errorf("select project: %w", err) +// lockProject returns with a shared project lease held for the requested +// membership. Same-project operations can therefore run concurrently, while a +// membership switch waits for all operations using the previous project to +// finish. The loop closes the gap between switching and reacquiring the shared +// lease if another waiter changes the membership first. +func (idx *Indexer) lockProject(projectDir string) (func(), error) { + for { + idx.projectMu.RLock() + selectedPath := projectDir + if selectedPath == "" { + selectedPath = idx.projectPath + } + if idx.projectPath == selectedPath { + return idx.projectMu.RUnlock, nil + } + idx.projectMu.RUnlock() + + idx.projectMu.Lock() + selectedPath = projectDir + if selectedPath == "" { + selectedPath = idx.projectPath + } + if idx.projectPath != selectedPath { + if err := idx.store.UseProject(selectedPath); err != nil { + idx.projectMu.Unlock() + return nil, fmt.Errorf("select project: %w", err) + } + idx.projectPath = selectedPath + } + idx.projectMu.Unlock() } - idx.projectPath = projectDir - return nil } // isBinaryContent reports whether data appears to be binary by checking diff --git a/internal/index/index_concurrency_test.go b/internal/index/index_concurrency_test.go index 2bcfdc5d..029bb1ea 100644 --- a/internal/index/index_concurrency_test.go +++ b/internal/index/index_concurrency_test.go @@ -364,6 +364,74 @@ done: <-indexDone } +func TestConcurrentProjectSelectionKeepsSearchAndStatusScoped(t *testing.T) { + const dims = 4 + projectA := t.TempDir() + projectB := t.TempDir() + writeGoFile(t, projectA, "alpha.go", `package alpha + +func Alpha() {} +`) + writeGoFile(t, projectB, "beta.go", `package beta + +func Beta() {} +`) + + idx, err := NewIndexer(filepath.Join(t.TempDir(), "shared.db"), &mockEmbedder{dims: dims, model: "project-race"}, 512) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + if _, err := idx.Index(context.Background(), projectA, false, nil); err != nil { + t.Fatalf("index project A: %v", err) + } + if _, err := idx.Index(context.Background(), projectB, false, nil); err != nil { + t.Fatalf("index project B: %v", err) + } + + type projectCase struct { + path string + file string + } + cases := []projectCase{{path: projectA, file: "alpha.go"}, {path: projectB, file: "beta.go"}} + start := make(chan struct{}) + errs := make(chan error, 40) + var wg sync.WaitGroup + for i := range 40 { + tc := cases[i%len(cases)] + wg.Add(1) + go func() { + defer wg.Done() + <-start + results, searchErr := idx.Search(context.Background(), tc.path, []float32{0.1, 0.1, 0.1, 0.1}, 5, 0, "") + if searchErr != nil { + errs <- fmt.Errorf("search %s: %w", tc.path, searchErr) + return + } + if len(results) == 0 || results[0].FilePath != tc.file { + errs <- fmt.Errorf("search %s returned wrong membership: %+v", tc.path, results) + return + } + status, statusErr := idx.Status(tc.path) + if statusErr != nil { + errs <- fmt.Errorf("status %s: %w", tc.path, statusErr) + return + } + if status.ProjectPath != tc.path || status.IndexedFiles != 1 { + errs <- fmt.Errorf("status %s returned wrong membership: %+v", tc.path, status) + } + }() + } + close(start) + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Error(err) + } + } +} + // TestSearch_RespectsContextCancellation verifies that a blocked Search call // returns promptly when its context is cancelled. func TestSearch_RespectsContextCancellation(t *testing.T) { From 15181383f62f90fd733cf398b88804b98f52d162 Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Sat, 8 Aug 2026 22:46:23 +0200 Subject: [PATCH 3/9] perf(store): scope shared collection garbage collection --- .gitattributes | 2 +- e2e_cli_test.go | 4 +- internal/sqlitevec/lib.go | 14 ++- internal/sqlitevec/lib_test.go | 4 +- internal/store/shared.go | 203 +++++++++++++++++++++++++++------ internal/store/shared_test.go | 97 ++++++++++++++++ internal/store/store.go | 19 ++- internal/store/store_test.go | 21 ++++ 8 files changed, 317 insertions(+), 47 deletions(-) diff --git a/.gitattributes b/.gitattributes index 0165b268..b2dceba6 100644 --- a/.gitattributes +++ b/.gitattributes @@ -3,7 +3,7 @@ docs/plans linguist-generated # Keep the vendored sqlite-vec release source byte-for-byte identical to # upstream, including its existing whitespace. -internal/sqlitevec/sqlite-vec.c -whitespace linguist-vendored +internal/sqlitevec/sqlite-vec.[ch] -text -whitespace linguist-vendored # Windows batch files require CRLF line endings — the cmd.exe parser # has 512-byte boundary bugs with bare-LF files (GOTO/CALL label diff --git a/e2e_cli_test.go b/e2e_cli_test.go index a31fcc11..4d9e8bc4 100644 --- a/e2e_cli_test.go +++ b/e2e_cli_test.go @@ -139,7 +139,9 @@ func TestE2E_CLI_IndexForceReindex(t *testing.T) { // openIndexDB opens the SQLite index database for a given project path and dataHome. func openIndexDB(t *testing.T, dataHome, projectPath string) *sql.DB { t.Helper() - sqlite_vec.Auto() + if err := sqlite_vec.Auto(); err != nil { + t.Fatal(err) + } dbPath := config.DBPathForProjectBase(dataHome, projectPath, "all-minilm") db, err := sql.Open("sqlite3", dbPath) if err != nil { diff --git a/internal/sqlitevec/lib.go b/internal/sqlitevec/lib.go index 4550f58f..0075521a 100644 --- a/internal/sqlitevec/lib.go +++ b/internal/sqlitevec/lib.go @@ -4,22 +4,30 @@ package sqlitevec // #cgo CFLAGS: -DSQLITE_CORE // #cgo linux LDFLAGS: -lm +// /* SQLITE_CORE intentionally resolves sqlite-vec's SQLite symbols against +// the same mattn/go-sqlite3 amalgamation linked into this process. This matches +// the upstream sqlite-vec Go bindings and avoids loading a second SQLite ABI. */ // #include "sqlite-vec.h" import "C" import ( "bytes" "encoding/binary" + "fmt" ) // Auto registers sqlite-vec for every SQLite connection opened afterward. -func Auto() { - C.sqlite3_auto_extension((*[0]byte)(C.sqlite3_vec_init)) +func Auto() error { + if rc := C.sqlite3_auto_extension((*[0]byte)(C.sqlite3_vec_init)); rc != C.SQLITE_OK { + return fmt.Errorf("register sqlite-vec auto extension: sqlite error %d", int(rc)) + } + return nil } // Cancel cancels the automatic sqlite-vec extension registration. func Cancel() { - C.sqlite3_cancel_auto_extension((*[0]byte)(C.sqlite3_vec_init)) + // Cancellation is best-effort cleanup; registration may already be absent. + _ = C.sqlite3_cancel_auto_extension((*[0]byte)(C.sqlite3_vec_init)) } // SerializeFloat32 encodes a vector as sqlite-vec's little-endian float BLOB. diff --git a/internal/sqlitevec/lib_test.go b/internal/sqlitevec/lib_test.go index 61514d91..cf954d71 100644 --- a/internal/sqlitevec/lib_test.go +++ b/internal/sqlitevec/lib_test.go @@ -13,7 +13,9 @@ import ( ) func TestBundledVersion(t *testing.T) { - Auto() + if err := Auto(); err != nil { + t.Fatal(err) + } db, err := sql.Open("sqlite3", ":memory:") if err != nil { t.Fatal(err) diff --git a/internal/store/shared.go b/internal/store/shared.go index a323fa73..b08f4a68 100644 --- a/internal/store/shared.go +++ b/internal/store/shared.go @@ -13,6 +13,7 @@ import ( "crypto/sha256" "database/sql" "encoding/hex" + "errors" "fmt" "math" "os" @@ -26,6 +27,11 @@ import ( const sharedSchemaVersion = "1" +// ErrVectorVanished marks a cross-process race where a vector observed during +// MissingChunkInputs was garbage-collected before StoreFileRevision began its +// write transaction. Callers can re-embed the complete file and retry once. +var ErrVectorVanished = errors.New("shared vector vanished") + // CollectionStats describes both project-local references and collection-wide // physical storage. Chunk references may exceed UniqueVectors because vectors // are content-addressed across file revisions and projects. @@ -71,8 +77,18 @@ func openCollection(dsn string, dimensions int, vectorStorage string) (*Store, e // Legacy per-worktree databases remain readable during the lazy migration // window. New profile paths always create the shared schema; an existing // legacy path is upgraded by normal indexing without making it unreadable. - if legacy, _ := checkTableExists(db, "files"); legacy { - if shared, _ := checkTableExists(db, "collection_meta"); !shared { + legacy, err := checkTableExists(db, "files") + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("check files table: %w", err) + } + if legacy { + shared, err := checkTableExists(db, "collection_meta") + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("check collection_meta table: %w", err) + } + if !shared { _ = db.Close() return openStore(dsn, dimensions) } @@ -216,7 +232,9 @@ func createCollectionSchema(db *sql.DB, dimensions int, vectorStorage string) er func (s *Store) IsShared() bool { return s.shared } // UseProject selects (and, if necessary, creates) a project membership in the -// shared collection. Calling it repeatedly for the same path is cheap. +// shared collection. Calling it repeatedly for the same path is cheap. Callers +// must serialize membership switches against Store operations; Indexer does so +// through its project lease. func (s *Store) UseProject(projectPath string) error { if !s.shared { return nil @@ -347,10 +365,11 @@ func (s *Store) AttachExistingFileRevision(relativePath, contentHash string) (bo if err != nil { return false, err } - if err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID); err != nil { + displaced, err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID) + if err != nil { return false, err } - if err := gcUnreferencedTx(tx); err != nil { + if err := gcRevisionsTx(tx, displaced); err != nil { return false, err } return true, tx.Commit() @@ -359,14 +378,43 @@ func (s *Store) AttachExistingFileRevision(relativePath, contentHash string) (bo // MissingChunkInputs returns the chunk positions whose exact embedding inputs // are absent from the collection. Callers only need to embed these positions. func (s *Store) MissingChunkInputs(chunks []chunker.Chunk) ([]int, error) { - missing := make([]int, 0, len(chunks)) + const queryBatchSize = 256 + hashes := make([][sha256.Size]byte, len(chunks)) + present := make(map[[sha256.Size]byte]struct{}, len(chunks)) for i, c := range chunks { - h := embeddingInputHash(c) - var exists bool - if err := s.reader().QueryRow(`SELECT EXISTS(SELECT 1 FROM vector_keys WHERE input_hash = ?)`, h[:]).Scan(&exists); err != nil { - return nil, err + hashes[i] = embeddingInputHash(c) + } + for start := 0; start < len(hashes); start += queryBatchSize { + end := min(start+queryBatchSize, len(hashes)) + marks := make([]string, end-start) + args := make([]any, end-start) + for i := start; i < end; i++ { + marks[i-start] = "?" + args[i-start] = hashes[i][:] } - if !exists { + rows, err := s.reader().Query(`SELECT input_hash FROM vector_keys WHERE input_hash IN (`+strings.Join(marks, ",")+`)`, args...) + if err != nil { + return nil, fmt.Errorf("query shared vector inputs: %w", err) + } + for rows.Next() { + var blob []byte + if err := rows.Scan(&blob); err != nil { + _ = rows.Close() + return nil, fmt.Errorf("scan shared vector input: %w", err) + } + if len(blob) == sha256.Size { + var hash [sha256.Size]byte + copy(hash[:], blob) + present[hash] = struct{}{} + } + } + if err := errors.Join(rows.Err(), rows.Close()); err != nil { + return nil, fmt.Errorf("read shared vector inputs: %w", err) + } + } + missing := make([]int, 0, len(chunks)) + for i, hash := range hashes { + if _, ok := present[hash]; !ok { missing = append(missing, i) } } @@ -401,7 +449,7 @@ func (s *Store) StoreFileRevision(relativePath, contentHash string, chunks []chu if err == sql.ErrNoRows { vec, ok := vectors[i] if !ok { - return false, fmt.Errorf("missing embedding for chunk %d (%s)", i, c.ID) + return false, fmt.Errorf("%w: missing embedding for chunk %d (%s)", ErrVectorVanished, i, c.ID) } result, err := tx.Exec(`INSERT INTO vector_keys(input_hash) VALUES (?)`, h[:]) if err != nil { @@ -470,19 +518,31 @@ func (s *Store) StoreFileRevision(relativePath, contentHash string, chunks []chu } } - if err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID); err != nil { + displaced, err := replaceProjectFileTx(tx, s.projectID, relativePath, revisionID) + if err != nil { return false, err } - if err := gcUnreferencedTx(tx); err != nil { + if err := gcRevisionsTx(tx, displaced); err != nil { return false, err } return inserted > 0, tx.Commit() } -func replaceProjectFileTx(tx *sql.Tx, projectID int64, path string, revisionID int64) error { - _, err := tx.Exec(`INSERT INTO project_files(project_id, relative_path, file_revision_id) VALUES (?, ?, ?) +func replaceProjectFileTx(tx *sql.Tx, projectID int64, path string, revisionID int64) (int64, error) { + var displaced int64 + err := tx.QueryRow(`SELECT file_revision_id FROM project_files WHERE project_id = ? AND relative_path = ?`, projectID, path).Scan(&displaced) + if err != nil && err != sql.ErrNoRows { + return 0, err + } + _, err = tx.Exec(`INSERT INTO project_files(project_id, relative_path, file_revision_id) VALUES (?, ?, ?) ON CONFLICT(project_id, relative_path) DO UPDATE SET file_revision_id = excluded.file_revision_id`, projectID, path, revisionID) - return err + if err != nil { + return 0, err + } + if displaced == revisionID { + return 0, nil + } + return displaced, nil } func (s *Store) serializeVector(vector []float32) ([]byte, error) { @@ -521,21 +581,27 @@ func (s *Store) upsertSharedFile(path, hash string) error { return err } } - _, err := s.db.Exec(`INSERT OR IGNORE INTO file_revisions(relative_path, content_hash, complete) VALUES (?, ?, 0)`, path, hashBlob(hash)) + tx, err := s.db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + _, err = tx.Exec(`INSERT OR IGNORE INTO file_revisions(relative_path, content_hash, complete) VALUES (?, ?, 0)`, path, hashBlob(hash)) if err != nil { return err } var revisionID int64 - if err := s.db.QueryRow(`SELECT id FROM file_revisions WHERE relative_path = ? AND content_hash = ?`, path, hashBlob(hash)).Scan(&revisionID); err != nil { + if err := tx.QueryRow(`SELECT id FROM file_revisions WHERE relative_path = ? AND content_hash = ?`, path, hashBlob(hash)).Scan(&revisionID); err != nil { return err } - return replaceProjectFileDB(s.db, s.projectID, path, revisionID) -} - -func replaceProjectFileDB(db *sql.DB, projectID int64, path string, revisionID int64) error { - _, err := db.Exec(`INSERT INTO project_files(project_id, relative_path, file_revision_id) VALUES (?, ?, ?) - ON CONFLICT(project_id, relative_path) DO UPDATE SET file_revision_id = excluded.file_revision_id`, projectID, path, revisionID) - return err + displaced, err := replaceProjectFileTx(tx, s.projectID, path, revisionID) + if err != nil { + return err + } + if err := gcRevisionsTx(tx, displaced); err != nil { + return err + } + return tx.Commit() } func (s *Store) insertSharedChunks(chunks []chunker.Chunk, vectors [][]float32) error { @@ -550,7 +616,10 @@ func (s *Store) insertSharedChunks(chunks []chunker.Chunk, vectors [][]float32) var hash []byte if err := s.db.QueryRow(`SELECT fr.content_hash FROM project_files pf JOIN file_revisions fr ON fr.id = pf.file_revision_id WHERE pf.project_id = ? AND pf.relative_path = ?`, s.projectID, path).Scan(&hash); err != nil { - return err + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("no file revision registered for %q; call UpsertFile first: %w", path, err) + } + return fmt.Errorf("read file revision for %q: %w", path, err) } fileChunks := make([]chunker.Chunk, len(positions)) fileVectors := make(map[int][]float32, len(positions)) @@ -571,16 +640,76 @@ func (s *Store) removeProjectFile(path string) error { return err } defer func() { _ = tx.Rollback() }() + var displaced int64 + err = tx.QueryRow(`SELECT file_revision_id FROM project_files WHERE project_id = ? AND relative_path = ?`, s.projectID, path).Scan(&displaced) + if err != nil && err != sql.ErrNoRows { + return err + } if _, err := tx.Exec(`DELETE FROM project_files WHERE project_id = ? AND relative_path = ?`, s.projectID, path); err != nil { return err } - if err := gcUnreferencedTx(tx); err != nil { + if err := gcRevisionsTx(tx, displaced); err != nil { return err } return tx.Commit() } -func gcUnreferencedTx(tx *sql.Tx) error { +func gcRevisionsTx(tx *sql.Tx, revisionIDs ...int64) error { + seen := make(map[int64]struct{}, len(revisionIDs)) + for _, revisionID := range revisionIDs { + if revisionID == 0 { + continue + } + if _, ok := seen[revisionID]; ok { + continue + } + seen[revisionID] = struct{}{} + var referenced bool + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM project_files WHERE file_revision_id = ?)`, revisionID).Scan(&referenced); err != nil { + return err + } + if referenced { + continue + } + rows, err := tx.Query(`SELECT DISTINCT vector_id FROM chunk_defs WHERE file_revision_id = ?`, revisionID) + if err != nil { + return err + } + var vectorIDs []int64 + for rows.Next() { + var vectorID int64 + if err := rows.Scan(&vectorID); err != nil { + _ = rows.Close() + return err + } + vectorIDs = append(vectorIDs, vectorID) + } + if err := errors.Join(rows.Err(), rows.Close()); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM file_revisions WHERE id = ?`, revisionID); err != nil { + return err + } + for _, vectorID := range vectorIDs { + var used bool + if err := tx.QueryRow(`SELECT EXISTS(SELECT 1 FROM chunk_defs WHERE vector_id = ?)`, vectorID).Scan(&used); err != nil { + return err + } + if used { + continue + } + if _, err := tx.Exec(`DELETE FROM vec_vectors WHERE vector_id = ?`, vectorID); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM vector_keys WHERE id = ?`, vectorID); err != nil { + return err + } + } + } + return nil +} + +func gcAllUnreferencedTx(tx *sql.Tx) error { if _, err := tx.Exec(`DELETE FROM file_revisions WHERE NOT EXISTS ( SELECT 1 FROM project_files WHERE file_revision_id = file_revisions.id)`); err != nil { return err @@ -599,7 +728,7 @@ func gcUnreferencedTx(tx *sql.Tx) error { } ids = append(ids, id) } - if err := rows.Close(); err != nil { + if err := errors.Join(rows.Err(), rows.Close()); err != nil { return err } for _, id := range ids { @@ -819,7 +948,7 @@ func (s *Store) CleanupStaleProjects(cutoff time.Time) (CleanupStats, error) { stale = append(stale, p) } } - if err := rows.Close(); err != nil { + if err := errors.Join(rows.Err(), rows.Close()); err != nil { return CleanupStats{}, err } tx, err := s.db.Begin() @@ -832,7 +961,7 @@ func (s *Store) CleanupStaleProjects(cutoff time.Time) (CleanupStats, error) { return CleanupStats{}, err } } - if err := gcUnreferencedTx(tx); err != nil { + if err := gcAllUnreferencedTx(tx); err != nil { return CleanupStats{}, err } if err := tx.Commit(); err != nil { @@ -864,26 +993,28 @@ func CleanupCollectionAt(dbPath string, cutoff time.Time) (CleanupStats, bool, e if err != nil { return CleanupStats{}, false, err } + defer func() { + if db != nil { + _ = db.Close() + } + }() var shared bool if err := db.QueryRow(`SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'collection_meta')`).Scan(&shared); err != nil { - _ = db.Close() return CleanupStats{}, false, err } if !shared { - _ = db.Close() return CleanupStats{}, false, nil } var dimensions int var storage string if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vec_dimensions'`).Scan(&dimensions); err != nil { - _ = db.Close() return CleanupStats{}, true, err } if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vector_storage'`).Scan(&storage); err != nil { - _ = db.Close() return CleanupStats{}, true, err } _ = db.Close() + db = nil s, err := openCollection(dbPath, dimensions, storage) if err != nil { return CleanupStats{}, true, err diff --git a/internal/store/shared_test.go b/internal/store/shared_test.go index d98a8373..a547bdc9 100644 --- a/internal/store/shared_test.go +++ b/internal/store/shared_test.go @@ -11,6 +11,7 @@ import ( "os" "path/filepath" "strconv" + "strings" "sync" "testing" "time" @@ -18,6 +19,73 @@ import ( "github.com/ory/lumen/internal/chunker" ) +func TestSharedReplacementGarbageCollectsDisplacedRevision(t *testing.T) { + s, err := NewCollection(":memory:", 4, "int8", t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + oldChunk := chunker.Chunk{ID: "old", FilePath: "main.go", Symbol: "Old", Kind: "function", StartLine: 1, EndLine: 1, Content: "func Old() {}"} + newChunk := chunker.Chunk{ID: "new", FilePath: "main.go", Symbol: "New", Kind: "function", StartLine: 1, EndLine: 1, Content: "func New() {}"} + if _, err := s.StoreFileRevision("main.go", "old", []chunker.Chunk{oldChunk}, map[int][]float32{0: {1, 0, 0, 0}}); err != nil { + t.Fatal(err) + } + if _, err := s.StoreFileRevision("main.go", "new", []chunker.Chunk{newChunk}, map[int][]float32{0: {0, 1, 0, 0}}); err != nil { + t.Fatal(err) + } + stats, err := s.CollectionStats() + if err != nil { + t.Fatal(err) + } + if stats.UniqueVectors != 1 || stats.ChunkReferences != 1 { + t.Fatalf("displaced revision was not collected: %+v", stats) + } +} + +func TestMissingChunkInputsBatchesLargeQueries(t *testing.T) { + s, err := NewCollection(":memory:", 4, "int8", t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + chunks := make([]chunker.Chunk, 300) + vectors := make(map[int][]float32, len(chunks)) + for i := range chunks { + chunks[i] = chunker.Chunk{ID: strconv.Itoa(i), FilePath: "large.go", Content: "input " + strconv.Itoa(i)} + vectors[i] = []float32{1, 0, 0, 0} + } + if _, err := s.StoreFileRevision("large.go", "large", chunks, vectors); err != nil { + t.Fatal(err) + } + missing, err := s.MissingChunkInputs(chunks) + if err != nil { + t.Fatal(err) + } + if len(missing) != 0 { + t.Fatalf("missing = %v, want none", missing) + } + chunks[299].Content = "changed" + missing, err = s.MissingChunkInputs(chunks) + if err != nil { + t.Fatal(err) + } + if len(missing) != 1 || missing[0] != 299 { + t.Fatalf("missing = %v, want [299]", missing) + } +} + +func TestInsertSharedChunksRequiresRegisteredRevision(t *testing.T) { + s, err := NewCollection(":memory:", 4, "int8", t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + err = s.InsertChunks([]chunker.Chunk{{ID: "x", FilePath: "missing.go", Content: "x"}}, [][]float32{{1, 0, 0, 0}}) + if err == nil || !strings.Contains(err.Error(), `no file revision registered for "missing.go"; call UpsertFile first`) { + t.Fatalf("error = %v", err) + } +} + func TestSharedCollectionReusesRevisionsAndVectors(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "index.db") projectA := filepath.Join(t.TempDir(), "worktree-a") @@ -118,6 +186,35 @@ func TestSharedCollectionFloat32Override(t *testing.T) { } } +func TestSharedCollectionValidatesStorageProfile(t *testing.T) { + if _, err := NewCollection(":memory:", 4, "float16", t.TempDir()); err == nil { + t.Fatal("expected unsupported vector storage to fail") + } + dbPath := filepath.Join(t.TempDir(), "index.db") + s, err := NewCollection(dbPath, 4, "int8", t.TempDir()) + if err != nil { + t.Fatal(err) + } + if err := s.Close(); err != nil { + t.Fatal(err) + } + for _, tc := range []struct { + name string + dimensions int + storage string + }{ + {"dimensions", 5, "int8"}, + {"storage", 4, "float32"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := NewCollection(dbPath, tc.dimensions, tc.storage, t.TempDir()) + if err == nil || !strings.Contains(err.Error(), "collection profile mismatch") { + t.Fatalf("error = %v, want profile mismatch", err) + } + }) + } +} + func TestSharedSearchAdaptivelyExpandsSparseProjectCandidates(t *testing.T) { s, err := NewCollection(":memory:", 4, "int8", "/project-a") if err != nil { diff --git a/internal/store/store.go b/internal/store/store.go index 796f047d..185ab0aa 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -39,7 +39,9 @@ const MetaLastAccessedAt = "last_accessed_at" const accessStampBusyTimeoutMS = 250 func init() { - sqlite_vec.Auto() + if err := sqlite_vec.Auto(); err != nil { + panic(err) + } } // IsCorruptionErr reports whether err indicates SQLite database corruption. @@ -84,7 +86,8 @@ type StoreStats struct { //nolint:revive // StoreStats is intentionally named to } // Store manages SQLite + sqlite-vec storage for code chunks and their -// embedding vectors. +// embedding vectors. A shared Store's selected project is mutable; callers +// must serialize UseProject with all operations that depend on that selection. type Store struct { db *sql.DB readDB *sql.DB // separate read-only connection; nil for :memory: databases @@ -130,7 +133,8 @@ func New(dsn string, dimensions int) (*Store, error) { // NewCollection opens a repository-scoped shared collection and selects the // project membership identified by projectPath. vectorStorage must be int8 or // float32. Multiple Store instances may safely select different worktrees in -// the same database. +// the same database. If schema setup detects corruption, on-disk database and +// sidecar files are removed and creation is retried once. func NewCollection(dsn string, dimensions int, vectorStorage, projectPath string) (*Store, error) { if vectorStorage != "int8" && vectorStorage != "float32" { return nil, fmt.Errorf("unsupported vector storage %q", vectorStorage) @@ -211,11 +215,16 @@ func openStore(dsn string, dimensions int) (*Store, error) { // A low-level caller may open a database already upgraded to the shared // schema (for example metadata tooling during a rolling upgrade). Return a // shared view instead of attempting to overlay the legacy tables. - if shared, _ := checkTableExists(db, "collection_meta"); shared { + shared, err := checkTableExists(db, "collection_meta") + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("check collection_meta table: %w", err) + } + if shared { var storage string if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = 'vector_storage'`).Scan(&storage); err != nil { _ = db.Close() - return nil, err + return nil, fmt.Errorf("read vector_storage: %w", err) } _ = db.Close() return openCollection(dsn, dimensions, storage) diff --git a/internal/store/store_test.go b/internal/store/store_test.go index c24180b1..e8f2ccd9 100644 --- a/internal/store/store_test.go +++ b/internal/store/store_test.go @@ -47,6 +47,27 @@ func TestNewStore_CreatesSchema(t *testing.T) { } } +func TestOpenStoreReportsMissingSharedVectorStorage(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "index.db") + db, err := sql.Open("sqlite3", dbPath) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`CREATE TABLE collection_meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)`); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`INSERT INTO collection_meta(key, value) VALUES ('vec_dimensions', '4')`); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + _, err = New(dbPath, 4) + if err == nil || !strings.Contains(err.Error(), "read vector_storage") { + t.Fatalf("error = %v, want read vector_storage context", err) + } +} + func TestStore_SetGetMeta(t *testing.T) { s, err := New(":memory:", 4) if err != nil { From c6dc73360fb4770439aef17710201d803ac87fc2 Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Sat, 8 Aug 2026 22:46:29 +0200 Subject: [PATCH 4/9] perf(index): batch embeddings across shared files --- internal/config/config.go | 2 + internal/config/config_test.go | 21 ++- internal/config/service.go | 2 +- internal/index/index.go | 26 +-- internal/index/index_concurrency_test.go | 10 +- internal/index/index_test.go | 34 +++- internal/index/migrate.go | 11 +- internal/index/migrate_test.go | 64 +++++++- internal/index/shared.go | 194 ++++++++++++++++++----- internal/index/shared_batch_test.go | 170 ++++++++++++++++++++ 10 files changed, 466 insertions(+), 68 deletions(-) create mode 100644 internal/index/shared_batch_test.go diff --git a/internal/config/config.go b/internal/config/config.go index 235b0ca7..e11f21a0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -79,6 +79,8 @@ func DBPathForProjectProfileBase(dataDir, projectPath, model string, dimensions // scope explicit in the key so subdirectory collections can be added // without another on-disk format change. scope = "." + } else if resolved, resolveErr := filepath.EvalSymlinks(identity); resolveErr == nil { + identity = filepath.Clean(resolved) } profile := identity + "\x00" + scope + "\x00" + model + "\x00" + strconv.Itoa(dimensions) + "\x00" + vectorStorage + "\x00" + diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d5b771fe..2d163e36 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -107,6 +107,23 @@ func TestDBPathForProjectProfileSharesGitWorktrees(t *testing.T) { } } +func TestDBPathForProjectProfileResolvesNonGitSymlinks(t *testing.T) { + realProject := filepath.Join(t.TempDir(), "project") + if err := os.MkdirAll(realProject, 0o755); err != nil { + t.Fatal(err) + } + link := filepath.Join(t.TempDir(), "project-link") + if err := os.Symlink(realProject, link); err != nil { + t.Fatal(err) + } + dataDir := t.TempDir() + realPath := DBPathForProjectProfileBase(dataDir, realProject, "model", 768, "int8", 512) + linkPath := DBPathForProjectProfileBase(dataDir, link, "model", 768, "int8", 512) + if realPath != linkPath { + t.Fatalf("symlinked non-Git project should share identity: %q != %q", realPath, linkPath) + } +} + func TestXDGConfigDir(t *testing.T) { t.Run("uses XDG_CONFIG_HOME when set", func(t *testing.T) { t.Setenv("XDG_CONFIG_HOME", "/custom/config") @@ -135,8 +152,8 @@ func TestVectorStorageConfiguration(t *testing.T) { t.Fatalf("VectorStorage() = %q, want int8", got) } }) - t.Run("accepts float32 override", func(t *testing.T) { - t.Setenv("LUMEN_VECTOR_STORAGE", "FLOAT32") + t.Run("normalizes float32 override", func(t *testing.T) { + t.Setenv("LUMEN_VECTOR_STORAGE", " FLOAT32 ") cfg, err := NewConfigService("") if err != nil { t.Fatal(err) diff --git a/internal/config/service.go b/internal/config/service.go index 9b408262..560aec14 100644 --- a/internal/config/service.go +++ b/internal/config/service.go @@ -312,7 +312,7 @@ func (s *ConfigService) MaxChunkTokens() int { func (s *ConfigService) VectorStorage() string { s.mu.RLock() defer s.mu.RUnlock() - return s.k.String("vector_storage") + return strings.ToLower(strings.TrimSpace(s.k.String("vector_storage"))) } func (s *ConfigService) FreshnessTTL() time.Duration { diff --git a/internal/index/index.go b/internal/index/index.go index fcb1bcfe..1c2fdec5 100644 --- a/internal/index/index.go +++ b/internal/index/index.go @@ -100,6 +100,7 @@ type Indexer struct { projectPath string legacyVectors map[[32]byte][]float32 legacySource string + embedBatchSize int } // SetLogger attaches a logger to the indexer for structured diagnostic output. @@ -135,6 +136,7 @@ func NewIndexerForProject(dsn string, emb embedder.Embedder, maxChunkTokens int, dsn: dsn, vectorStorage: vectorStorage, projectPath: projectPath, + embedBatchSize: 256, }, nil } @@ -574,12 +576,15 @@ func (idx *Indexer) indexWithTree(ctx context.Context, projectDir, oldRootHash s return stats, nil } -// LastIndexedAt returns the time the index was last successfully updated, as -// stored in the last_indexed_at metadata field. Returns (zero, false) if the -// field is absent or unparseable (e.g. the index has never been run). -func (idx *Indexer) LastIndexedAt() (time.Time, bool) { - idx.projectMu.RLock() - defer idx.projectMu.RUnlock() +// LastIndexedAt returns the time projectDir was last successfully updated, as +// stored in its last_indexed_at metadata field. Returns (zero, false) if the +// field is absent or unparseable (e.g. the project has never been indexed). +func (idx *Indexer) LastIndexedAt(projectDir string) (time.Time, bool) { + releaseProject, err := idx.lockProject(projectDir) + if err != nil { + return time.Time{}, false + } + defer releaseProject() val, err := idx.store.GetMeta("last_indexed_at") if err != nil || val == "" { return time.Time{}, false @@ -598,16 +603,15 @@ func (idx *Indexer) LastIndexedAt() (time.Time, bool) { // IsFresh does not acquire the indexer mutex; it reads through the store's // read-only connection (SQLite WAL isolation). func (idx *Indexer) IsFresh(projectDir string) (bool, error) { + curTree, err := merkle.BuildTree(projectDir, makeSkip(projectDir)) + if err != nil { + return false, fmt.Errorf("build merkle tree: %w", err) + } releaseProject, err := idx.lockProject(projectDir) if err != nil { return false, err } defer releaseProject() - curTree, err := merkle.BuildTree(projectDir, makeSkip(projectDir)) - if err != nil { - return false, fmt.Errorf("build merkle tree: %w", err) - } - storedHash, err := idx.store.GetMeta("root_hash") if err != nil && err != sql.ErrNoRows { return false, fmt.Errorf("get root_hash: %w", err) diff --git a/internal/index/index_concurrency_test.go b/internal/index/index_concurrency_test.go index 029bb1ea..3f4c2408 100644 --- a/internal/index/index_concurrency_test.go +++ b/internal/index/index_concurrency_test.go @@ -395,8 +395,16 @@ func Beta() {} } cases := []projectCase{{path: projectA, file: "alpha.go"}, {path: projectB, file: "beta.go"}} start := make(chan struct{}) - errs := make(chan error, 40) + errs := make(chan error, 41) var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + <-start + if _, indexErr := idx.Index(context.Background(), projectA, false, nil); indexErr != nil { + errs <- fmt.Errorf("concurrent index project A: %w", indexErr) + } + }() for i := range 40 { tc := cases[i%len(cases)] wg.Add(1) diff --git a/internal/index/index_test.go b/internal/index/index_test.go index b2b34dde..d56cd02b 100644 --- a/internal/index/index_test.go +++ b/internal/index/index_test.go @@ -69,6 +69,35 @@ func Hello(name string) { fmt.Println("hello", name) } +func TestIndexerLastIndexedAtIsProjectScoped(t *testing.T) { + projectA, projectB := t.TempDir(), t.TempDir() + idx, err := NewIndexerForProject(":memory:", &mockEmbedder{dims: 4, model: "test-model"}, 512, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + timeA := time.Now().Add(-time.Hour).UTC().Truncate(time.Second) + timeB := time.Now().UTC().Truncate(time.Second) + if err := idx.store.SetMeta("last_indexed_at", timeA.Format(time.RFC3339)); err != nil { + t.Fatal(err) + } + release, err := idx.lockProject(projectB) + if err != nil { + t.Fatal(err) + } + if err := idx.store.SetMeta("last_indexed_at", timeB.Format(time.RFC3339)); err != nil { + release() + t.Fatal(err) + } + release() + if got, ok := idx.LastIndexedAt(projectA); !ok || !got.Equal(timeA) { + t.Fatalf("project A LastIndexedAt = %v, %v; want %v, true", got, ok, timeA) + } + if got, ok := idx.LastIndexedAt(projectB); !ok || !got.Equal(timeB) { + t.Fatalf("project B LastIndexedAt = %v, %v; want %v, true", got, ok, timeB) + } +} + // Goodbye prints a farewell. func Goodbye(name string) { fmt.Println("bye", name) @@ -554,7 +583,7 @@ func TestIndexer_LastIndexedAt_ReturnsFalseWhenNotIndexed(t *testing.T) { } defer func() { _ = idx.Close() }() - _, ok := idx.LastIndexedAt() + _, ok := idx.LastIndexedAt("") if ok { t.Fatal("expected ok=false for an index with no last_indexed_at metadata") } @@ -578,7 +607,7 @@ func TestIndexer_LastIndexedAt_ReturnsTimeAfterIndex(t *testing.T) { } after := time.Now().Add(time.Second) - at, ok := idx.LastIndexedAt() + at, ok := idx.LastIndexedAt(projectDir) if !ok { t.Fatal("expected ok=true after Index was called") } @@ -619,7 +648,6 @@ func TestIndexer_StaleUnsupportedExtensionNotCountedAsRemoved(t *testing.T) { // The test passing without error means the ghost record was not propagated. } - // TestIndexer_StaleUnsupportedExtensionDeletedFromDB verifies that after a // reindex, stale file records with unsupported extensions (e.g. .md from // donor seeding) are purged from the DB. diff --git a/internal/index/migrate.go b/internal/index/migrate.go index 6a7b6779..f22016e4 100644 --- a/internal/index/migrate.go +++ b/internal/index/migrate.go @@ -10,6 +10,7 @@ import ( "database/sql" "encoding/binary" "encoding/hex" + "errors" "fmt" "math" "os" @@ -31,15 +32,11 @@ func (idx *Indexer) PrepareLegacyMigration(projectDir, legacyPath string) error } return err } - db, err := sql.Open("sqlite3", legacyPath) + db, err := sql.Open("sqlite3", "file:"+legacyPath+"?mode=ro&_query_only=1") if err != nil { return err } defer func() { _ = db.Close() }() - if _, err := db.Exec(`PRAGMA query_only=ON`); err != nil { - return err - } - legacyByChunk := make(map[string][]float32) rows, err := db.Query(`SELECT c.id, v.embedding FROM chunks c JOIN vec_chunks v ON v.id = c.id`) if err != nil { @@ -57,7 +54,7 @@ func (idx *Indexer) PrepareLegacyMigration(projectDir, legacyPath string) error legacyByChunk[id] = vector } } - if err := rows.Close(); err != nil { + if err := errors.Join(rows.Err(), rows.Close()); err != nil { return err } @@ -94,7 +91,7 @@ func (idx *Indexer) PrepareLegacyMigration(projectDir, legacyPath string) error } } } - if err := fileRows.Close(); err != nil { + if err := errors.Join(fileRows.Err(), fileRows.Close()); err != nil { return err } idx.legacyVectors = recovered diff --git a/internal/index/migrate_test.go b/internal/index/migrate_test.go index 50bd1009..545bfb73 100644 --- a/internal/index/migrate_test.go +++ b/internal/index/migrate_test.go @@ -24,7 +24,7 @@ func TestLegacyMigrationReusesUnchangedVectors(t *testing.T) { } emb := &mockEmbedder{dims: 4, model: "test-model"} newPath := filepath.Join(t.TempDir(), "shared.db") - idx, err := NewIndexerForProject(newPath, emb, 0, "int8", projectDir) + idx, err := NewIndexerForProject(newPath, emb, 512, "int8", projectDir) if err != nil { t.Fatal(err) } @@ -53,10 +53,22 @@ func TestLegacyMigrationReusesUnchangedVectors(t *testing.T) { if err := legacy.Close(); err != nil { t.Fatal(err) } + legacyBefore, err := os.ReadFile(legacyPath) + if err != nil { + t.Fatal(err) + } + legacyDigest := sha256.Sum256(legacyBefore) if err := idx.PrepareLegacyMigration(projectDir, legacyPath); err != nil { t.Fatal(err) } + legacyAfter, err := os.ReadFile(legacyPath) + if err != nil { + t.Fatal(err) + } + if got := sha256.Sum256(legacyAfter); got != legacyDigest { + t.Fatal("PrepareLegacyMigration modified the read-only legacy database") + } stats, err := idx.Index(context.Background(), projectDir, false, nil) if err != nil { t.Fatal(err) @@ -71,3 +83,53 @@ func TestLegacyMigrationReusesUnchangedVectors(t *testing.T) { t.Fatalf("legacy database should be removed after verification, stat err=%v", err) } } + +func TestLegacyMigrationReembedsChangedContentAndRemovesSource(t *testing.T) { + projectDir := t.TempDir() + oldContent := []byte("package demo\n\nfunc Before() {}\n") + newContent := []byte("package demo\n\nfunc After() {}\n") + if err := os.WriteFile(filepath.Join(projectDir, "main.go"), newContent, 0o644); err != nil { + t.Fatal(err) + } + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexerForProject(filepath.Join(t.TempDir(), "shared.db"), emb, 512, "int8", projectDir) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + legacyPath := filepath.Join(t.TempDir(), "index.db") + legacy, err := store.New(legacyPath, 4) + if err != nil { + t.Fatal(err) + } + oldHash := sha256.Sum256(oldContent) + if err := legacy.UpsertFile("main.go", hex.EncodeToString(oldHash[:])); err != nil { + t.Fatal(err) + } + chunks, err := idx.chunker.Chunk("main.go", oldContent) + if err != nil { + t.Fatal(err) + } + vectors := make([][]float32, len(chunks)) + for i := range vectors { + vectors[i] = []float32{1, 0, 0, 0} + } + if err := legacy.InsertChunks(chunks, vectors); err != nil { + t.Fatal(err) + } + if err := legacy.Close(); err != nil { + t.Fatal(err) + } + if err := idx.PrepareLegacyMigration(projectDir, legacyPath); err != nil { + t.Fatal(err) + } + if _, err := idx.Index(context.Background(), projectDir, false, nil); err != nil { + t.Fatal(err) + } + if emb.callCount == 0 { + t.Fatal("changed content should be embedded") + } + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("legacy database should be removed after successful reindex, stat err=%v", err) + } +} diff --git a/internal/index/shared.go b/internal/index/shared.go index 2e78f561..9c42e925 100644 --- a/internal/index/shared.go +++ b/internal/index/shared.go @@ -11,6 +11,7 @@ package index import ( "context" "crypto/sha256" + "errors" "fmt" "os" "path/filepath" @@ -18,6 +19,7 @@ import ( "strconv" "time" + "github.com/ory/lumen/internal/chunker" "github.com/ory/lumen/internal/merkle" "github.com/ory/lumen/internal/store" ) @@ -75,6 +77,108 @@ func (idx *Indexer) indexSharedWithTree(ctx context.Context, projectDir, _ strin progress(0, len(filesToIndex), fmt.Sprintf("Found %d files to index", len(filesToIndex))) } + type pendingSharedFile struct { + relativePath string + contentHash string + chunks []chunker.Chunk + vectors map[int][]float32 + remaining int + fileIndex int + skipped bool + } + type chunkRef struct { + file *pendingSharedFile + position int + } + + embedBatchSize := idx.embedBatchSize + if embedBatchSize <= 0 { + embedBatchSize = 256 + } + var pending []*pendingSharedFile + var batchTexts []string + var batchRefs []chunkRef + + embedAllFileChunks := func(file *pendingSharedFile) (map[int][]float32, error) { + vectors := make(map[int][]float32, len(file.chunks)) + for start := 0; start < len(file.chunks); start += embedBatchSize { + end := min(start+embedBatchSize, len(file.chunks)) + texts := make([]string, end-start) + for i := start; i < end; i++ { + texts[i-start] = store.EmbeddingInput(file.chunks[i]) + } + embedded, err := idx.emb.Embed(ctx, texts) + if err != nil { + return nil, fmt.Errorf("re-embed %s: %w", file.relativePath, err) + } + if len(embedded) != len(texts) { + return nil, fmt.Errorf("re-embed %s returned %d vectors for %d inputs", file.relativePath, len(embedded), len(texts)) + } + for i, vector := range embedded { + vectors[start+i] = vector + } + } + return vectors, nil + } + + storeFile := func(file *pendingSharedFile) (bool, error) { + created, err := idx.store.StoreFileRevision(file.relativePath, file.contentHash, file.chunks, file.vectors) + if !errors.Is(err, store.ErrVectorVanished) { + return created, err + } + // A concurrent cleanup can remove a vector after MissingChunkInputs + // observes it. Re-embedding every position makes the retry independent + // of all shared vector rows and closes that TOCTOU window. + vectors, embedErr := embedAllFileChunks(file) + if embedErr != nil { + return false, embedErr + } + created, err = idx.store.StoreFileRevision(file.relativePath, file.contentHash, file.chunks, vectors) + return created, err + } + + drainReady := func() error { + for len(pending) > 0 && pending[0].remaining == 0 { + file := pending[0] + created, err := storeFile(file) + if err != nil { + return fmt.Errorf("store file revision %s: %w", file.relativePath, err) + } + if created || force { + stats.ChunksCreated += len(file.chunks) + } + if !file.skipped { + stats.IndexedFiles++ + } + pending = pending[1:] + } + return nil + } + + flushBatch := func() error { + if len(batchTexts) == 0 { + return drainReady() + } + embedded, err := idx.emb.Embed(ctx, batchTexts) + if err != nil { + return fmt.Errorf("embed shared batch: %w", err) + } + if len(embedded) != len(batchRefs) { + return fmt.Errorf("embed shared batch returned %d vectors for %d inputs", len(embedded), len(batchRefs)) + } + for i, ref := range batchRefs { + ref.file.vectors[ref.position] = embedded[i] + ref.file.remaining-- + } + if progress != nil { + last := batchRefs[len(batchRefs)-1].file + progress(last.fileIndex+1, len(filesToIndex), fmt.Sprintf("Embedded %d shared chunks", len(batchRefs))) + } + batchTexts = batchTexts[:0] + batchRefs = batchRefs[:0] + return drainReady() + } + for fileIndex, relativePath := range filesToIndex { if err := ctx.Err(); err != nil { return stats, err @@ -115,8 +219,15 @@ func (idx *Indexer) indexSharedWithTree(ctx context.Context, projectDir, _ strin idx.logger.Warn("skipping unchunkable file", "path", relativePath, "error", err) } stats.FilesSkipped++ - if _, err := idx.store.StoreFileRevision(relativePath, contentHash, nil, nil); err != nil { - return stats, fmt.Errorf("record skipped file %s: %w", relativePath, err) + pending = append(pending, &pendingSharedFile{ + relativePath: relativePath, + contentHash: contentHash, + vectors: map[int][]float32{}, + fileIndex: fileIndex, + skipped: true, + }) + if err := drainReady(); err != nil { + return stats, err } continue } @@ -136,65 +247,64 @@ func (idx *Indexer) indexSharedWithTree(ctx context.Context, projectDir, _ strin return stats, fmt.Errorf("check shared vectors for %s: %w", relativePath, err) } } - vectors := make(map[int][]float32, len(missing)) - const embedBatchSize = 256 + file := &pendingSharedFile{ + relativePath: relativePath, + contentHash: contentHash, + chunks: chunks, + vectors: make(map[int][]float32, len(missing)), + fileIndex: fileIndex, + } + pending = append(pending, file) var needsEmbedding []int for _, position := range missing { h := sha256.Sum256([]byte(store.EmbeddingInput(chunks[position]))) if vector, ok := idx.legacyVectors[h]; ok { - vectors[position] = vector + file.vectors[position] = vector } else { needsEmbedding = append(needsEmbedding, position) } } - for start := 0; start < len(needsEmbedding); start += embedBatchSize { - end := min(start+embedBatchSize, len(needsEmbedding)) - positions := needsEmbedding[start:end] - texts := make([]string, len(positions)) - for i, position := range positions { - texts[i] = store.EmbeddingInput(chunks[position]) - } - embedded, err := idx.emb.Embed(ctx, texts) - if err != nil { - return stats, fmt.Errorf("embed %s: %w", relativePath, err) - } - if len(embedded) != len(positions) { - return stats, fmt.Errorf("embed %s returned %d vectors for %d inputs", relativePath, len(embedded), len(positions)) - } - for i, position := range positions { - vectors[position] = embedded[i] - } - if progress != nil { - progress(fileIndex+1, len(filesToIndex), fmt.Sprintf("Embedded %d chunks for %s", len(positions), relativePath)) + file.remaining = len(needsEmbedding) + for _, position := range needsEmbedding { + batchTexts = append(batchTexts, store.EmbeddingInput(chunks[position])) + batchRefs = append(batchRefs, chunkRef{file: file, position: position}) + if len(batchTexts) == embedBatchSize { + if err := flushBatch(); err != nil { + return stats, err + } } } - created, err := idx.store.StoreFileRevision(relativePath, contentHash, chunks, vectors) - if err != nil { - return stats, fmt.Errorf("store file revision %s: %w", relativePath, err) - } - if created || force { - stats.ChunksCreated += len(chunks) + if err := drainReady(); err != nil { + return stats, err } - stats.IndexedFiles++ + } + if err := flushBatch(); err != nil { + return stats, err + } + if len(pending) != 0 { + return stats, fmt.Errorf("shared embedding queue left %d files pending", len(pending)) } if len(filesToIndex) > 0 { idx.store.Analyze() } - if err := idx.store.SetMeta("root_hash", curTree.RootHash); err != nil { - return stats, err + metadata := []struct{ key, value string }{ + {"embedding_model", idx.emb.ModelName()}, + {"project_path", projectDir}, + {"last_indexed_at", time.Now().UTC().Format(time.RFC3339)}, + {"total_files", strconv.Itoa(stats.TotalFiles)}, + {"vector_storage", idx.vectorStorage}, } - for key, value := range map[string]string{ - "embedding_model": idx.emb.ModelName(), - "project_path": projectDir, - "last_indexed_at": time.Now().UTC().Format(time.RFC3339), - "total_files": strconv.Itoa(stats.TotalFiles), - "vector_storage": idx.vectorStorage, - } { - if err := idx.store.SetMeta(key, value); err != nil { - return stats, fmt.Errorf("store %s metadata: %w", key, err) + for _, item := range metadata { + if err := idx.store.SetMeta(item.key, item.value); err != nil { + return stats, fmt.Errorf("store %s metadata: %w", item.key, err) } } + // root_hash is the commit marker and must be written after every other + // durable file, vector, and metadata update. + if err := idx.store.SetMeta("root_hash", curTree.RootHash); err != nil { + return stats, fmt.Errorf("store root_hash metadata: %w", err) + } if progress != nil && len(filesToIndex) > 0 { progress(len(filesToIndex), len(filesToIndex), fmt.Sprintf("Indexing complete: %d files, %d new chunks", len(filesToIndex), stats.ChunksCreated)) } diff --git a/internal/index/shared_batch_test.go b/internal/index/shared_batch_test.go new file mode 100644 index 00000000..8c237238 --- /dev/null +++ b/internal/index/shared_batch_test.go @@ -0,0 +1,170 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. + +package index + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/ory/lumen/internal/chunker" + "github.com/ory/lumen/internal/store" +) + +func writeSharedBatchFixture(t *testing.T, projectDir string, files int) { + t.Helper() + for i := 0; i < files; i++ { + name := fmt.Sprintf("file_%d.go", i) + writeGoFile(t, projectDir, name, fmt.Sprintf("package demo\n\nfunc File%d() {}\n", i)) + } +} + +func TestSharedIndexBatchesEmbeddingsAcrossFiles(t *testing.T) { + projectDir := t.TempDir() + writeSharedBatchFixture(t, projectDir, 5) + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexerForProject(filepath.Join(t.TempDir(), "index.db"), emb, 512, "int8", projectDir) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + + stats, err := idx.Index(context.Background(), projectDir, false, nil) + if err != nil { + t.Fatal(err) + } + if stats.IndexedFiles != 5 { + t.Fatalf("IndexedFiles = %d, want 5", stats.IndexedFiles) + } + if emb.callCount != 1 { + t.Fatalf("Embed calls = %d, want 1 cross-file batch", emb.callCount) + } +} + +type failSecondBatchEmbedder struct { + calls int +} + +func (e *failSecondBatchEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) { + e.calls++ + if e.calls == 2 { + return nil, errors.New("injected batch failure") + } + vectors := make([][]float32, len(texts)) + for i := range vectors { + vectors[i] = []float32{1, 0, 0, 0} + } + return vectors, nil +} + +func (*failSecondBatchEmbedder) Dimensions() int { return 4 } +func (*failSecondBatchEmbedder) ModelName() string { return "test-model" } + +func TestSharedIndexBatchFailureDoesNotCommitRootHash(t *testing.T) { + projectDir := t.TempDir() + writeSharedBatchFixture(t, projectDir, 5) + emb := &failSecondBatchEmbedder{} + idx, err := NewIndexerForProject(filepath.Join(t.TempDir(), "index.db"), emb, 512, "int8", projectDir) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + idx.embedBatchSize = 2 + + if _, err := idx.Index(context.Background(), projectDir, false, nil); err == nil { + t.Fatal("expected injected embedding failure") + } + if _, err := idx.store.GetMeta("root_hash"); !errors.Is(err, sql.ErrNoRows) { + t.Fatalf("root_hash error = %v, want sql.ErrNoRows", err) + } + hashes, err := idx.store.GetFileHashes() + if err != nil { + t.Fatal(err) + } + if len(hashes) != 2 { + t.Fatalf("durable revisions = %d, want first flushed batch of 2", len(hashes)) + } +} + +type fixedChunker struct { + chunks []chunker.Chunk +} + +func (c fixedChunker) Chunk(string, []byte) ([]chunker.Chunk, error) { + return append([]chunker.Chunk(nil), c.chunks...), nil +} + +type callbackEmbedder struct { + calls int + callback func() error +} + +func (e *callbackEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) { + e.calls++ + if e.callback != nil { + callback := e.callback + e.callback = nil + if err := callback(); err != nil { + return nil, err + } + } + vectors := make([][]float32, len(texts)) + for i := range vectors { + vectors[i] = []float32{1, 0, 0, 0} + } + return vectors, nil +} + +func (*callbackEmbedder) Dimensions() int { return 4 } +func (*callbackEmbedder) ModelName() string { return "test-model" } + +func TestSharedIndexRetriesWhenPreviouslySharedVectorVanishes(t *testing.T) { + projectA, projectB := t.TempDir(), t.TempDir() + if err := os.WriteFile(filepath.Join(projectA, "main.go"), []byte("package demo\n"), 0o644); err != nil { + t.Fatal(err) + } + chunkA := chunker.Chunk{ID: "a", FilePath: "main.go", Symbol: "A", Kind: "function", StartLine: 1, EndLine: 10, Content: strings.Repeat("shared input ", 80)} + chunkB := chunker.Chunk{ID: "b", FilePath: "main.go", Symbol: "B", Kind: "function", StartLine: 11, EndLine: 20, Content: strings.Repeat("new input ", 80)} + dbPath := filepath.Join(t.TempDir(), "index.db") + emb := &callbackEmbedder{} + idx, err := NewIndexerForProject(dbPath, emb, 512, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + idx.chunker = fixedChunker{chunks: []chunker.Chunk{chunkA, chunkB}} + if _, err := idx.store.StoreFileRevision("main.go", "old", []chunker.Chunk{chunkA}, map[int][]float32{0: {1, 0, 0, 0}}); err != nil { + t.Fatal(err) + } + keeper, err := store.NewCollection(dbPath, 4, "int8", projectB) + if err != nil { + t.Fatal(err) + } + defer func() { _ = keeper.Close() }() + if attached, err := keeper.AttachExistingFileRevision("main.go", "old"); err != nil || !attached { + t.Fatalf("attach keeper: %v, %v", attached, err) + } + if err := idx.store.DeleteFileChunks("main.go"); err != nil { + t.Fatal(err) + } + emb.callback = func() error { return keeper.DeleteFileChunks("main.go") } + + stats, err := idx.Index(context.Background(), projectA, false, nil) + if err != nil { + t.Fatal(err) + } + if emb.calls != 2 { + t.Fatalf("Embed calls = %d, want initial missing batch plus full-file retry", emb.calls) + } + if stats.ChunksCreated != 2 { + t.Fatalf("ChunksCreated = %d, want 2", stats.ChunksCreated) + } +} From 777d71a56848813811ac170d8b8f4eb1d9abc7a9 Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Sat, 8 Aug 2026 22:46:34 +0200 Subject: [PATCH 5/9] fix(stdio): harden background index maintenance --- cmd/clean.go | 35 ++++++++++++++---- cmd/clean_test.go | 36 +++++++++++++++++++ cmd/index.go | 6 ++-- cmd/stdio.go | 54 +++++++++++++++++++++------- cmd/stdio_test.go | 80 +++++++++++++++++++++++++++++++++++++++-- docs/INDEX_STORAGE.md | 5 +-- skills/reindex/SKILL.md | 6 ++-- 7 files changed, 194 insertions(+), 28 deletions(-) diff --git a/cmd/clean.go b/cmd/clean.go index 444f6ff8..c9caf6d1 100644 --- a/cmd/clean.go +++ b/cmd/clean.go @@ -15,10 +15,13 @@ package cmd import ( + "bytes" "fmt" "io" + "log/slog" "os" "path/filepath" + "strings" "time" "github.com/ory/lumen/internal/config" @@ -39,6 +42,7 @@ const dailyCleanupInterval = 24 * time.Hour var ( removeIndexDir = os.RemoveAll cleanupCollectionAt = store.CleanupCollectionAt + tryAcquireExclusive = indexlock.TryAcquire ) func init() { @@ -90,8 +94,11 @@ func runClean(cmd *cobra.Command, _ []string) error { } // cleanIndexes removes every stale index directory directly under dataDir, -// reporting each decision on stderr and a summary on stdout. now is injected so -// the age cutoff is testable. Failures to remove a single directory are +// reporting each decision on the injected stderr and a summary on the injected +// stdout. The injected writers deliberately keep this reusable by both the +// interactive CLI and the MCP background cleanup without mutating pterm's +// process-global state. now is injected so the age cutoff is testable. Failures +// to remove a single directory are // reported and the sweep continues; the first such failure is returned once // every directory has been considered. func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.Time) error { @@ -144,8 +151,12 @@ func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.T // exclusive collection lock for the entire database cleanup and removal. func cleanIndex(stderr io.Writer, name, hashDir string, days int, cutoff time.Time) (bool, store.CleanupStats, error) { dbPath := filepath.Join(hashDir, "index.db") - lock, lockErr := indexlock.TryAcquire(indexlock.LockPathForDB(dbPath)) - if lockErr != nil || lock == nil { + lock, lockErr := tryAcquireExclusive(indexlock.LockPathForDB(dbPath)) + if lockErr != nil { + _, _ = fmt.Fprintf(stderr, "Failed to acquire index lock for %s: %v\n", name, lockErr) + return false, store.CleanupStats{}, fmt.Errorf("acquire index lock for %s: %w", name, lockErr) + } + if lock == nil { _, _ = fmt.Fprintf(stderr, "Keeping %s: an indexer is currently running.\n", name) return false, store.CleanupStats{}, nil } @@ -249,16 +260,26 @@ func pluralY(n int) string { // runDailyCleanup performs the MCP-startup maintenance sweep at most once per // day. The stamp is deliberately outside collection directories so it is not // mistaken for an index by cleanIndexes. -func runDailyCleanup(dataDir string, now time.Time) { +func runDailyCleanup(dataDir string, now time.Time, logger *slog.Logger) { + if logger == nil { + logger = slog.New(slog.NewTextHandler(io.Discard, nil)) + } stampPath := filepath.Join(dataDir, ".last-cleanup") if info, err := os.Stat(stampPath); err == nil && now.Sub(info.ModTime()) < dailyCleanupInterval { + logger.Debug("daily cleanup skipped: stamp is fresh", "stamp_path", stampPath) return } if err := os.MkdirAll(dataDir, 0o755); err != nil { + logger.Warn("daily cleanup: create data directory", "path", dataDir, "error", err) return } - if err := cleanIndexes(io.Discard, io.Discard, dataDir, defaultCleanDays, now); err != nil { + var stderr, stdout bytes.Buffer + if err := cleanIndexes(&stderr, &stdout, dataDir, defaultCleanDays, now); err != nil { + logger.Warn("daily cleanup failed", "error", err, "details", strings.TrimSpace(stderr.String())) return } - _ = os.WriteFile(stampPath, []byte(now.UTC().Format(time.RFC3339)), 0o600) + logger.Info("daily cleanup complete", "summary", strings.TrimSpace(stdout.String()), "details", strings.TrimSpace(stderr.String())) + if err := os.WriteFile(stampPath, []byte(now.UTC().Format(time.RFC3339)), 0o600); err != nil { + logger.Warn("daily cleanup: write stamp", "path", stampPath, "error", err) + } } diff --git a/cmd/clean_test.go b/cmd/clean_test.go index b0aaee1b..aa538803 100644 --- a/cmd/clean_test.go +++ b/cmd/clean_test.go @@ -18,8 +18,10 @@ import ( "bytes" "database/sql" "errors" + "log/slog" "os" "path/filepath" + "strings" "testing" "time" @@ -307,6 +309,40 @@ func TestClean_LeavesNonIndexFilesAlone(t *testing.T) { assert.FileExists(t, logPath, "debug.log must not be removed") } +func TestCleanIndexReportsLockAcquisitionErrors(t *testing.T) { + original := tryAcquireExclusive + t.Cleanup(func() { tryAcquireExclusive = original }) + tryAcquireExclusive = func(string) (*indexlock.Lock, error) { + return nil, errors.New("permission denied") + } + var stderr bytes.Buffer + removed, _, err := cleanIndex(&stderr, "abc", t.TempDir(), 30, time.Now()) + if err == nil || removed { + t.Fatalf("removed=%v err=%v", removed, err) + } + if !strings.Contains(stderr.String(), "Failed to acquire index lock") || strings.Contains(stderr.String(), "currently running") { + t.Fatalf("unexpected stderr: %s", stderr.String()) + } +} + +func TestRunDailyCleanupUsesProvidedLoggerAndStampsSuccess(t *testing.T) { + dataDir := t.TempDir() + var logs bytes.Buffer + logger := slog.New(slog.NewTextHandler(&logs, nil)) + now := time.Date(2026, 8, 8, 12, 0, 0, 0, time.UTC) + runDailyCleanup(dataDir, now, logger) + stamp, err := os.ReadFile(filepath.Join(dataDir, ".last-cleanup")) + if err != nil { + t.Fatal(err) + } + if string(stamp) != now.Format(time.RFC3339) { + t.Fatalf("stamp = %q, want %q", stamp, now.Format(time.RFC3339)) + } + if !strings.Contains(logs.String(), "daily cleanup complete") { + t.Fatalf("provided logger did not receive cleanup summary: %s", logs.String()) + } +} + // TestClean_HandlesMultipleModelsPerProject verifies each model's index is aged // independently, since switching models creates a separate index directory. func TestClean_HandlesMultipleModelsPerProject(t *testing.T) { diff --git a/cmd/index.go b/cmd/index.go index 8241dcc2..bcd4fc83 100644 --- a/cmd/index.go +++ b/cmd/index.go @@ -237,8 +237,10 @@ func setupIndexerForProject(cfg *config.ConfigService, emb *embedder.FailoverEmb idx.SetLogger(logger) if projectPath != "" { legacyPath := config.LegacyDBPathForProject(projectPath, emb.ModelName()) - if err := idx.PrepareLegacyMigration(projectPath, legacyPath); err != nil && logger != nil { - logger.Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + if err := idx.PrepareLegacyMigration(projectPath, legacyPath); err != nil { + if logger != nil { + logger.Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + } } } return idx, nil diff --git a/cmd/stdio.go b/cmd/stdio.go index c4ca1e13..29077cb9 100644 --- a/cmd/stdio.go +++ b/cmd/stdio.go @@ -98,7 +98,7 @@ type IndexStatusOutput struct { TotalChunks int `json:"total_chunks"` UniqueVectors int `json:"unique_vectors"` SharedReferences int `json:"shared_references"` - DeduplicationRate float64 `json:"deduplication_ratio"` + DeduplicationRate float64 `json:"deduplication_rate"` VectorStorage string `json:"vector_storage"` DatabaseBytes int64 `json:"database_bytes"` ReclaimableBytes int64 `json:"reclaimable_bytes"` @@ -151,6 +151,14 @@ const backgroundReindexMaxDuration = 10 * time.Minute // is identical across all four code paths. const staleIndexWarning = "Index is being updated in the background. Results may be incomplete or outdated. Use grep/glob/find for code search until indexing finishes (usually a few minutes; longer for large repositories)." +var ( + tryAcquire = indexlock.TryAcquire + tryAcquireShared = indexlock.TryAcquireShared + prepareMigrationFunc = func(idx *index.Indexer, projectDir, legacyPath string) error { + return idx.PrepareLegacyMigration(projectDir, legacyPath) + } +) + type cacheEntry struct { idx *index.Indexer effectiveRoot string @@ -235,6 +243,20 @@ func (ic *indexerCache) getFreshnessTTL() time.Duration { return defaultFreshnessTTL } +func (ic *indexerCache) maxChunkTokens() int { + if ic.cfg == nil { + return 0 + } + return ic.cfg.MaxChunkTokens() +} + +func (ic *indexerCache) vectorStorage() string { + if ic.cfg == nil { + return "int8" + } + return ic.cfg.VectorStorage() +} + // getReindexTimeout returns the effective reindex timeout, checking the override // field first, then cfg, then the default constant. func (ic *indexerCache) getReindexTimeout() time.Duration { @@ -516,21 +538,16 @@ func (ic *indexerCache) getOrCreate(projectPath string, preferredRoot string, mo seed: ic.seedFunc, }) - idx, err := index.NewIndexerForProject(dbPath, ic.embedder, ic.cfg.MaxChunkTokens(), ic.cfg.VectorStorage(), effectiveRoot) + idx, err := index.NewIndexerForProject(dbPath, ic.embedder, ic.maxChunkTokens(), ic.vectorStorage(), effectiveRoot) if err != nil { return nil, "", "", fmt.Errorf("create indexer: %w", err) } idx.SetLogger(ic.logger()) - legacyPath := config.LegacyDBPathForProject(effectiveRoot, modelName) - if err := idx.PrepareLegacyMigration(effectiveRoot, legacyPath); err != nil { - ic.logger().Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) - } - // Pre-populate the freshness TTL if the index was recently stamped by // background pre-warming (SessionStart hook). This avoids a redundant // merkle walk on the very first search in a new session. entry := cacheEntry{idx: idx, effectiveRoot: effectiveRoot, model: modelName} - if lastAt, ok := idx.LastIndexedAt(); ok { + if lastAt, ok := idx.LastIndexedAt(effectiveRoot); ok { ttl := ic.getFreshnessTTL() if time.Since(lastAt) < ttl { entry.lastCheckedAt = lastAt @@ -834,7 +851,7 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn ic.mu.Unlock() }() - collectionLock, lockErr := indexlock.TryAcquireShared(collectionLockPath) + collectionLock, lockErr := tryAcquireShared(collectionLockPath) if lockErr != nil { ic.logger().Warn("background reindex: failed to acquire lock", "project", projectDir, "err", lockErr) done <- freshResult{skipped: true} @@ -847,8 +864,14 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn return } defer collectionLock.Release() - projectLock, lockErr := indexlock.TryAcquire(projectLockPath) - if lockErr != nil || projectLock == nil { + projectLock, lockErr := tryAcquire(projectLockPath) + if lockErr != nil { + ic.logger().Warn("background reindex: failed to acquire project lock", "project", projectDir, "err", lockErr) + done <- freshResult{skipped: true} + return + } + if projectLock == nil { + ic.logger().Debug("background reindex: project lock held by another process, skipping", "project", projectDir) done <- freshResult{skipped: true} return } @@ -857,7 +880,7 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn // If a recent external process (e.g. lumen index from SessionStart) // already updated the index within freshnessTTL, trust the DB timestamp // and skip the expensive merkle tree walk. - if lastAt, ok := idx.LastIndexedAt(); ok { + if lastAt, ok := idx.LastIndexedAt(projectDir); ok { ttl := ic.getFreshnessTTL() if ttl == 0 { ttl = defaultFreshnessTTL @@ -873,6 +896,11 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn } } + legacyPath := config.LegacyDBPathForProject(projectDir, modelName) + if err := prepareMigrationFunc(idx, projectDir, legacyPath); err != nil { + ic.logger().Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + } + ensureFresh := ic.ensureFreshFunc if ensureFresh == nil { ensureFresh = func(ctx context.Context, idx *index.Indexer, dir string, p index.ProgressFunc) (bool, index.Stats, error) { @@ -1446,7 +1474,7 @@ func runStdio(_ *cobra.Command, _ []string) error { "backend", cfg.Servers()[0].Backend, "freshness_ttl", cfg.FreshnessTTL().String(), ) - runDailyCleanup(filepath.Join(config.XDGDataDir(), "lumen"), time.Now()) + runDailyCleanup(filepath.Join(config.XDGDataDir(), "lumen"), time.Now(), logger) closeCtx, closeFn := context.WithCancel(context.Background()) indexers := &indexerCache{ diff --git a/cmd/stdio_test.go b/cmd/stdio_test.go index e7469e67..975b1975 100644 --- a/cmd/stdio_test.go +++ b/cmd/stdio_test.go @@ -16,6 +16,7 @@ package cmd import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -46,6 +47,16 @@ var ( discardLog = slog.New(slog.NewTextHandler(io.Discard, nil)) ) +func TestIndexStatusOutputDeduplicationRateJSONTag(t *testing.T) { + data, err := json.Marshal(IndexStatusOutput{DeduplicationRate: 0.5}) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(data), `"deduplication_rate":0.5`) || strings.Contains(string(data), "deduplication_ratio") { + t.Fatalf("unexpected JSON: %s", data) + } +} + // assertGolden compares got against the golden file at path. If -update-golden // is set, it writes got to the golden file instead. func assertGolden(t *testing.T, goldenPath, got string) { @@ -1382,6 +1393,64 @@ func TestGetOrCreate_SeedCancelledByClose(t *testing.T) { } } +func TestGetOrCreateSupportsNilConfig(t *testing.T) { + t.Setenv("XDG_DATA_HOME", t.TempDir()) + projectDir := t.TempDir() + ic := &indexerCache{embedder: &stubEmbedder{}, log: discardLog} + idx, _, _, err := ic.getOrCreate(projectDir, "") + if err != nil { + t.Fatal(err) + } + if ic.maxChunkTokens() != 0 || ic.vectorStorage() != "int8" { + t.Fatalf("nil-config defaults = %d, %q", ic.maxChunkTokens(), ic.vectorStorage()) + } + if err := idx.Close(); err != nil { + t.Fatal(err) + } +} + +func TestLegacyMigrationPreparationRunsInBackgroundIndexing(t *testing.T) { + t.Setenv("XDG_DATA_HOME", t.TempDir()) + projectDir := t.TempDir() + cfg := newTestConfigService(t, 512) + original := prepareMigrationFunc + t.Cleanup(func() { prepareMigrationFunc = original }) + prepareCalls := 0 + prepareMigrationFunc = func(_ *index.Indexer, gotProject, _ string) error { + prepareCalls++ + if gotProject != projectDir { + t.Fatalf("project = %q, want %q", gotProject, projectDir) + } + return nil + } + ic := &indexerCache{ + embedder: &stubEmbedder{}, + cfg: cfg, + log: discardLog, + ensureFreshFunc: func(_ context.Context, _ *index.Indexer, _ string, _ index.ProgressFunc) (bool, index.Stats, error) { + if prepareCalls != 1 { + t.Fatalf("PrepareLegacyMigration calls before EnsureFresh = %d, want 1", prepareCalls) + } + return false, index.Stats{}, nil + }, + } + idx, effectiveRoot, _, err := ic.getOrCreate(projectDir, "") + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + if prepareCalls != 0 { + t.Fatalf("PrepareLegacyMigration ran during getOrCreate: %d calls", prepareCalls) + } + input := SemanticSearchInput{Cwd: projectDir, Path: projectDir, Query: "test"} + if _, err := ic.ensureIndexed(idx, input, effectiveRoot, ic.dbPath(effectiveRoot, "stub"), nil); err != nil { + t.Fatal(err) + } + if prepareCalls != 1 { + t.Fatalf("PrepareLegacyMigration calls = %d, want 1", prepareCalls) + } +} + func TestFormatSearchResults_IncludesSeedWarning(t *testing.T) { out := SemanticSearchOutput{ Results: nil, @@ -1627,7 +1696,7 @@ func TestEnsureIndexed_SkipsMerkleWalkWhenRecentlyIndexedExternally(t *testing.T tmpDir := t.TempDir() dbPath := filepath.Join(tmpDir, "test.db") - idx, err := index.NewIndexer(dbPath, &stubEmbedder{}, 512) + idx, err := index.NewIndexerForProject(dbPath, &stubEmbedder{}, 512, "int8", tmpDir) if err != nil { t.Fatal(err) } @@ -1635,7 +1704,14 @@ func TestEnsureIndexed_SkipsMerkleWalkWhenRecentlyIndexedExternally(t *testing.T // Simulate an external process (e.g. lumen index from SessionStart) having // recently written last_indexed_at to the DB. - if err := writeDBWithLastIndexedAt(t, dbPath, time.Now().Add(-5*time.Second)); err != nil { + external, err := store.NewCollection(dbPath, 4, "int8", tmpDir) + if err != nil { + t.Fatal(err) + } + if err := external.SetMeta("last_indexed_at", time.Now().Add(-5*time.Second).UTC().Format(time.RFC3339)); err != nil { + t.Fatal(err) + } + if err := external.Close(); err != nil { t.Fatal(err) } diff --git a/docs/INDEX_STORAGE.md b/docs/INDEX_STORAGE.md index 586b90b9..76477e0a 100644 --- a/docs/INDEX_STORAGE.md +++ b/docs/INDEX_STORAGE.md @@ -162,13 +162,14 @@ Use the narrowest operation that matches the problem: lumen index . # refresh changed files lumen index --force . # reprocess every file in the current project lumen clean # reclaim stale memberships and unreferenced data -lumen clean --days 0 # wipe every cached index on the host +lumen clean --days 0 # wipe every cached index not held by an active indexer ``` `--force` does not wipe other worktrees from a shared collection. It rebuilds the current project's memberships and chunk definitions while the collection continues to deduplicate physical vectors. Use the full wipe only when you -intend to rebuild all Lumen indexes on the machine. +intend to rebuild all Lumen indexes on the machine. Collections held by an +active indexer lock are kept and counted as skipped. To delete indexes manually, stop active Lumen indexers and remove the Lumen data directory. No source-tree files are stored there, and no files are added to the diff --git a/skills/reindex/SKILL.md b/skills/reindex/SKILL.md index 2f9988e1..b36ca6fc 100644 --- a/skills/reindex/SKILL.md +++ b/skills/reindex/SKILL.md @@ -19,8 +19,10 @@ Refresh or rebuild the bundled Lumen index for the current project. run one via the shell: - `lumen index --force .` — reprocesses every file for the current project without wiping other worktrees or their shared vectors. Prefer this. - - `lumen clean --days 0 && lumen index .` — deletes every cached index on the - host before rebuilding. Use only when the user asks for a full wipe. + - `lumen clean --days 0 && lumen index .` — deletes every cached index not + held by an active indexer lock before rebuilding; locked collections are + skipped and counted in the summary. Use only when the user asks for a full + wipe. - `lumen clean` — removes indexes for projects that no longer exist or have not been used in 30 days. Use to reclaim disk space, not to rebuild the current project. From 048e96d34932bb46fa6c7fcf8a0aba260ab44b4a Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Sat, 8 Aug 2026 23:15:05 +0200 Subject: [PATCH 6/9] fix shared collection E2E assertions --- e2e_cli_test.go | 225 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 162 insertions(+), 63 deletions(-) diff --git a/e2e_cli_test.go b/e2e_cli_test.go index 4d9e8bc4..23698a74 100644 --- a/e2e_cli_test.go +++ b/e2e_cli_test.go @@ -79,6 +79,7 @@ func runCLIWithDataHome(t *testing.T, dataHome string, args ...string) (stdout, "OLLAMA_HOST=" + ollamaHost, "LUMEN_EMBED_MODEL=all-minilm", "XDG_DATA_HOME=" + dataHome, + "XDG_CONFIG_HOME=" + filepath.Join(dataHome, "config"), "HOME=" + os.Getenv("HOME"), "PATH=" + os.Getenv("PATH"), } @@ -136,8 +137,14 @@ func TestE2E_CLI_IndexForceReindex(t *testing.T) { } } -// openIndexDB opens the SQLite index database for a given project path and dataHome. -func openIndexDB(t *testing.T, dataHome, projectPath string) *sql.DB { +type projectIndexDB struct { + *sql.DB + projectID int64 +} + +// openIndexDB opens the shared SQLite collection and resolves the membership +// for projectPath so raw SQL assertions remain project-local. +func openIndexDB(t *testing.T, dataHome, projectPath string) *projectIndexDB { t.Helper() if err := sqlite_vec.Auto(); err != nil { t.Fatal(err) @@ -148,7 +155,15 @@ func openIndexDB(t *testing.T, dataHome, projectPath string) *sql.DB { t.Fatalf("open index db: %v", err) } t.Cleanup(func() { db.Close() }) - return db + absoluteProjectPath, err := filepath.Abs(projectPath) + if err != nil { + t.Fatalf("resolve project path: %v", err) + } + var projectID int64 + if err := db.QueryRow("SELECT id FROM projects WHERE path = ?", filepath.Clean(absoluteProjectPath)).Scan(&projectID); err != nil { + t.Fatalf("resolve project membership: %v", err) + } + return &projectIndexDB{DB: db, projectID: projectID} } func TestE2E_CLI_SQLVerifySchema(t *testing.T) { @@ -165,7 +180,10 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { db := openIndexDB(t, dataHome, projectPath) // --- Verify tables exist --- - expectedTables := []string{"files", "chunks", "project_meta", "vec_chunks"} + expectedTables := []string{ + "collection_meta", "projects", "project_meta", "file_revisions", + "project_files", "vector_keys", "chunk_defs", "vec_vectors", + } for _, table := range expectedTables { var count int err := db.QueryRow("SELECT count(*) FROM sqlite_master WHERE name = ?", table).Scan(&count) @@ -177,17 +195,23 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { } } - // --- Verify files table --- + // --- Verify project files --- var fileCount int - if err := db.QueryRow("SELECT count(*) FROM files").Scan(&fileCount); err != nil { - t.Fatalf("count files: %v", err) + if err := db.QueryRow("SELECT count(*) FROM project_files WHERE project_id = ?", db.projectID).Scan(&fileCount); err != nil { + t.Fatalf("count project files: %v", err) } if fileCount != 7 { t.Errorf("expected 7 files, got %d", fileCount) } // All file paths should end in .go, .svelte, or .swift and have valid hashes. - rows, err := db.Query("SELECT path, hash FROM files ORDER BY path") + rows, err := db.Query(` + SELECT pf.relative_path, hex(fr.content_hash) + FROM project_files pf + JOIN file_revisions fr ON fr.id = pf.file_revision_id + WHERE pf.project_id = ? + ORDER BY pf.relative_path + `, db.projectID) if err != nil { t.Fatalf("query files: %v", err) } @@ -215,10 +239,15 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { } } - // --- Verify chunks table --- + // --- Verify project chunks --- var chunkCount int - if err := db.QueryRow("SELECT count(*) FROM chunks").Scan(&chunkCount); err != nil { - t.Fatalf("count chunks: %v", err) + if err := db.QueryRow(` + SELECT count(*) + FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&chunkCount); err != nil { + t.Fatalf("count project chunks: %v", err) } if chunkCount < 15 { t.Errorf("expected at least 15 chunks (fixture has ~20), got %d", chunkCount) @@ -226,11 +255,12 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { // Every chunk should reference a valid file and have sensible fields. chunkRows, err := db.Query(` - SELECT c.id, c.file_path, c.symbol, c.kind, c.start_line, c.end_line - FROM chunks c - JOIN files f ON c.file_path = f.path - ORDER BY c.file_path, c.start_line - `) + SELECT cd.chunk_key, pf.relative_path, cd.symbol, cd.kind, cd.start_line, cd.end_line + FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? + ORDER BY pf.relative_path, cd.start_line + `, db.projectID) if err != nil { t.Fatalf("query chunks: %v", err) } @@ -272,18 +302,24 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { t.Error("expected at least one interface chunk") } - // --- Verify vec_chunks has same count as chunks --- + // --- Verify every project chunk has a physical vector --- var vecCount int - if err := db.QueryRow("SELECT count(*) FROM vec_chunks").Scan(&vecCount); err != nil { - t.Fatalf("count vec_chunks: %v", err) + if err := db.QueryRow(` + SELECT count(*) + FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + JOIN vec_vectors vv ON vv.vector_id = cd.vector_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&vecCount); err != nil { + t.Fatalf("count project chunk vectors: %v", err) } if vecCount != chunkCount { - t.Errorf("vec_chunks count (%d) should match chunks count (%d)", vecCount, chunkCount) + t.Errorf("project chunk vector count (%d) should match chunk count (%d)", vecCount, chunkCount) } // --- Verify project_meta --- meta := make(map[string]string) - metaRows, err := db.Query("SELECT key, value FROM project_meta") + metaRows, err := db.Query("SELECT key, value FROM project_meta WHERE project_id = ?", db.projectID) if err != nil { t.Fatalf("query project_meta: %v", err) } @@ -303,11 +339,13 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { t.Errorf("root_hash should be 64 hex chars, got %d", len(rh)) } - // vec_dimensions should be "384" (all-minilm). - if vd, ok := meta["vec_dimensions"]; !ok { - t.Error("project_meta missing vec_dimensions") - } else if vd != "384" { - t.Errorf("expected vec_dimensions=384, got %s", vd) + // vec_dimensions is collection-scoped and should be "384" (all-minilm). + var vecDimensions string + if err := db.QueryRow("SELECT value FROM collection_meta WHERE key = 'vec_dimensions'").Scan(&vecDimensions); err != nil { + t.Fatalf("query vec_dimensions: %v", err) + } + if vecDimensions != "384" { + t.Errorf("expected vec_dimensions=384, got %s", vecDimensions) } // embedding_model should be "all-minilm". @@ -317,12 +355,12 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { t.Errorf("expected embedding_model=all-minilm, got %s", em) } - // --- Verify no orphan chunks (chunks without matching files) --- + // --- Verify no orphan chunks (chunks without matching revisions) --- var orphans int if err := db.QueryRow(` - SELECT count(*) FROM chunks c - LEFT JOIN files f ON c.file_path = f.path - WHERE f.path IS NULL + SELECT count(*) FROM chunk_defs cd + LEFT JOIN file_revisions fr ON fr.id = cd.file_revision_id + WHERE fr.id IS NULL `).Scan(&orphans); err != nil { t.Fatalf("query orphan chunks: %v", err) } @@ -340,7 +378,12 @@ func TestE2E_CLI_SQLVerifySchema(t *testing.T) { } for symbol, expectedKind := range knownSymbols { var kind string - err := db.QueryRow("SELECT kind FROM chunks WHERE symbol = ?", symbol).Scan(&kind) + err := db.QueryRow(` + SELECT cd.kind + FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? AND cd.symbol = ? + `, db.projectID, symbol).Scan(&kind) if err == sql.ErrNoRows { t.Errorf("expected symbol %q to exist in chunks", symbol) } else if err != nil { @@ -364,29 +407,34 @@ func TestE2E_CLI_SQLVerifyKNN(t *testing.T) { db := openIndexDB(t, dataHome, projectPath) - // Grab the embedding vector of ValidateToken from vec_chunks. - var tokenID string - if err := db.QueryRow("SELECT id FROM chunks WHERE symbol = 'ValidateToken'").Scan(&tokenID); err != nil { + // Grab the embedding vector of ValidateToken from the shared vector table. + var tokenID int64 + if err := db.QueryRow(` + SELECT cd.vector_id + FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? AND cd.symbol = 'ValidateToken' + `, db.projectID).Scan(&tokenID); err != nil { t.Fatalf("get ValidateToken id: %v", err) } // Use ValidateToken's own vector as the query vector — it should be // the top result (distance ≈ 0, score ≈ 1). var vecBlob []byte - if err := db.QueryRow("SELECT embedding FROM vec_chunks WHERE id = ?", tokenID).Scan(&vecBlob); err != nil { + if err := db.QueryRow("SELECT embedding FROM vec_vectors WHERE vector_id = ?", tokenID).Scan(&vecBlob); err != nil { t.Fatalf("get ValidateToken embedding: %v", err) } // Run raw KNN query. rows, err := db.Query(` - SELECT c.symbol, c.kind, v.distance - FROM vec_chunks v - JOIN chunks c ON v.id = c.id - WHERE v.embedding MATCH ? - AND v.k = 5 + SELECT cd.symbol, cd.kind, v.distance + FROM vec_vectors v + JOIN chunk_defs cd ON cd.vector_id = v.vector_id + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE v.embedding MATCH vec_int8(?) AND v.k = 5 AND pf.project_id = ? ORDER BY v.distance LIMIT 5 - `, vecBlob) + `, vecBlob, db.projectID) if err != nil { t.Fatalf("KNN query: %v", err) } @@ -418,9 +466,11 @@ func TestE2E_CLI_SQLVerifyKNN(t *testing.T) { t.Errorf("self-similarity distance should be ≈ 0, got %f", results[0].distance) } - // All distances should be non-negative and ordered ascending. + // Cosine distance can land a few ulps below zero for an identical int8 + // vector, so allow a small numerical tolerance around the valid range. + const distanceTolerance = 1e-6 for i, r := range results { - if r.distance < 0 { + if r.distance < -distanceTolerance { t.Errorf("result[%d] %s: distance should be >= 0, got %f", i, r.symbol, r.distance) } if i > 0 && r.distance < results[i-1].distance { @@ -432,7 +482,7 @@ func TestE2E_CLI_SQLVerifyKNN(t *testing.T) { // Scores (1 - distance) should all be in (0, 1]. for i, r := range results { score := 1.0 - r.distance - if score <= 0 || score > 1 { + if score <= 0 || score > 1+distanceTolerance { t.Errorf("result[%d] %s: score should be in (0, 1], got %f", i, r.symbol, score) } } @@ -454,8 +504,16 @@ func TestE2E_CLI_SQLVerifyIncremental(t *testing.T) { // Count initial state. var initialFiles, initialChunks int - db.QueryRow("SELECT count(*) FROM files").Scan(&initialFiles) - db.QueryRow("SELECT count(*) FROM chunks").Scan(&initialChunks) + if err := db.QueryRow("SELECT count(*) FROM project_files WHERE project_id = ?", db.projectID).Scan(&initialFiles); err != nil { + t.Fatalf("count initial project files: %v", err) + } + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&initialChunks); err != nil { + t.Fatalf("count initial project chunks: %v", err) + } if initialFiles != 7 { t.Fatalf("expected 7 initial files, got %d", initialFiles) @@ -463,7 +521,9 @@ func TestE2E_CLI_SQLVerifyIncremental(t *testing.T) { // Get initial root hash. var hash1 string - db.QueryRow("SELECT value FROM project_meta WHERE key = 'root_hash'").Scan(&hash1) + if err := db.QueryRow("SELECT value FROM project_meta WHERE project_id = ? AND key = 'root_hash'", db.projectID).Scan(&hash1); err != nil { + t.Fatalf("read initial root hash: %v", err) + } // Add a new file. newCode := "package project\n\n// Shutdown stops the server.\nfunc Shutdown() error { return nil }\n" @@ -484,36 +544,65 @@ func TestE2E_CLI_SQLVerifyIncremental(t *testing.T) { // Verify file count increased. var newFileCount int - db.QueryRow("SELECT count(*) FROM files").Scan(&newFileCount) + if err := db.QueryRow("SELECT count(*) FROM project_files WHERE project_id = ?", db.projectID).Scan(&newFileCount); err != nil { + t.Fatalf("count project files after addition: %v", err) + } if newFileCount != 8 { t.Errorf("expected 8 files after adding one, got %d", newFileCount) } // Verify chunk count increased. var newChunkCount int - db.QueryRow("SELECT count(*) FROM chunks").Scan(&newChunkCount) + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&newChunkCount); err != nil { + t.Fatalf("count project chunks after addition: %v", err) + } if newChunkCount <= initialChunks { t.Errorf("expected more chunks after adding file: before=%d, after=%d", initialChunks, newChunkCount) } // Verify new symbol exists. var shutdownExists int - db.QueryRow("SELECT count(*) FROM chunks WHERE symbol = 'Shutdown'").Scan(&shutdownExists) + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? AND cd.symbol = 'Shutdown' + `, db.projectID).Scan(&shutdownExists); err != nil { + t.Fatalf("query Shutdown symbol: %v", err) + } if shutdownExists == 0 { t.Error("expected Shutdown symbol in chunks after adding file") } - // Verify vec_chunks stayed in sync. + // Verify every project chunk still has a physical vector. var vecCount, chunkCount int - db.QueryRow("SELECT count(*) FROM vec_chunks").Scan(&vecCount) - db.QueryRow("SELECT count(*) FROM chunks").Scan(&chunkCount) + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + JOIN vec_vectors vv ON vv.vector_id = cd.vector_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&vecCount); err != nil { + t.Fatalf("count project chunk vectors: %v", err) + } + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? + `, db.projectID).Scan(&chunkCount); err != nil { + t.Fatalf("count project chunks: %v", err) + } if vecCount != chunkCount { - t.Errorf("vec_chunks (%d) out of sync with chunks (%d)", vecCount, chunkCount) + t.Errorf("project chunk vectors (%d) out of sync with chunks (%d)", vecCount, chunkCount) } // Root hash should have changed. var hash2 string - db.QueryRow("SELECT value FROM project_meta WHERE key = 'root_hash'").Scan(&hash2) + if err := db.QueryRow("SELECT value FROM project_meta WHERE project_id = ? AND key = 'root_hash'", db.projectID).Scan(&hash2); err != nil { + t.Fatalf("read updated root hash: %v", err) + } if hash2 == hash1 { t.Error("root_hash should change after adding a file") } @@ -529,27 +618,37 @@ func TestE2E_CLI_SQLVerifyIncremental(t *testing.T) { db = openIndexDB(t, dataHome, tmpDir) - // Verify file removed from files table. + // Verify file removed from the project's membership. var dbFileExists int - db.QueryRow("SELECT count(*) FROM files WHERE path LIKE '%database.go'").Scan(&dbFileExists) + if err := db.QueryRow("SELECT count(*) FROM project_files WHERE project_id = ? AND relative_path LIKE '%database.go'", db.projectID).Scan(&dbFileExists); err != nil { + t.Fatalf("query removed project file: %v", err) + } if dbFileExists != 0 { t.Error("database.go should be removed from files table after deletion") } // Verify QueryUsers chunks are gone. var queryUsersExists int - db.QueryRow("SELECT count(*) FROM chunks WHERE symbol = 'QueryUsers'").Scan(&queryUsersExists) + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + JOIN project_files pf ON pf.file_revision_id = cd.file_revision_id + WHERE pf.project_id = ? AND cd.symbol = 'QueryUsers' + `, db.projectID).Scan(&queryUsersExists); err != nil { + t.Fatalf("query removed QueryUsers symbol: %v", err) + } if queryUsersExists != 0 { t.Error("QueryUsers chunks should be removed after deleting database.go") } // Verify no orphan chunks. var orphans int - db.QueryRow(` - SELECT count(*) FROM chunks c - LEFT JOIN files f ON c.file_path = f.path - WHERE f.path IS NULL - `).Scan(&orphans) + if err := db.QueryRow(` + SELECT count(*) FROM chunk_defs cd + LEFT JOIN file_revisions fr ON fr.id = cd.file_revision_id + WHERE fr.id IS NULL + `).Scan(&orphans); err != nil { + t.Fatalf("query orphan chunks after deletion: %v", err) + } if orphans != 0 { t.Errorf("found %d orphan chunks after file deletion", orphans) } From 0256cd8286e4cb4b09d6272fb0153528525cfd4a Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Sun, 9 Aug 2026 08:29:17 +0200 Subject: [PATCH 7/9] fix(index): seed shared collection metadata --- internal/index/seed.go | 51 ++++++++++++++++++++++++++++++++++++- internal/index/seed_test.go | 25 +++++++++++++----- 2 files changed, 68 insertions(+), 8 deletions(-) diff --git a/internal/index/seed.go b/internal/index/seed.go index 936b45e5..727f58ae 100644 --- a/internal/index/seed.go +++ b/internal/index/seed.go @@ -138,7 +138,20 @@ func setSeedProjectPath(ctx context.Context, dbPath, projectPath string) error { if err != nil { return err } - if _, err := db.ExecContext(ctx, + + var shared bool + if err := db.QueryRowContext(ctx, + `SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'collection_meta')`, + ).Scan(&shared); err != nil { + _ = db.Close() + return err + } + if shared { + if err := setSharedSeedProjectPath(ctx, db, projectPath); err != nil { + _ = db.Close() + return err + } + } else if _, err := db.ExecContext(ctx, `INSERT INTO project_meta (key, value) VALUES ('project_path', ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, projectPath, @@ -153,6 +166,42 @@ func setSeedProjectPath(ctx context.Context, dbPath, projectPath string) error { return db.Close() } +// setSharedSeedProjectPath translates the legacy single-owner metadata update +// to the repository-scoped schema. The completed donor project retains all of +// its project_meta and project_files rows; only its canonical path changes. +func setSharedSeedProjectPath(ctx context.Context, db *sql.DB, projectPath string) error { + if abs, err := filepath.Abs(projectPath); err == nil { + projectPath = filepath.Clean(abs) + } + + var donorID int64 + err := db.QueryRowContext(ctx, ` + SELECT pm.project_id + FROM project_meta pm + JOIN projects p ON p.id = pm.project_id + WHERE pm.key = 'root_hash' AND pm.value <> '' + ORDER BY p.last_accessed_at DESC, p.id DESC + LIMIT 1`, + ).Scan(&donorID) + if err != nil { + return err + } + + var existingID int64 + err = db.QueryRowContext(ctx, `SELECT id FROM projects WHERE path = ?`, projectPath).Scan(&existingID) + switch { + case err == nil && existingID == donorID: + return nil + case err == nil: + return fmt.Errorf("seed project path %q already belongs to project %d", projectPath, existingID) + case !errors.Is(err, sql.ErrNoRows): + return err + } + + _, err = db.ExecContext(ctx, `UPDATE projects SET path = ? WHERE id = ?`, projectPath, donorID) + return err +} + func removeSQLiteFiles(path string) { for _, suffix := range []string{"", "-wal", "-shm"} { _ = os.Remove(path + suffix) diff --git a/internal/index/seed_test.go b/internal/index/seed_test.go index 45f06112..f396c4a8 100644 --- a/internal/index/seed_test.go +++ b/internal/index/seed_test.go @@ -61,25 +61,30 @@ func Hello() {} } // Verify the seeded DB works. - idx2, err := NewIndexer(dstPath, emb, 0) + idx2, err := NewIndexerForProject(dstPath, emb, 0, "int8", seedProjectDir) if err != nil { t.Fatal(err) } defer func() { _ = idx2.Close() }() - status, err := idx2.Status(projectDir) + status, err := idx2.Status(seedProjectDir) if err != nil { t.Fatal(err) } if status.IndexedFiles == 0 { t.Fatal("expected seeded DB to have indexed files") } - seedMeta, err := store.ReadMetaAt(dstPath, "project_path") + seedDB, err := sql.Open("sqlite3", sqliteFileDSN(dstPath, "ro")) if err != nil { t.Fatal(err) } - if seedMeta["project_path"] != seedProjectDir { - t.Fatalf("seeded project_path = %q, want %q", seedMeta["project_path"], seedProjectDir) + defer func() { _ = seedDB.Close() }() + var seededProjectPath string + if err := seedDB.QueryRow(`SELECT path FROM projects WHERE path = ?`, seedProjectDir).Scan(&seededProjectPath); err != nil { + t.Fatal(err) + } + if seededProjectPath != seedProjectDir { + t.Fatalf("seeded project path = %q, want %q", seededProjectPath, seedProjectDir) } } @@ -109,8 +114,13 @@ func TestSeedFromDonor_SnapshotsCommittedWALWithActiveWriter(t *testing.T) { if _, err := writer.Exec("PRAGMA journal_mode=WAL"); err != nil { t.Fatal(err) } + var projectID int64 + if err := writer.QueryRow(`SELECT id FROM projects WHERE path = ?`, projectDir).Scan(&projectID); err != nil { + t.Fatal(err) + } if _, err := writer.Exec( - `INSERT INTO project_meta (key, value) VALUES ('snapshot_marker', 'committed')`, + `INSERT INTO project_meta (project_id, key, value) VALUES (?, 'snapshot_marker', 'committed')`, + projectID, ); err != nil { t.Fatal(err) } @@ -123,7 +133,8 @@ func TestSeedFromDonor_SnapshotsCommittedWALWithActiveWriter(t *testing.T) { } t.Cleanup(func() { _ = tx.Rollback() }) if _, err := tx.Exec( - `INSERT INTO project_meta (key, value) VALUES ('snapshot_uncommitted', 'hidden')`, + `INSERT INTO project_meta (project_id, key, value) VALUES (?, 'snapshot_uncommitted', 'hidden')`, + projectID, ); err != nil { t.Fatal(err) } From 727eff332f504ac5541abb6ab5706450df20bdcb Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:19:09 +0200 Subject: [PATCH 8/9] test: update language search snapshots --- .../TestLang_Java-form_input_validation | 12 ++++------- .../TestLang_PHP-authentication_guard_session | 10 +++++----- .../TestLang_PHP-model_relationships | 14 ++++++------- .../TestLang_Python-exception_error_handling | 12 +++++------ ...stLang_Ruby-authentication_callback_filter | 15 ++++++-------- .../TestLang_Ruby-database_record_query_scope | 20 ++++++++----------- ...Script-platform_detection_operating_system | 15 +++++++------- 7 files changed, 42 insertions(+), 56 deletions(-) diff --git a/testdata/snapshots/TestLang_Java-form_input_validation b/testdata/snapshots/TestLang_Java-form_input_validation index b4d4f8db..692a5995 100644 --- a/testdata/snapshots/TestLang_Java-form_input_validation +++ b/testdata/snapshots/TestLang_Java-form_input_validation @@ -1,9 +1,5 @@ -results: 7 -Owner.java:135-145 getPet (method) -Owner.java:166-176 Owner (type) -OwnerController.java:77-83 processCreationForm+OwnerController (type) -OwnerController.java:141-154 processUpdateOwnerForm+OwnerController (type) -PetController.java:50-54 VIEWS_PETS_CREATE_OR_UPDATE_FORM+owners+types (var) +results: 4 +OwnerController.java:78-83 OwnerController (type) +OwnerController.java:142-154 OwnerController+processUpdateOwnerForm (type) PetController.java:93-145 initPetBinder+initCreationForm+PetController+processCreationForm+initUpdateForm+processUpdateForm (method) -VisitController.java:92-97 processNewVisitForm (method) - +VisitController.java:92-97 VisitController+processNewVisitForm (type) diff --git a/testdata/snapshots/TestLang_PHP-authentication_guard_session b/testdata/snapshots/TestLang_PHP-authentication_guard_session index 4f30bfcd..c35027e3 100644 --- a/testdata/snapshots/TestLang_PHP-authentication_guard_session +++ b/testdata/snapshots/TestLang_PHP-authentication_guard_session @@ -1,7 +1,7 @@ -results: 5 +results: 6 AuthManager.php:9-23 AuthManager (type) -AuthManager.php:47-69 __construct+AuthManager+guard (type) -AuthManager.php:116-166 AuthManager+createSessionDriver+createTokenDriver (type) +AuthManager.php:60-69 AuthManager+guard (type) +AuthManager.php:116-139 AuthManager+createSessionDriver (type) +AuthManager.php:148-166 createSessionDriver+AuthManager+createTokenDriver (method) AuthManager.php:180-190 AuthManager+getConfig (type) -AuthManager.php:296-316 AuthManager+hasResolvedGuards+forgetGuards (method) - +AuthManager.php:296-316 hasResolvedGuards+AuthManager+forgetGuards (method) diff --git a/testdata/snapshots/TestLang_PHP-model_relationships b/testdata/snapshots/TestLang_PHP-model_relationships index 09ec417c..4c9d3cac 100644 --- a/testdata/snapshots/TestLang_PHP-model_relationships +++ b/testdata/snapshots/TestLang_PHP-model_relationships @@ -1,12 +1,10 @@ -results: 10 +results: 9 BelongsTo.php:22-53 BelongsTo+child+foreignKey+ownerKey+relationName (type) -BelongsTo.php:95-100 addConstraints (method) BelongsTo.php:118-127 BelongsTo+getEagerModelKeys (type) -BelongsTo.php:154-159 match (method) -BelongsTo.php:192-200 BelongsTo+associate (type) +BelongsTo.php:154-159 BelongsTo+match (type) +BelongsTo.php:192-200 associate+BelongsTo (method) BelongsTo.php:213-233 BelongsTo+touch (type) Builder.php:815-822 Builder+eagerLoadRelations (type) -Builder.php:852-856 eagerLoadRelation (method) -Model.php:1081-1095 push (method) -Model.php:1538-1547 Model+newQueryWithoutRelationships (type) - +Builder.php:852-856 Builder+eagerLoadRelation (type) +Model.php:1081-1095 push+Model (type) +Model.php:1538-1547 newQueryWithoutRelationships+Model (method) diff --git a/testdata/snapshots/TestLang_Python-exception_error_handling b/testdata/snapshots/TestLang_Python-exception_error_handling index f8baf089..a3cbc2f3 100644 --- a/testdata/snapshots/TestLang_Python-exception_error_handling +++ b/testdata/snapshots/TestLang_Python-exception_error_handling @@ -1,9 +1,9 @@ -results: 7 +results: 8 django-exceptions.py:147-152 __init__+ValidationError (function) -django-exceptions.py:166-222 __init__+ValidationError+update_error_dict+__iter__ (function) +django-exceptions.py:166-205 __init__+ValidationError (function) +django-exceptions.py:212-218 __iter__ (function) django-exceptions.py:245-247 __hash__ (function) flask-app.py:828-836 Flask+handle_http_exception (function) -flask-app.py:867-872 handle_user_exception (function) -flask-app.py:882-891 handle_user_exception (function) -flask-app.py:948-953 Flask+log_exception (type) - +flask-app.py:867-872 Flask+handle_user_exception (type) +flask-app.py:882-891 Flask+handle_user_exception (type) +flask-app.py:948-953 Flask (type) diff --git a/testdata/snapshots/TestLang_Ruby-authentication_callback_filter b/testdata/snapshots/TestLang_Ruby-authentication_callback_filter index af0c8ae4..10ca756e 100644 --- a/testdata/snapshots/TestLang_Ruby-authentication_callback_filter +++ b/testdata/snapshots/TestLang_Ruby-authentication_callback_filter @@ -1,12 +1,9 @@ -results: 10 -application.rb:772-781 Rails.Application+Application.coerce_same_site_protection (type) -base.rb:108-119 ActionCable (type) -base.rb:306-324 Base.action_signature+Channel.Base (function) -cache.rb:442-451 Store.fetch+Cache.Store (function) -callbacks.rb:86-89 ActiveRecord (type) +results: 8 +application.rb:772-781 Rails+Rails.Application+Application.coerce_same_site_protection (type) +base.rb:306-324 Base.action_signature+ActionCable+ActionCable.Channel+Channel.Base (function) +cache.rb:442-447 Store.fetch (function) callbacks.rb:101-110 ActiveRecord (type) callbacks.rb:125-131 ActiveRecord (type) -callbacks.rb:281-287 Callbacks.CALLBACKS+ActiveRecord (var) +callbacks.rb:281-287 Callbacks.CALLBACKS+ActiveRecord+ActiveRecord.Callbacks (var) metal.rb:8-17 ActionController (type) -sinatra-base.rb:1483-1509 Sinatra.Base+Base.before+Base.after (type) - +sinatra-base.rb:1490-1509 Base.after+Sinatra+Sinatra.Base (type) diff --git a/testdata/snapshots/TestLang_Ruby-database_record_query_scope b/testdata/snapshots/TestLang_Ruby-database_record_query_scope index 10c97d53..7c2c6479 100644 --- a/testdata/snapshots/TestLang_Ruby-database_record_query_scope +++ b/testdata/snapshots/TestLang_Ruby-database_record_query_scope @@ -1,12 +1,8 @@ -results: 10 -associations.rb:764-768 ActiveRecord (type) -associations.rb:1482-1491 ActiveRecord (type) -associations.rb:1692-1699 ActiveRecord (type) -associations.rb:1932-1939 ActiveRecord (type) -relation.rb:470-474 ActiveRecord.Relation+Relation.cache_version (function) -relation.rb:497-510 Relation.compute_cache_version (function) -relation.rb:537-561 ActiveRecord.Relation+Relation.scoping (type) -relation.rb:1339-1352 ActiveRecord.Relation (type) -relation.rb:1362-1376 ActiveRecord.Relation (type) -relation.rb:1397-1412 Relation._scoping+ActiveRecord.Relation (type) - +results: 7 +associations.rb:764-768 ActiveRecord+ActiveRecord.Associations (type) +associations.rb:1482-1491 ActiveRecord+ActiveRecord.Associations+Associations.ClassMethods (type) +associations.rb:1692-1699 ActiveRecord+ActiveRecord.Associations+Associations.ClassMethods (type) +relation.rb:502-510 ActiveRecord+ActiveRecord.Relation+Relation.compute_cache_version (type) +relation.rb:537-552 ActiveRecord+ActiveRecord.Relation+Relation.scoping (type) +relation.rb:1362-1376 ActiveRecord+ActiveRecord.Relation (type) +relation.rb:1397-1412 Relation._scoping+ActiveRecord+ActiveRecord.Relation (type) diff --git a/testdata/snapshots/TestLang_TypeScript-platform_detection_operating_system b/testdata/snapshots/TestLang_TypeScript-platform_detection_operating_system index 498b1c03..58beb54c 100644 --- a/testdata/snapshots/TestLang_TypeScript-platform_detection_operating_system +++ b/testdata/snapshots/TestLang_TypeScript-platform_detection_operating_system @@ -1,12 +1,11 @@ results: 10 -path.ts:293-306 win32 (const) -path.ts:404-420 win32 (const) +path.ts:293-306 resolve+win32 (method) +path.ts:404-420 win32+normalize (const) path.ts:453-473 isAbsolute+win32 (const) -path.ts:656-662 win32 (const) -path.ts:723-732 toNamespacedPath+win32 (method) -path.ts:791-800 win32 (const) -path.ts:1036-1044 win32 (const) +path.ts:656-662 relative+win32 (method) +path.ts:723-732 win32+toNamespacedPath (const) +path.ts:791-800 dirname+win32 (method) +path.ts:1036-1044 parse+win32 (method) platform.ts:37-49 INodeProcess (interface) -platform.ts:123-256 PlatformName+Platform+OperatingSystem+PlatformToString+locale+platformLocale+setTimeout0+OS+isChrome+isFirefox+isSafari+isEdge+isAndroid (function) +platform.ts:123-256 PlatformName+Platform+OperatingSystem+PlatformToString+locale+platformLocale (function) platform.ts:279-281 isTahoeOrNewer (function) - From 192607256dca357e16c7a4bbddb7c3891aa6a459 Mon Sep 17 00:00:00 2001 From: aeneasr <3372410+aeneasr@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:18:02 +0200 Subject: [PATCH 9/9] fix: address shared index review feedback --- cmd/clean.go | 137 ++++++++++++++++++++++----------- cmd/clean_test.go | 32 ++++++-- cmd/stdio.go | 37 +++++++-- cmd/stdio_test.go | 68 +++++++++++++++- e2e_cli_test.go | 2 +- internal/config/config.go | 15 ++-- internal/config/config_test.go | 10 +++ internal/index/index_test.go | 58 +++++++------- internal/index/seed.go | 42 +++++++--- internal/index/seed_test.go | 47 +++++++++++ internal/store/shared.go | 46 +++++------ internal/store/shared_test.go | 64 +++++++++++++++ internal/store/store.go | 4 +- 13 files changed, 429 insertions(+), 133 deletions(-) diff --git a/cmd/clean.go b/cmd/clean.go index c9caf6d1..86753bae 100644 --- a/cmd/clean.go +++ b/cmd/clean.go @@ -15,18 +15,17 @@ package cmd import ( - "bytes" "fmt" "io" "log/slog" "os" "path/filepath" - "strings" "time" "github.com/ory/lumen/internal/config" "github.com/ory/lumen/internal/indexlock" "github.com/ory/lumen/internal/store" + "github.com/ory/lumen/internal/tui" "github.com/spf13/cobra" ) @@ -90,31 +89,83 @@ func runClean(cmd *cobra.Command, _ []string) error { return fmt.Errorf("--days must not exceed %d, got %d", maxCleanDays, days) } dataDir := filepath.Join(config.XDGDataDir(), "lumen") - return cleanIndexes(cmd.ErrOrStderr(), cmd.OutOrStdout(), dataDir, days, time.Now()) + reporter := interactiveCleanReporter{progress: tui.NewProgress(os.Stderr)} + summary, err := cleanIndexes(reporter, dataDir, days, time.Now()) + if output := formatCleanSummary(summary); output != "" { + fmt.Printf("%s", output) + } + return err +} + +type cleanReporter interface { + Info(string) + Error(string) +} + +type interactiveCleanReporter struct { + progress *tui.Progress +} + +func (r interactiveCleanReporter) Info(message string) { + r.progress.Info(message) +} + +func (interactiveCleanReporter) Error(message string) { + fmt.Fprintf(os.Stderr, "%s\n", message) +} + +type slogCleanReporter struct { + logger *slog.Logger +} + +func (r slogCleanReporter) Info(message string) { + r.logger.Info("daily cleanup detail", "message", message) +} + +func (r slogCleanReporter) Error(message string) { + r.logger.Warn("daily cleanup issue", "message", message) +} + +type cleanSummary struct { + noData bool + removed int + skipped int + projectsRemoved int + vectorsRemoved int + bytesReclaimed int64 +} + +func formatCleanSummary(summary cleanSummary) string { + if summary.noData { + return "" + } + output := fmt.Sprintf("Removed %d index director%s, skipped %d.\n", + summary.removed, pluralY(summary.removed), summary.skipped) + if summary.projectsRemoved > 0 || summary.vectorsRemoved > 0 || summary.bytesReclaimed > 0 { + output += fmt.Sprintf("Shared cleanup: %d projects, %d vectors, %d bytes reclaimed.\n", + summary.projectsRemoved, summary.vectorsRemoved, summary.bytesReclaimed) + } + return output } -// cleanIndexes removes every stale index directory directly under dataDir, -// reporting each decision on the injected stderr and a summary on the injected -// stdout. The injected writers deliberately keep this reusable by both the -// interactive CLI and the MCP background cleanup without mutating pterm's -// process-global state. now is injected so the age cutoff is testable. Failures -// to remove a single directory are -// reported and the sweep continues; the first such failure is returned once -// every directory has been considered. -func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.Time) error { +// cleanIndexes removes every stale index directory directly under dataDir and +// reports individual decisions through reporter. The returned summary is +// rendered by the caller using the output strategy for its execution context. +// Failures to remove a single directory are reported and the sweep continues; +// the first such failure is returned after every directory has been considered. +func cleanIndexes(reporter cleanReporter, dataDir string, days int, now time.Time) (cleanSummary, error) { + var summary cleanSummary entries, err := os.ReadDir(dataDir) if err != nil { if os.IsNotExist(err) { - _, _ = fmt.Fprintln(stderr, "No index data found — nothing to clean.") - return nil + reporter.Info("No index data found — nothing to clean.") + summary.noData = true + return summary, nil } - return fmt.Errorf("read data dir: %w", err) + return summary, fmt.Errorf("read data dir: %w", err) } cutoff := now.Add(-time.Duration(days) * 24 * time.Hour) - removed, skipped := 0, 0 - projectsRemoved, vectorsRemoved := 0, 0 - var bytesReclaimed int64 var firstErr error for _, entry := range entries { @@ -124,40 +175,33 @@ func cleanIndexes(stderr, stdout io.Writer, dataDir string, days int, now time.T continue } hashDir := filepath.Join(dataDir, entry.Name()) - wasRemoved, sharedStats, cleanErr := cleanIndex(stderr, entry.Name(), hashDir, days, cutoff) - projectsRemoved += sharedStats.ProjectsRemoved - vectorsRemoved += sharedStats.VectorsRemoved - bytesReclaimed += sharedStats.BytesReclaimed + wasRemoved, sharedStats, cleanErr := cleanIndex(reporter, entry.Name(), hashDir, days, cutoff) + summary.projectsRemoved += sharedStats.ProjectsRemoved + summary.vectorsRemoved += sharedStats.VectorsRemoved + summary.bytesReclaimed += sharedStats.BytesReclaimed if wasRemoved { - removed++ + summary.removed++ } else { - skipped++ + summary.skipped++ } if cleanErr != nil && firstErr == nil { firstErr = cleanErr } } - - _, _ = fmt.Fprintf(stdout, "Removed %d index director%s, skipped %d.\n", - removed, pluralY(removed), skipped) - if projectsRemoved > 0 || vectorsRemoved > 0 || bytesReclaimed > 0 { - _, _ = fmt.Fprintf(stdout, "Shared cleanup: %d projects, %d vectors, %d bytes reclaimed.\n", - projectsRemoved, vectorsRemoved, bytesReclaimed) - } - return firstErr + return summary, firstErr } // cleanIndex cleans one legacy index or shared collection while retaining the // exclusive collection lock for the entire database cleanup and removal. -func cleanIndex(stderr io.Writer, name, hashDir string, days int, cutoff time.Time) (bool, store.CleanupStats, error) { +func cleanIndex(reporter cleanReporter, name, hashDir string, days int, cutoff time.Time) (bool, store.CleanupStats, error) { dbPath := filepath.Join(hashDir, "index.db") lock, lockErr := tryAcquireExclusive(indexlock.LockPathForDB(dbPath)) if lockErr != nil { - _, _ = fmt.Fprintf(stderr, "Failed to acquire index lock for %s: %v\n", name, lockErr) + reporter.Error(fmt.Sprintf("Failed to acquire index lock for %s: %v", name, lockErr)) return false, store.CleanupStats{}, fmt.Errorf("acquire index lock for %s: %w", name, lockErr) } if lock == nil { - _, _ = fmt.Fprintf(stderr, "Keeping %s: an indexer is currently running.\n", name) + reporter.Info(fmt.Sprintf("Keeping %s: an indexer is currently running.", name)) return false, store.CleanupStats{}, nil } defer lock.Release() @@ -165,16 +209,17 @@ func cleanIndex(stderr io.Writer, name, hashDir string, days int, cutoff time.Ti sharedStats, shared, sharedErr := cleanupCollectionAt(dbPath, cutoff) if shared { if sharedErr != nil { - _, _ = fmt.Fprintf(stderr, "Failed to clean shared collection %s: %v\n", name, sharedErr) + reporter.Error(fmt.Sprintf("Failed to clean shared collection %s: %v", name, sharedErr)) return false, store.CleanupStats{}, fmt.Errorf("clean shared collection %s: %w", name, sharedErr) } if sharedStats.ProjectsLeft > 0 { - _, _ = fmt.Fprintf(stderr, "Cleaned %s: removed %d projects and %d vectors.\n", name, sharedStats.ProjectsRemoved, sharedStats.VectorsRemoved) + reporter.Info(fmt.Sprintf("Cleaned %s: removed %d projects and %d vectors.", name, sharedStats.ProjectsRemoved, sharedStats.VectorsRemoved)) return false, sharedStats, nil } // Empty collections have no future owner and can be removed as a // directory, reclaiming sidecars and metadata in one operation. if err := removeIndexDir(hashDir); err != nil { + reporter.Error(fmt.Sprintf("Failed to remove %s: %v", hashDir, err)) return false, sharedStats, fmt.Errorf("remove empty collection %s: %w", hashDir, err) } return true, sharedStats, nil @@ -185,10 +230,10 @@ func cleanIndex(stderr io.Writer, name, hashDir string, days int, cutoff time.Ti return false, store.CleanupStats{}, nil } if err := removeIndexDir(hashDir); err != nil { - _, _ = fmt.Fprintf(stderr, "Failed to remove %s: %v\n", hashDir, err) + reporter.Error(fmt.Sprintf("Failed to remove %s: %v", hashDir, err)) return false, store.CleanupStats{}, fmt.Errorf("remove %s: %w", hashDir, err) } - _, _ = fmt.Fprintf(stderr, "Removed %s (%s).\n", name, reason) + reporter.Info(fmt.Sprintf("Removed %s (%s).", name, reason)) return true, store.CleanupStats{}, nil } @@ -273,12 +318,18 @@ func runDailyCleanup(dataDir string, now time.Time, logger *slog.Logger) { logger.Warn("daily cleanup: create data directory", "path", dataDir, "error", err) return } - var stderr, stdout bytes.Buffer - if err := cleanIndexes(&stderr, &stdout, dataDir, defaultCleanDays, now); err != nil { - logger.Warn("daily cleanup failed", "error", err, "details", strings.TrimSpace(stderr.String())) + summary, err := cleanIndexes(slogCleanReporter{logger: logger}, dataDir, defaultCleanDays, now) + if err != nil { + logger.Warn("daily cleanup failed", "error", err) return } - logger.Info("daily cleanup complete", "summary", strings.TrimSpace(stdout.String()), "details", strings.TrimSpace(stderr.String())) + logger.Info("daily cleanup complete", + "indexes_removed", summary.removed, + "indexes_skipped", summary.skipped, + "projects_removed", summary.projectsRemoved, + "vectors_removed", summary.vectorsRemoved, + "bytes_reclaimed", summary.bytesReclaimed, + ) if err := os.WriteFile(stampPath, []byte(now.UTC().Format(time.RFC3339)), 0o600); err != nil { logger.Warn("daily cleanup: write stamp", "path", stampPath, "error", err) } diff --git a/cmd/clean_test.go b/cmd/clean_test.go index aa538803..b017628c 100644 --- a/cmd/clean_test.go +++ b/cmd/clean_test.go @@ -18,6 +18,7 @@ import ( "bytes" "database/sql" "errors" + "fmt" "log/slog" "os" "path/filepath" @@ -86,10 +87,25 @@ func projectDir(t *testing.T, name string) string { // runCleanIndexes invokes the cleanup sweep against the data dir under tmp. func runCleanIndexes(t *testing.T, tmp string, days int) (stdout, stderr string, err error) { t.Helper() - outBuf := new(bytes.Buffer) - errBuf := new(bytes.Buffer) - err = cleanIndexes(errBuf, outBuf, filepath.Join(tmp, "lumen"), days, cleanNow) - return outBuf.String(), errBuf.String(), err + reporter := newBufferCleanReporter() + summary, err := cleanIndexes(reporter, filepath.Join(tmp, "lumen"), days, cleanNow) + return formatCleanSummary(summary), reporter.output.String(), err +} + +type bufferCleanReporter struct { + output bytes.Buffer +} + +func newBufferCleanReporter() *bufferCleanReporter { + return &bufferCleanReporter{} +} + +func (r *bufferCleanReporter) Info(message string) { + _, _ = fmt.Fprintln(&r.output, message) +} + +func (r *bufferCleanReporter) Error(message string) { + _, _ = fmt.Fprintln(&r.output, message) } // runCleanCmd invokes runClean through a command carrying the real clean flags. @@ -315,13 +331,13 @@ func TestCleanIndexReportsLockAcquisitionErrors(t *testing.T) { tryAcquireExclusive = func(string) (*indexlock.Lock, error) { return nil, errors.New("permission denied") } - var stderr bytes.Buffer - removed, _, err := cleanIndex(&stderr, "abc", t.TempDir(), 30, time.Now()) + reporter := newBufferCleanReporter() + removed, _, err := cleanIndex(reporter, "abc", t.TempDir(), 30, time.Now()) if err == nil || removed { t.Fatalf("removed=%v err=%v", removed, err) } - if !strings.Contains(stderr.String(), "Failed to acquire index lock") || strings.Contains(stderr.String(), "currently running") { - t.Fatalf("unexpected stderr: %s", stderr.String()) + if !strings.Contains(reporter.output.String(), "Failed to acquire index lock") || strings.Contains(reporter.output.String(), "currently running") { + t.Fatalf("unexpected stderr: %s", reporter.output.String()) } } diff --git a/cmd/stdio.go b/cmd/stdio.go index 29077cb9..fff11a35 100644 --- a/cmd/stdio.go +++ b/cmd/stdio.go @@ -154,6 +154,7 @@ const staleIndexWarning = "Index is being updated in the background. Results may var ( tryAcquire = indexlock.TryAcquire tryAcquireShared = indexlock.TryAcquireShared + runDailyCleanupFunc = runDailyCleanup prepareMigrationFunc = func(idx *index.Indexer, projectDir, legacyPath string) error { return idx.PrepareLegacyMigration(projectDir, legacyPath) } @@ -214,6 +215,7 @@ type indexerCache struct { mu sync.RWMutex cache map[string]cacheEntry reindexing map[string]bool // projects with an active background reindex goroutine + migrationPrepared map[string]bool // project/model keys already scanned for legacy vectors embedder embedder.Embedder cfg *config.ConfigService freshnessTTL time.Duration // override for tests; 0 reads from cfg, then defaultFreshnessTTL @@ -300,6 +302,29 @@ func (ic *indexerCache) logger() *slog.Logger { return ic.log } +// markMigrationPrepared reports whether this is the first legacy-migration +// preparation attempt for key. Failed attempts are intentionally remembered: +// a later background refresh should rebuild missing vectors instead of +// repeatedly scanning the same legacy database. +func (ic *indexerCache) markMigrationPrepared(key string) bool { + ic.mu.Lock() + defer ic.mu.Unlock() + if ic.migrationPrepared == nil { + ic.migrationPrepared = make(map[string]bool) + } + if ic.migrationPrepared[key] { + return false + } + ic.migrationPrepared[key] = true + return true +} + +func (ic *indexerCache) startDailyCleanup(dataDir string, now time.Time, logger *slog.Logger) { + ic.wg.Go(func() { + runDailyCleanupFunc(dataDir, now, logger) + }) +} + // Close cancels all background reindex goroutines, waits for them to drain // (up to 30 seconds), then closes all cached indexers. Call on MCP server // shutdown. @@ -332,6 +357,7 @@ func (ic *indexerCache) Close() { } } ic.cache = nil + ic.migrationPrepared = nil } // findEffectiveRoot walks up the directory tree from path's parent to find an @@ -896,9 +922,11 @@ func (ic *indexerCache) ensureIndexed(idx *index.Indexer, input SemanticSearchIn } } - legacyPath := config.LegacyDBPathForProject(projectDir, modelName) - if err := prepareMigrationFunc(idx, projectDir, legacyPath); err != nil { - ic.logger().Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + if ic.markMigrationPrepared(reindexKey) { + legacyPath := config.LegacyDBPathForProject(projectDir, modelName) + if err := prepareMigrationFunc(idx, projectDir, legacyPath); err != nil { + ic.logger().Warn("legacy index migration unavailable; rebuilding missing vectors", "path", legacyPath, "error", err) + } } ensureFresh := ic.ensureFreshFunc @@ -1474,8 +1502,6 @@ func runStdio(_ *cobra.Command, _ []string) error { "backend", cfg.Servers()[0].Backend, "freshness_ttl", cfg.FreshnessTTL().String(), ) - runDailyCleanup(filepath.Join(config.XDGDataDir(), "lumen"), time.Now(), logger) - closeCtx, closeFn := context.WithCancel(context.Background()) indexers := &indexerCache{ embedder: emb, @@ -1485,6 +1511,7 @@ func runStdio(_ *cobra.Command, _ []string) error { closeFn: closeFn, } defer indexers.Close() + indexers.startDailyCleanup(filepath.Join(config.XDGDataDir(), "lumen"), time.Now(), logger) server := mcp.NewServer(&mcp.Implementation{ Name: "lumen", diff --git a/cmd/stdio_test.go b/cmd/stdio_test.go index 975b1975..1bbd9b47 100644 --- a/cmd/stdio_test.go +++ b/cmd/stdio_test.go @@ -1419,7 +1419,8 @@ func TestLegacyMigrationPreparationRunsInBackgroundIndexing(t *testing.T) { prepareMigrationFunc = func(_ *index.Indexer, gotProject, _ string) error { prepareCalls++ if gotProject != projectDir { - t.Fatalf("project = %q, want %q", gotProject, projectDir) + t.Errorf("project = %q, want %q", gotProject, projectDir) + return nil } return nil } @@ -1429,7 +1430,8 @@ func TestLegacyMigrationPreparationRunsInBackgroundIndexing(t *testing.T) { log: discardLog, ensureFreshFunc: func(_ context.Context, _ *index.Indexer, _ string, _ index.ProgressFunc) (bool, index.Stats, error) { if prepareCalls != 1 { - t.Fatalf("PrepareLegacyMigration calls before EnsureFresh = %d, want 1", prepareCalls) + t.Errorf("PrepareLegacyMigration calls before EnsureFresh = %d, want 1", prepareCalls) + return false, index.Stats{}, nil } return false, index.Stats{}, nil }, @@ -1449,6 +1451,68 @@ func TestLegacyMigrationPreparationRunsInBackgroundIndexing(t *testing.T) { if prepareCalls != 1 { t.Fatalf("PrepareLegacyMigration calls = %d, want 1", prepareCalls) } + + deadline := time.Now().Add(time.Second) + for { + ic.mu.RLock() + active := ic.reindexing[cacheKey(effectiveRoot, "stub")] + ic.mu.RUnlock() + if !active { + break + } + if time.Now().After(deadline) { + t.Fatal("background reindex state was not cleared") + } + time.Sleep(time.Millisecond) + } + ic.mu.Lock() + for key, entry := range ic.cache { + entry.lastCheckedAt = time.Time{} + ic.cache[key] = entry + } + ic.mu.Unlock() + if _, err := ic.ensureIndexed(idx, input, effectiveRoot, ic.dbPath(effectiveRoot, "stub"), nil); err != nil { + t.Fatal(err) + } + if prepareCalls != 1 { + t.Fatalf("PrepareLegacyMigration calls after second refresh = %d, want 1", prepareCalls) + } +} + +func TestStartDailyCleanupIsAsyncAndTracked(t *testing.T) { + original := runDailyCleanupFunc + started := make(chan struct{}) + release := make(chan struct{}) + runDailyCleanupFunc = func(string, time.Time, *slog.Logger) { + close(started) + <-release + } + t.Cleanup(func() { runDailyCleanupFunc = original }) + + ic := &indexerCache{} + ic.startDailyCleanup(t.TempDir(), time.Now(), discardLog) + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("daily cleanup did not start") + } + + closeDone := make(chan struct{}) + go func() { + ic.Close() + close(closeDone) + }() + select { + case <-closeDone: + t.Fatal("Close returned before tracked daily cleanup completed") + case <-time.After(20 * time.Millisecond): + } + close(release) + select { + case <-closeDone: + case <-time.After(time.Second): + t.Fatal("Close did not return after daily cleanup completed") + } } func TestFormatSearchResults_IncludesSeedWarning(t *testing.T) { diff --git a/e2e_cli_test.go b/e2e_cli_test.go index 23698a74..848b9234 100644 --- a/e2e_cli_test.go +++ b/e2e_cli_test.go @@ -154,7 +154,7 @@ func openIndexDB(t *testing.T, dataHome, projectPath string) *projectIndexDB { if err != nil { t.Fatalf("open index db: %v", err) } - t.Cleanup(func() { db.Close() }) + t.Cleanup(func() { _ = db.Close() }) absoluteProjectPath, err := filepath.Abs(projectPath) if err != nil { t.Fatalf("resolve project path: %v", err) diff --git a/internal/config/config.go b/internal/config/config.go index e11f21a0..f1d65533 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -47,13 +47,16 @@ func DBPathForProjectBase(dataDir, projectPath, model string) string { return DBPathForProjectProfileBase(dataDir, projectPath, model, dimensions, "int8", 512) } -// ModelDimensions resolves dimensions for a model in the built-in registry. -func ModelDimensions(model string) (int, bool) { - canonical := model +func canonicalModel(model string) string { if resolved, ok := models.ModelAliases[model]; ok { - canonical = resolved + return resolved } - spec, ok := models.KnownModels[canonical] + return model +} + +// ModelDimensions resolves dimensions for a model in the built-in registry. +func ModelDimensions(model string) (int, bool) { + spec, ok := models.KnownModels[canonicalModel(model)] return spec.Dims, ok } @@ -82,7 +85,7 @@ func DBPathForProjectProfileBase(dataDir, projectPath, model string, dimensions } else if resolved, resolveErr := filepath.EvalSymlinks(identity); resolveErr == nil { identity = filepath.Clean(resolved) } - profile := identity + "\x00" + scope + "\x00" + model + "\x00" + + profile := identity + "\x00" + scope + "\x00" + canonicalModel(model) + "\x00" + strconv.Itoa(dimensions) + "\x00" + vectorStorage + "\x00" + strconv.Itoa(maxChunkTokens) + "\x00" + IndexVersion hash := fmt.Sprintf("%x", sha256.Sum256([]byte(profile))) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 2d163e36..c249f39d 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -124,6 +124,16 @@ func TestDBPathForProjectProfileResolvesNonGitSymlinks(t *testing.T) { } } +func TestDBPathForProjectProfileCanonicalizesModelAliases(t *testing.T) { + dataDir := t.TempDir() + project := t.TempDir() + aliasPath := DBPathForProjectProfileBase(dataDir, project, "text-embedding-nomic-embed-code", 3584, "int8", 512) + canonicalPath := DBPathForProjectProfileBase(dataDir, project, "nomic-ai/nomic-embed-code-GGUF", 3584, "int8", 512) + if aliasPath != canonicalPath { + t.Fatalf("alias and canonical model should share a collection: %q != %q", aliasPath, canonicalPath) + } +} + func TestXDGConfigDir(t *testing.T) { t.Run("uses XDG_CONFIG_HOME when set", func(t *testing.T) { t.Setenv("XDG_CONFIG_HOME", "/custom/config") diff --git a/internal/index/index_test.go b/internal/index/index_test.go index d56cd02b..ab8fcfef 100644 --- a/internal/index/index_test.go +++ b/internal/index/index_test.go @@ -69,35 +69,6 @@ func Hello(name string) { fmt.Println("hello", name) } -func TestIndexerLastIndexedAtIsProjectScoped(t *testing.T) { - projectA, projectB := t.TempDir(), t.TempDir() - idx, err := NewIndexerForProject(":memory:", &mockEmbedder{dims: 4, model: "test-model"}, 512, "int8", projectA) - if err != nil { - t.Fatal(err) - } - defer func() { _ = idx.Close() }() - timeA := time.Now().Add(-time.Hour).UTC().Truncate(time.Second) - timeB := time.Now().UTC().Truncate(time.Second) - if err := idx.store.SetMeta("last_indexed_at", timeA.Format(time.RFC3339)); err != nil { - t.Fatal(err) - } - release, err := idx.lockProject(projectB) - if err != nil { - t.Fatal(err) - } - if err := idx.store.SetMeta("last_indexed_at", timeB.Format(time.RFC3339)); err != nil { - release() - t.Fatal(err) - } - release() - if got, ok := idx.LastIndexedAt(projectA); !ok || !got.Equal(timeA) { - t.Fatalf("project A LastIndexedAt = %v, %v; want %v, true", got, ok, timeA) - } - if got, ok := idx.LastIndexedAt(projectB); !ok || !got.Equal(timeB) { - t.Fatalf("project B LastIndexedAt = %v, %v; want %v, true", got, ok, timeB) - } -} - // Goodbye prints a farewell. func Goodbye(name string) { fmt.Println("bye", name) @@ -131,6 +102,35 @@ func Goodbye(name string) { } } +func TestIndexerLastIndexedAtIsProjectScoped(t *testing.T) { + projectA, projectB := t.TempDir(), t.TempDir() + idx, err := NewIndexerForProject(":memory:", &mockEmbedder{dims: 4, model: "test-model"}, 512, "int8", projectA) + if err != nil { + t.Fatal(err) + } + defer func() { _ = idx.Close() }() + timeA := time.Now().Add(-time.Hour).UTC().Truncate(time.Second) + timeB := time.Now().UTC().Truncate(time.Second) + if err := idx.store.SetMeta("last_indexed_at", timeA.Format(time.RFC3339)); err != nil { + t.Fatal(err) + } + release, err := idx.lockProject(projectB) + if err != nil { + t.Fatal(err) + } + if err := idx.store.SetMeta("last_indexed_at", timeB.Format(time.RFC3339)); err != nil { + release() + t.Fatal(err) + } + release() + if got, ok := idx.LastIndexedAt(projectA); !ok || !got.Equal(timeA) { + t.Fatalf("project A LastIndexedAt = %v, %v; want %v, true", got, ok, timeA) + } + if got, ok := idx.LastIndexedAt(projectB); !ok || !got.Equal(timeB) { + t.Fatalf("project B LastIndexedAt = %v, %v; want %v, true", got, ok, timeB) + } +} + func TestIndexer_IncrementalIndex(t *testing.T) { projectDir := t.TempDir() writeGoFile(t, projectDir, "main.go", `package main diff --git a/internal/index/seed.go b/internal/index/seed.go index 727f58ae..c942fbc7 100644 --- a/internal/index/seed.go +++ b/internal/index/seed.go @@ -78,8 +78,25 @@ func SeedFromDonorContext(ctx context.Context, donorPath, dstPath, projectPath s if err != nil { return false, fmt.Errorf("open donor: %w", err) } + var shared bool + if err := db.QueryRowContext(ctx, + `SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'collection_meta')`, + ).Scan(&shared); err != nil { + _ = db.Close() + return false, fmt.Errorf("detect donor schema: %w", err) + } var rootHash sql.NullString - if err := db.QueryRowContext(ctx, "SELECT value FROM project_meta WHERE key = 'root_hash'").Scan(&rootHash); err != nil && !errors.Is(err, sql.ErrNoRows) { + rootHashQuery := "SELECT value FROM project_meta WHERE key = 'root_hash'" + if shared { + rootHashQuery = ` + SELECT pm.value + FROM project_meta pm + JOIN projects p ON p.id = pm.project_id + WHERE pm.key = 'root_hash' AND pm.value <> '' + ORDER BY p.last_accessed_at DESC, p.id DESC + LIMIT 1` + } + if err := db.QueryRowContext(ctx, rootHashQuery).Scan(&rootHash); err != nil && !errors.Is(err, sql.ErrNoRows) { _ = db.Close() return false, fmt.Errorf("read donor metadata: %w", err) } @@ -133,37 +150,38 @@ func sqliteFileDSN(path, mode string) string { }).String() } -func setSeedProjectPath(ctx context.Context, dbPath, projectPath string) error { +func setSeedProjectPath(ctx context.Context, dbPath, projectPath string) (err error) { db, err := sql.Open("sqlite3", sqliteFileDSN(dbPath, "rw")) if err != nil { - return err + return fmt.Errorf("open seed snapshot: %w", err) } + defer func() { + if closeErr := db.Close(); closeErr != nil { + err = errors.Join(err, fmt.Errorf("close seed snapshot: %w", closeErr)) + } + }() var shared bool if err := db.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'collection_meta')`, ).Scan(&shared); err != nil { - _ = db.Close() - return err + return fmt.Errorf("detect seed schema: %w", err) } if shared { if err := setSharedSeedProjectPath(ctx, db, projectPath); err != nil { - _ = db.Close() - return err + return fmt.Errorf("stamp shared seed project path: %w", err) } } else if _, err := db.ExecContext(ctx, `INSERT INTO project_meta (key, value) VALUES ('project_path', ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, projectPath, ); err != nil { - _ = db.Close() - return err + return fmt.Errorf("stamp seed project path: %w", err) } if _, err := db.ExecContext(ctx, "PRAGMA wal_checkpoint(TRUNCATE)"); err != nil { - _ = db.Close() - return err + return fmt.Errorf("checkpoint seed snapshot: %w", err) } - return db.Close() + return nil } // setSharedSeedProjectPath translates the legacy single-owner metadata update diff --git a/internal/index/seed_test.go b/internal/index/seed_test.go index f396c4a8..66ac13f4 100644 --- a/internal/index/seed_test.go +++ b/internal/index/seed_test.go @@ -88,6 +88,53 @@ func Hello() {} } } +func TestSeedFromDonor_SelectsCompleteSharedProject(t *testing.T) { + incompleteProject := t.TempDir() + completeProject := t.TempDir() + writeGoFile(t, completeProject, "main.go", "package main\n\nfunc Complete() {}\n") + + donorPath := filepath.Join(t.TempDir(), "donor.db") + emb := &mockEmbedder{dims: 4, model: "test-model"} + idx, err := NewIndexerForProject(donorPath, emb, 512, "int8", incompleteProject) + if err != nil { + t.Fatal(err) + } + if err := idx.store.SetMeta("root_hash", ""); err != nil { + _ = idx.Close() + t.Fatal(err) + } + if _, err := idx.Index(context.Background(), completeProject, false, nil); err != nil { + _ = idx.Close() + t.Fatal(err) + } + if err := idx.Close(); err != nil { + t.Fatal(err) + } + + destinationProject := t.TempDir() + dstPath := filepath.Join(t.TempDir(), "seeded.db") + seeded, err := SeedFromDonor(donorPath, dstPath, destinationProject) + if err != nil { + t.Fatal(err) + } + if !seeded { + t.Fatal("expected complete shared project to be selected as donor") + } + + db, err := sql.Open("sqlite3", sqliteFileDSN(dstPath, "ro")) + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + var projectPath string + if err := db.QueryRow(`SELECT path FROM projects WHERE path = ?`, destinationProject).Scan(&projectPath); err != nil { + t.Fatal(err) + } + if projectPath != destinationProject { + t.Fatalf("seeded project path = %q, want %q", projectPath, destinationProject) + } +} + func TestSeedFromDonor_SnapshotsCommittedWALWithActiveWriter(t *testing.T) { projectDir := t.TempDir() writeGoFile(t, projectDir, "main.go", "package main\n\nfunc Hello() {}\n") diff --git a/internal/store/shared.go b/internal/store/shared.go index b08f4a68..43fe856f 100644 --- a/internal/store/shared.go +++ b/internal/store/shared.go @@ -61,13 +61,13 @@ func openCollection(dsn string, dimensions int, vectorStorage string) (*Store, e } db.SetMaxOpenConns(1) for _, pragma := range []string{ + "PRAGMA busy_timeout=120000", "PRAGMA auto_vacuum=INCREMENTAL", "PRAGMA journal_mode=WAL", "PRAGMA foreign_keys=ON", "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA temp_store=MEMORY", - "PRAGMA busy_timeout=120000", } { if _, err := db.Exec(pragma); err != nil { _ = db.Close() @@ -194,36 +194,28 @@ func createCollectionSchema(db *sql.DB, dimensions int, vectorStorage string) er "vector_storage": vectorStorage, } for key, value := range want { + if _, err := db.Exec(`INSERT OR IGNORE INTO collection_meta(key, value) VALUES (?, ?)`, key, value); err != nil { + return fmt.Errorf("initialize collection profile %s: %w", key, err) + } var existing string - err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = ?`, key).Scan(&existing) - switch { - case err == sql.ErrNoRows: - if _, err := db.Exec(`INSERT INTO collection_meta(key, value) VALUES (?, ?)`, key, value); err != nil { - return err - } - case err != nil: - return err - case existing != value: + if err := db.QueryRow(`SELECT value FROM collection_meta WHERE key = ?`, key).Scan(&existing); err != nil { + return fmt.Errorf("read collection profile %s: %w", key, err) + } + if existing != value { return fmt.Errorf("collection profile mismatch for %s: stored %q, requested %q", key, existing, value) } } - exists, err := checkTableExists(db, "vec_vectors") - if err != nil { - return err + elementType := "int8" + if vectorStorage == "float32" { + elementType = "float" } - if !exists { - elementType := "int8" - if vectorStorage == "float32" { - elementType = "float" - } - stmt := fmt.Sprintf(`CREATE VIRTUAL TABLE vec_vectors USING vec0( - vector_id INTEGER PRIMARY KEY, - embedding %s[%d] distance_metric=cosine - )`, elementType, dimensions) - if _, err := db.Exec(stmt); err != nil { - return fmt.Errorf("create vec_vectors: %w", err) - } + stmt := fmt.Sprintf(`CREATE VIRTUAL TABLE IF NOT EXISTS vec_vectors USING vec0( + vector_id INTEGER PRIMARY KEY, + embedding %s[%d] distance_metric=cosine + )`, elementType, dimensions) + if _, err := db.Exec(stmt); err != nil { + return fmt.Errorf("create vec_vectors: %w", err) } return nil } @@ -498,7 +490,9 @@ func (s *Store) StoreFileRevision(relativePath, contentHash string, chunks []chu } h := embeddingInputHash(chunks[position]) var vectorID int64 - if err := tx.QueryRow(`SELECT id FROM vector_keys WHERE input_hash = ?`, h[:]).Scan(&vectorID); err != nil { + if err := tx.QueryRow(`SELECT id FROM vector_keys WHERE input_hash = ?`, h[:]).Scan(&vectorID); errors.Is(err, sql.ErrNoRows) { + return false, fmt.Errorf("%w: missing vector key for chunk %d (%s)", ErrVectorVanished, position, chunks[position].ID) + } else if err != nil { return false, err } blob, err := s.serializeVector(vec) diff --git a/internal/store/shared_test.go b/internal/store/shared_test.go index a547bdc9..1a8a40e7 100644 --- a/internal/store/shared_test.go +++ b/internal/store/shared_test.go @@ -7,6 +7,7 @@ package store import ( "context" + "errors" "math/rand" "os" "path/filepath" @@ -299,6 +300,66 @@ func TestSharedCollectionConcurrentRevisionInsertionIsIdempotent(t *testing.T) { } } +func TestSharedCollectionConcurrentFirstOpenIsIdempotent(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "index.db") + const openers = 8 + start := make(chan struct{}) + stores := make(chan *Store, openers) + errs := make(chan error, openers) + var wg sync.WaitGroup + for range openers { + project := t.TempDir() + wg.Add(1) + go func() { + defer wg.Done() + <-start + s, err := NewCollection(dbPath, 4, "int8", project) + if err != nil { + errs <- err + return + } + stores <- s + }() + } + close(start) + wg.Wait() + close(stores) + close(errs) + for s := range stores { + if !s.IsShared() { + t.Error("concurrent open returned a legacy store") + } + if err := s.Close(); err != nil { + t.Errorf("close concurrent store: %v", err) + } + } + for err := range errs { + t.Errorf("concurrent first open: %v", err) + } +} + +func TestSharedRefreshMapsMissingVectorKeyToSentinel(t *testing.T) { + s, err := NewCollection(":memory:", 4, "int8", t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = s.Close() }() + chunk := chunker.Chunk{ID: "gone", FilePath: "gone.go", Symbol: "Gone", Kind: "function", StartLine: 1, EndLine: 1, Content: "func Gone() {}"} + if _, err := s.StoreFileRevision("gone.go", "aa", []chunker.Chunk{chunk}, map[int][]float32{0: {1, 0, 0, 0}}); err != nil { + t.Fatal(err) + } + if _, err := s.db.Exec("PRAGMA foreign_keys=OFF"); err != nil { + t.Fatal(err) + } + if _, err := s.db.Exec("DELETE FROM vector_keys"); err != nil { + t.Fatal(err) + } + _, err = s.StoreFileRevision("gone.go", "aa", []chunker.Chunk{chunk}, map[int][]float32{0: {1, 0, 0, 0}}) + if !errors.Is(err, ErrVectorVanished) { + t.Fatalf("refresh error = %v, want ErrVectorVanished", err) + } +} + func TestSharedCleanupRemovesOnlyStaleMemberships(t *testing.T) { projectA, projectB := t.TempDir(), t.TempDir() s, err := NewCollection(":memory:", 4, "int8", projectA) @@ -404,6 +465,9 @@ func TestInt8RecallAt8AgainstFloat32(t *testing.T) { } func TestSharedInt8StorageAtMostTwentyPercentOfSeparateFloat32(t *testing.T) { + if testing.Short() { + t.Skip("storage size fixture writes three multi-megabyte databases") + } const ( dimensions = 768 chunkCount = 1000 diff --git a/internal/store/store.go b/internal/store/store.go index 185ab0aa..769e99c7 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -134,7 +134,9 @@ func New(dsn string, dimensions int) (*Store, error) { // project membership identified by projectPath. vectorStorage must be int8 or // float32. Multiple Store instances may safely select different worktrees in // the same database. If schema setup detects corruption, on-disk database and -// sidecar files are removed and creation is retried once. +// sidecar files are removed and creation is retried once. During lazy +// migration, opening a legacy per-worktree database returns a non-shared Store; +// callers that depend on project membership must check IsShared. func NewCollection(dsn string, dimensions int, vectorStorage, projectPath string) (*Store, error) { if vectorStorage != "int8" && vectorStorage != "float32" { return nil, fmt.Errorf("unsupported vector storage %q", vectorStorage)