From dbb87e091f836476bd066fd867f78670afe2b636 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 14 May 2026 15:02:02 +0300 Subject: [PATCH] Ensure schema parity for _StagingChunkRecord --- CHANGELOG.md | 2 +- haiku_rag_slim/haiku/rag/client/rebuild.py | 8 +++++++- tests/test_rebuild.py | 15 +++++++++++++++ 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77d3bbce..393b9be7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ### Fixed -- **`rebuild --embed-only` no longer buffers the entire corpus in memory.** The previous implementation accumulated every chunk's id, content, content_fts, metadata, and new embedding vector in a single Python list before flushing — on a multi-million-chunk corpus with a 3072-dim model this trivially exceeded 40 GB of RAM. The rebuild now stream-copies non-vector columns into a `chunks_rebuild_staging` table (1000 rows / page), recreates the chunks table fresh to honour vector-dim changes, then streams from staging one document at a time, embedding in batches of `embeddings.batch_size` and flushing to the new chunks table every 50 documents. Peak memory is bounded to one batch regardless of corpus size; LanceDB OSS does not support `rename_table`, so the staging copy is the only safe way to preserve chunk identity while the live table is recreated. +- **`rebuild --embed-only` no longer buffers the entire corpus in memory.** The previous implementation accumulated every chunk's id, content, content_fts, metadata, and new embedding vector in a single Python list before flushing. The rebuild now stream-copies non-vector columns into a `chunks_rebuild_staging` table (1000 rows / page), recreates the chunks table fresh to honour vector-dim changes, then streams from staging one document at a time, embedding in batches of `embeddings.batch_size` and flushing to the new chunks table every 50 documents. - **`rebuild --embed-only` is now idempotent across crashes.** A second table, `chunks_rebuild_marker`, is written immediately after phase 1 (staging copy) finishes. Its presence flips the next rebuild into resume mode: phase 1 is skipped, the live chunks table is recreated, and phase 2 (re-embed) runs from the existing staging snapshot. Cleanup drops the marker before the staging table, so an interruption between the two drops leaves a markerless staging that the next run discards harmlessly. A staging table without a marker is treated as a partial phase 1 and dropped (the live chunks table is still authoritative). Running a non-embed-only mode (FULL / RECHUNK / DESCRIPTIONS / TITLE_ONLY) after a crashed embed-only correctly discards the staging recovery state. Phase 1's pagination was switched from `offset/limit` to `to_batches`, removing the latent offset-drift risk and the O(N²) cost at high offsets. ## [0.46.0] - 2026-05-13 diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index f5f4897b..af0fa3d3 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -31,7 +31,13 @@ class _StagingChunkRecord(LanceModel): The staging table holds the original chunks' identity and content while the live ``chunks`` table is dropped and recreated with a potentially different vector dimension. The vector itself is omitted — re-embedding - is the whole point. + is the whole point — and ``content_fts`` is regenerated by + ``contextualize`` during phase 2. + + Mirrors ``ChunkRecordBase`` minus ``content_fts`` and ``vector``. Keep + in sync: ``test_staging_chunk_record_mirrors_chunk_record_schema`` + enforces parity so a new column on ``ChunkRecordBase`` can't silently + get dropped on every embed-only rebuild. """ id: str diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index f1fd8450..7af1469e 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -238,6 +238,21 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash( assert "chunks_rebuild_marker" not in tables +def test_staging_chunk_record_mirrors_chunk_record_schema(): + """``_StagingChunkRecord`` must hold every ``ChunkRecordBase`` field except + those that are re-derived (``content_fts``) or replaced (``vector``). + + If someone adds a column to ``ChunkRecordBase`` without updating + ``_StagingChunkRecord``, embed-only rebuilds will silently drop that + column on every crash-recovery cycle. This test fails loudly instead. + """ + from haiku.rag.client.rebuild import _StagingChunkRecord + from haiku.rag.store.engine import ChunkRecordBase + + expected = set(ChunkRecordBase.model_fields) - {"content_fts", "vector"} + assert set(_StagingChunkRecord.model_fields) == expected + + async def test_rebuild_drops_orphan_marker(temp_db_path): """Marker without staging is treated as corrupted and dropped.