fix(sqlite_exact): speed up query and status on large palaces - #2308
Conversation
sqlite_exact.query() loaded every document and metadata blob to do Python-loop cosine on the whole collection. On a 167k-drawer palace that made mempalace_search 6-9s and mempalace_status 7s (paging every metadata row because sqlite_exact had no facet_counts). Rank from the embedding column with vectorized numpy cosine, hydrate only the top-k documents, and cache the matrix plus wing/room/source_file on the long-lived handle. Add facet_counts and sqlite_wing_room_counts so status/list_wings use one GROUP BY. Advertise supports_metadata_facets and expose collection._backend so the MCP facet path runs through the embedding wrapper. Live palace (167k drawers / 166k closets): status 6997ms -> 1045ms search warm 6364ms -> 210-1600ms search cold 8236ms -> 5299ms (first matrix load)
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
igorls
left a comment
igorls
left a comment
There was a problem hiding this comment.
Read the whole diff and ran the numbers locally. The approach is right — ranking from the embedding column and hydrating only top-k is the correct shape, the SQL pushdown is properly guarded (_FACET_FIELD_RE on every interpolated key, so no injection surface), and _rows already used ORDER BY rowid so the stable-sort tie-breaking is preserved. Five things worth a look before merge, roughly in order of how much they matter.
1. The cache is cleared on every write, so the file-then-search cycle is always cold.
_cursor(write=True) does _vector_cache.clear(), and filing a single drawer is a write. On a box where hooks file drawers continuously, the warm path may be the exception rather than the rule.
Measured on a synthetic 40k-drawer / 384-dim palace:
COLD query: 99 ms
WARM query: 6 ms
query AFTER ONE WRITE: 90 ms (15x the warm path)
So the 210–1600 ms warm figure in the description is real but may be rare in practice, and 5299 ms cold is closer to the steady state for a machine that is actively filing. Worth considering an incremental update — append the new rows to the matrix instead of dropping all of it — since an upsert usually adds a handful of vectors to a 166k-row matrix.
2. The cache has no cap and no opt-out.
At 166k drawers x 384 float32 that is ~255 MB per collection, so ~510 MB across drawers + closets, pinned on the handle for the life of the hub and growing linearly with the palace.
The load transient is larger than the matrix itself — .fetchall() materializes every row (id, blob, three json_extract results) before np.stack allocates the copy:
40k drawers, 61 MB matrix -> +197 MB peak RSS on the cold load (~3.2x)
extrapolated to 166k -> ~255 MB resident, ~800 MB transient per collection
Given the memory work that just landed, a half-gig floor plus a multi-hundred-meg spike per cold load seems worth at least an env knob or a row-count ceiling above which the backend keeps the old streaming path.
3. np.stack raises when collections.dimension is NULL and rows have mixed widths.
_load_all_vectors only skips size-mismatched vectors when expected is not None. With a NULL dimension every row is kept and np.stack fails:
ValueError: all input arrays must have the same shape
The old per-query loop skipped mismatches (vec.size != q.size) and degraded gracefully. Reachability is narrow — _ensure_collection_dimension pins the dimension on first write and rejects mixed batches, so it needs a legacy or externally-modified DB — but it turns a soft skip into a hard failure. documents.dim is already a NOT NULL column, so filtering in SQL (WHERE dim = ?) or falling back to the first kept vector's size would close it.
4. Latent: a concurrent write can be swallowed by the cache.
_rank_vectors does get-miss, load, then store. A write committing between the load and the store clears a cache that is then repopulated with the pre-write snapshot, and nothing invalidates it until the next write — so a just-filed drawer stays invisible to search.
Not reachable through the hub today: _http_dispatch serializes every tool call on _HTTP_REQUEST_LOCK, and reads take handle.lock while writes take mine_palace_lock. But the cache's correctness currently depends on a lock in a different module rather than on anything local. A generation counter — same pattern as _capacity_cache_generation in the chroma backend, bump on clear, only store if unchanged — would make it self-contained.
5. Minor: the cosine is now computed entirely in float32.
dots, norms and the division all stay float32, where the old path did float(np.dot(...)) and divided in float64. Same formula, marginally different rounding, so near-ties can reorder. The ranking tests pass and this is almost certainly below any threshold that matters — noting it only so it is a deliberate choice rather than an accident.
Nothing here is an objection to the direction: mempalace_status at 6997 ms -> 1045 ms is a real win and the facet pushdown is clearly the right call.
sqlite_exact is the live backend on this machine (~167k drawers, ~166k closets, 1.6 GB).
query()was exact cosine over every row: it selectedid, document, metadata_json, embedding, JSON-parsed metadata, then dotted in a Python loop. MCPmempalace_searchdoes that twice (drawers + closets).mempalace_statuspaged every metadata row because sqlite_exact had nofacet_countsand_sqlite_taxonomywas chroma-only.This PR does not touch
chroma.py(in-flight work on develop, including #2307).Changes
wing/room/source_fileon the long-lived backend handle (invalidated on write).facet_counts+sqlite_wing_room_countsso status / list_wings / list_rooms use onejson_extractGROUP BY.supports_metadata_facetsand setcollection._backendso the MCP facet path works throughEmbeddingCollection.Exact cosine ranking is unchanged (existing ranking tests plus new ones).
Live palace (this box, MCP HTTP, 167k drawers / 166k closets)
mempalace_statusmempalace_searchcoldmempalace_searchwarmmempalace_searchwing=mempalacemempalace_kg_queryIsolated: FTS MATCH ~99 ms warm; embeddings-only numpy cosine ~100–186 ms once vectors are in RAM. The remaining warm-search spread is embed + two collection cosines + MCP JSON, not a full table scan of document text.
Tests
uv run pytest tests/test_sqlite_exact_backend.py tests/test_searcher.py tests/test_hybrid_search.py tests/test_mcp_server.py— 517 passed.