Make rebuild --embed-only idempotent across crashes

CHANGELOG.md

@@ -3,7 +3,8 @@
 
 ### 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. A leftover staging table from an interrupted previous rebuild is dropped at the top of `rebuild_database` — note that an interruption mid-rebuild may still leave the chunks table in a partial state until a follow-up adds idempotent recovery from staging.
+- **`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` 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
 

haiku_rag_slim/haiku/rag/client/rebuild.py

@@ -8,6 +8,7 @@ from lancedb.pydantic import LanceModel
 
 from haiku.rag.client.documents import check_source_accessible
 from haiku.rag.converters import get_converter
+from haiku.rag.store.engine import ChunkRecordBase
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.models.document_item import extract_items
@@ -20,7 +21,8 @@ logger = logging.getLogger(__name__)
 
 _REBUILD_BATCH_SIZE = 50
 _STAGING_TABLE_NAME = "chunks_rebuild_staging"
-_STAGING_COPY_PAGE_SIZE = 1000
+_STAGING_MARKER_TABLE_NAME = "chunks_rebuild_marker"
+_STAGING_COPY_BATCH_SIZE = 1000
 
 
 class _StagingChunkRecord(LanceModel):
@@ -39,6 +41,19 @@ class _StagingChunkRecord(LanceModel):
     order: int
 
 
+class _StagingMarkerRecord(LanceModel):
+    """Sentinel marking the staging table as complete.
+
+    The marker table is created only after ``_populate_staging_table`` writes
+    every chunk into staging. Its presence at the top of a rebuild means
+    phase 2 (the embed loop) was interrupted by an earlier crash, so staging
+    is the authoritative source for the original chunk identities and we
+    must resume from it instead of rerunning phase 1.
+    """
+
+    id: str
+
+
 async def rebuild_database(
     client: "HaikuRAG", mode: "RebuildMode | None" = None
 ) -> AsyncGenerator[str, None]:
@@ -51,9 +66,11 @@ async def rebuild_database(
     if mode is None:
         mode = RebuildMode.FULL
 
-    # If a previous embed-only rebuild was interrupted, a leftover staging
+    # Resolve any leftover staging/marker tables from a previously
-    # table may remain. Drop it so the new rebuild starts from a clean slate.
+    # interrupted rebuild. Returns True only when phase 1 was already
-    await _drop_leftover_staging_table(client)
+    # complete and the current mode is EMBED_ONLY, in which case we resume
+    # phase 2 from the existing staging table instead of recopying.
+    resume_from_staging = await _resolve_rebuild_recovery(client, mode)
 
     # Wait for any already-scheduled background vacuum before the destructive
     # table operations at the top of RECHUNK / FULL. Rebuild drops and
@@ -78,7 +95,9 @@ async def rebuild_database(
         async for doc_id in _rebuild_title_only(client, documents):
             yield doc_id
     elif mode == RebuildMode.EMBED_ONLY:
-        async for doc_id in _rebuild_embed_only(client, documents):
+        async for doc_id in _rebuild_embed_only(
+            client, documents, resume_from_staging=resume_from_staging
+        ):
             yield doc_id
     elif mode == RebuildMode.RECHUNK:
         await client.chunk_repository.delete_all()
@@ -145,29 +164,73 @@ async def _rebuild_title_only(
             yield doc.id
 
 
-async def _drop_leftover_staging_table(client: "HaikuRAG") -> None:
+async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool:
-    """Drop a stale staging table left over from an interrupted rebuild.
+    """Resolve any partially-completed rebuild state from a previous crash.
 
-    Same crash-recovery semantics as the other rebuild modes (RECHUNK/FULL
+    Returns ``True`` if ``_rebuild_embed_only`` should resume from the
-    also leave the database non-atomic if interrupted): the user is expected
+    existing staging table (phase 1 was already complete). In all other
-    to re-run rebuild from whatever state the database is in.
+    cases stale recovery tables are dropped and the rebuild starts fresh.
+
+    State at entry           → action
+    --------------------------------
+    no staging, no marker    → return False (normal start)
+    staging only             → drop staging (phase 1 was interrupted; ``chunks`` is intact)
+    marker only              → drop marker (corrupted state)
+    staging + marker, embed  → return True (resume phase 2 from staging)
+    staging + marker, other  → drop both (staging is for embed-only; user picked a different mode)
     """
-    tables = (await client.store.db.list_tables()).tables
+    from haiku.rag.client import RebuildMode
-    if _STAGING_TABLE_NAME not in tables:
+
-        return
+    db = client.store.db
+    tables = (await db.list_tables()).tables
+    has_staging = _STAGING_TABLE_NAME in tables
+    has_marker = _STAGING_MARKER_TABLE_NAME in tables
+
+    if not has_staging and not has_marker:
+        return False
+
+    if has_marker and not has_staging:
+        logger.warning(
+            "Found '%s' without staging table; dropping orphaned marker.",
+            _STAGING_MARKER_TABLE_NAME,
+        )
+        await db.drop_table(_STAGING_MARKER_TABLE_NAME)
+        return False
+
+    if not has_marker:
+        logger.warning(
+            "Dropping incomplete '%s' from an interrupted phase 1.",
+            _STAGING_TABLE_NAME,
+        )
+        await db.drop_table(_STAGING_TABLE_NAME)
+        return False
+
+    # has_staging and has_marker
+    if mode == RebuildMode.EMBED_ONLY:
+        logger.warning(
+            "Resuming interrupted embed-only rebuild: phase 2 will run from "
+            "existing '%s'.",
+            _STAGING_TABLE_NAME,
+        )
+        return True
+
     logger.warning(
-        "Dropping leftover '%s' table from a previous interrupted rebuild.",
+        "Dropping staging tables from a prior embed-only rebuild — current "
-        _STAGING_TABLE_NAME,
+        "mode (%s) does not consume them.",
+        mode.name,
     )
-    await client.store.db.drop_table(_STAGING_TABLE_NAME)
+    await db.drop_table(_STAGING_MARKER_TABLE_NAME)
+    await db.drop_table(_STAGING_TABLE_NAME)
+    return False
 
 
 async def _populate_staging_table(client: "HaikuRAG") -> None:
-    """Stream-copy the non-vector columns of the chunks table into staging.
+    """Stream the non-vector columns of the chunks table into staging.
 
-    Reads ``_STAGING_COPY_PAGE_SIZE`` rows per page so peak memory stays
+    Uses ``to_batches`` for a single streaming read (no offset/limit
-    bounded regardless of corpus size. The vector column is omitted — the
+    pagination drift), so peak memory stays bounded regardless of corpus
-    point of embed-only rebuild is to regenerate it.
+    size. The vector column is omitted — the point of embed-only rebuild is
+    to regenerate it.
     """
     db = client.store.db
     tables = (await db.list_tables()).tables
@@ -178,17 +241,15 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
     if "chunks" not in tables:
         return
 
-    offset = 0
+    stream = (
-    while True:
+        await client.store.chunks_table.query()
-        rows = (
+        .select(["id", "document_id", "content", "metadata", "order"])
-            await client.store.chunks_table.query()
+        .to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE)
-            .select(["id", "document_id", "content", "metadata", "order"])
+    )
-            .offset(offset)
+    async for batch in stream:
-            .limit(_STAGING_COPY_PAGE_SIZE)
+        rows = batch.to_pylist()
-            .to_arrow()
-        ).to_pylist()
         if not rows:
-            break
+            continue
         records = [
             _StagingChunkRecord(
                 id=r["id"],
@@ -200,7 +261,37 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
             for r in rows
         ]
         await staging.add(records)
-        offset += _STAGING_COPY_PAGE_SIZE
+
+
+async def _mark_phase1_complete(client: "HaikuRAG") -> None:
+    """Create the marker table that designates staging as authoritative.
+
+    Called after ``_populate_staging_table`` finishes. On crash recovery the
+    marker's presence flips ``_rebuild_embed_only`` into resume mode.
+    """
+    db = client.store.db
+    if _STAGING_MARKER_TABLE_NAME in (await db.list_tables()).tables:
+        return
+    marker = await db.create_table(
+        _STAGING_MARKER_TABLE_NAME, schema=_StagingMarkerRecord
+    )
+    await marker.add([_StagingMarkerRecord(id="phase1_complete")])
+
+
+async def _drop_staging_tables(client: "HaikuRAG") -> None:
+    """Drop the marker first, then the staging table.
+
+    Ordering matters: if a crash interrupts cleanup between the two drops,
+    the next rebuild sees ``staging`` without ``marker`` and treats it as a
+    partial phase 1 → drops staging harmlessly. The reverse order would
+    leak a marker pointing at nothing.
+    """
+    db = client.store.db
+    tables = (await db.list_tables()).tables
+    if _STAGING_MARKER_TABLE_NAME in tables:
+        await db.drop_table(_STAGING_MARKER_TABLE_NAME)
+    if _STAGING_TABLE_NAME in tables:
+        await db.drop_table(_STAGING_TABLE_NAME)
 
 
 async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chunk]:
@@ -232,82 +323,101 @@ async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chu
 
 
 async def _rebuild_embed_only(
-    client: "HaikuRAG", documents: list[Document]
+    client: "HaikuRAG",
+    documents: list[Document],
+    *,
+    resume_from_staging: bool = False,
 ) -> AsyncGenerator[str, None]:
     """Re-embed all chunks without changing chunk boundaries.
 
     Two-phase pattern that keeps peak memory bounded regardless of corpus
-    size:
+    size and is idempotent across crashes:
 
     1. Stream-copy the chunks table's non-vector columns into a staging
-       table. LanceDB OSS doesn't support ``rename_table``, so this is the
+       table, then write a marker row that designates staging as complete.
-       only safe way to preserve the original chunk identity while we drop
+       LanceDB OSS does not support ``rename_table``, so the staging copy
-       and recreate the live chunks table with a (potentially) new vector
+       is the only safe way to preserve chunk identity while the live
-       dim.
+       ``chunks`` table is dropped and recreated.
-    2. Drop and recreate ``chunks`` with the current schema, then stream
+    2. Drop-and-recreate ``chunks`` with the current schema, then stream
        from staging one document at a time, re-embed in batches of
        ``embeddings.batch_size``, and flush to the new chunks table every
        ``_REBUILD_BATCH_SIZE`` documents.
+
+    Cleanup runs only on success: a crash anywhere in phase 2 leaves both
+    staging and marker in place so the next rebuild can re-enter phase 2
+    via ``resume_from_staging=True``. The order of the success cleanup —
+    drop marker before staging — keeps an interruption between the two
+    drops recoverable: the next rebuild sees staging without marker and
+    treats it as a partial phase 1, which is harmless because phase 2 has
+    already finished writing the new chunks table.
     """
     from haiku.rag.embeddings import contextualize
 
     db = client.store.db
     batch_size = client._config.embeddings.batch_size
 
-    # Phase 1: copy chunks into staging. After this we can safely destroy
+    if not resume_from_staging:
-    # the live chunks table without losing chunk identity / content.
+        # Phase 1: copy chunks into staging, then mark it complete. After the
-    await _populate_staging_table(client)
+        # marker exists, a crash will resume phase 2 from staging.
+        await _populate_staging_table(client)
+        await _mark_phase1_complete(client)
 
-    # Recreate the chunks table fresh (handles vector-dim changes).
+    # Recreate the chunks table fresh (idempotent; handles vector-dim
+    # changes and discards any partial new chunks from a prior crashed
+    # phase 2).
     await client.store.recreate_embeddings_table()
 
     staging_table = await db.open_table(_STAGING_TABLE_NAME)
 
-    pending_records: list = []
+    pending_records: list[ChunkRecordBase] = []
     yielded_docs: set[str] = set()
 
-    try:
+    for doc in documents:
-        for doc in documents:
+        assert doc.id is not None
-            assert doc.id is not None
+        chunks = await _read_chunks_from_staging(staging_table, doc.id)
-            chunks = await _read_chunks_from_staging(staging_table, doc.id)
+        if not chunks:
-            if not chunks:
+            continue
-                continue
 
-            texts = contextualize(chunks)
+        texts = contextualize(chunks)
-            embeddings: list[list[float]] = []
+        embeddings: list[list[float]] = []
-            for i in range(0, len(texts), batch_size):
+        for i in range(0, len(texts), batch_size):
-                batch_embeddings = (
+            batch_embeddings = await client.chunk_repository.embedder.embed_documents(
-                    await client.chunk_repository.embedder.embed_documents(
+                texts[i : i + batch_size]
-                        texts[i : i + batch_size]
+            )
-                    )
+            embeddings.extend(batch_embeddings)
+
+        for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
+            pending_records.append(
+                client.store.ChunkRecord(
+                    id=chunk.id,
+                    document_id=chunk.document_id,
+                    content=chunk.content,
+                    content_fts=content_fts,
+                    metadata=json.dumps(chunk.metadata),
+                    order=chunk.order,
+                    vector=embedding,
                 )
-                embeddings.extend(batch_embeddings)
+            )
 
-            for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
+        yielded_docs.add(doc.id)
-                pending_records.append(
+        # Yield per-doc for progress reporting; the actual write batches up
-                    client.store.ChunkRecord(
+        # to _REBUILD_BATCH_SIZE docs. If the process is interrupted between
-                        id=chunk.id,
+        # yield and the next flush, the next rebuild resumes phase 2 from
-                        document_id=chunk.document_id,
+        # the staging table and redoes the batch (see _rebuild_rechunk for
-                        content=chunk.content,
+        # the original comment on the yield/flush gap).
-                        content_fts=content_fts,
+        yield doc.id
-                        metadata=json.dumps(chunk.metadata),
-                        order=chunk.order,
-                        vector=embedding,
-                    )
-                )
 
-            yielded_docs.add(doc.id)
+        if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
-            yield doc.id
-
-            if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
-                await client.store.chunks_table.add(pending_records)
-                pending_records = []
-
-        if pending_records:
             await client.store.chunks_table.add(pending_records)
-    finally:
+            pending_records = []
-        if _STAGING_TABLE_NAME in (await db.list_tables()).tables:
+
-            await db.drop_table(_STAGING_TABLE_NAME)
+    if pending_records:
+        await client.store.chunks_table.add(pending_records)
+
+    # Phase 2 finished. Drop the recovery state — marker first so a crash
+    # between the two drops leaves only staging behind, which the next
+    # rebuild discards harmlessly.
+    await _drop_staging_tables(client)
 
     # Yield docs with no chunks
     for doc in documents:

tests/test_rebuild.py

@@ -1,3 +1,4 @@
+import json
 import tempfile
 from pathlib import Path
 from typing import TypedDict
@@ -133,11 +134,11 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
 
 @pytest.mark.vcr()
 async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
-    """A staging table left behind by a previous interrupted rebuild gets dropped.
+    """Staging table without marker is treated as partial phase 1 and dropped.
 
-    Simulates a crash mid-rebuild by manually creating a ``chunks_rebuild_staging``
+    Simulates a phase-1 interruption by creating only the staging table (no
-    table before a fresh rebuild starts. ``rebuild_database`` should detect
+    marker). On the next rebuild ``_resolve_rebuild_recovery`` should drop
-    and drop it before doing anything else.
+    the partial staging — the live chunks table is still authoritative.
     """
     from haiku.rag.client.rebuild import _StagingChunkRecord
 
@@ -145,11 +146,13 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
         doc = await client.create_document(content=qa_corpus["document_extracted"][0])
         assert doc.id is not None
 
-        # Simulate a stale staging table from an interrupted rebuild.
+        # Simulate a partial phase 1 (staging exists, marker absent).
         await client.store.db.create_table(
             "chunks_rebuild_staging", schema=_StagingChunkRecord
         )
-        assert "chunks_rebuild_staging" in (await client.store.db.list_tables()).tables
+        tables = (await client.store.db.list_tables()).tables
+        assert "chunks_rebuild_staging" in tables
+        assert "chunks_rebuild_marker" not in tables
 
         processed_ids = [
             doc_id
@@ -157,9 +160,145 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
         ]
         assert doc.id in processed_ids
 
-        # Stale staging gone, fresh staging from this rebuild also gone.
         tables = (await client.store.db.list_tables()).tables
         assert "chunks_rebuild_staging" not in tables
+        assert "chunks_rebuild_marker" not in tables
+
+
+@pytest.mark.vcr()
+async def test_rebuild_resumes_phase2_from_staging_after_crash(
+    qa_corpus: Dataset, temp_db_path
+):
+    """Marker + staging present → phase 2 resumes from staging instead of
+    redoing phase 1.
+
+    Simulates a phase-2 crash: pre-populate staging with the original chunks,
+    create the marker, then drop the live chunks table entirely (the worst
+    case — crash right after ``recreate_embeddings_table`` succeeded but
+    before any phase-2 batch flushed). The rebuild must reconstruct the
+    chunks table from staging without losing data.
+    """
+    from haiku.rag.client.rebuild import (
+        _StagingChunkRecord,
+        _StagingMarkerRecord,
+    )
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.create_document(content=qa_corpus["document_extracted"][0])
+        assert doc.id is not None
+        original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        assert original_chunks
+        original_ids = {c.id for c in original_chunks}
+        original_contents = {c.id: c.content for c in original_chunks}
+
+        # Snapshot chunk data into staging (simulating phase 1's output).
+        staging = await client.store.db.create_table(
+            "chunks_rebuild_staging", schema=_StagingChunkRecord
+        )
+        await staging.add(
+            [
+                _StagingChunkRecord(
+                    id=c.id or "",
+                    document_id=c.document_id or "",
+                    content=c.content,
+                    metadata=json.dumps(c.metadata),
+                    order=c.order,
+                )
+                for c in original_chunks
+            ]
+        )
+
+        # Mark phase 1 complete (simulating the marker write that happens
+        # just before phase 2 starts).
+        marker = await client.store.db.create_table(
+            "chunks_rebuild_marker", schema=_StagingMarkerRecord
+        )
+        await marker.add([_StagingMarkerRecord(id="phase1_complete")])
+
+        # Wipe the live chunks table to simulate a worst-case phase-2 crash
+        # after recreate_embeddings_table but before any chunks were
+        # written.
+        await client.store.db.drop_table("chunks")
+
+        # Recovery: rebuild_database should detect marker+staging and have
+        # _rebuild_embed_only skip phase 1.
+        processed_ids = [
+            doc_id
+            async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
+        ]
+        assert doc.id in processed_ids
+
+        recovered = await client.chunk_repository.get_by_document_id(doc.id)
+        assert {c.id for c in recovered} == original_ids
+        for chunk in recovered:
+            assert chunk.content == original_contents[chunk.id]
+
+        tables = (await client.store.db.list_tables()).tables
+        assert "chunks_rebuild_staging" not in tables
+        assert "chunks_rebuild_marker" not in tables
+
+
+async def test_rebuild_drops_orphan_marker(temp_db_path):
+    """Marker without staging is treated as corrupted and dropped.
+
+    No embeddings are needed: ``_resolve_rebuild_recovery`` decides on
+    tables before any embed call, and the empty database has no documents
+    to embed.
+    """
+    from haiku.rag.client.rebuild import _StagingMarkerRecord
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        marker = await client.store.db.create_table(
+            "chunks_rebuild_marker", schema=_StagingMarkerRecord
+        )
+        await marker.add([_StagingMarkerRecord(id="phase1_complete")])
+
+        _ = [
+            doc_id
+            async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
+        ]
+
+        tables = (await client.store.db.list_tables()).tables
+        assert "chunks_rebuild_marker" not in tables
+        assert "chunks_rebuild_staging" not in tables
+
+
+@pytest.mark.vcr()
+async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
+    qa_corpus: Dataset, temp_db_path
+):
+    """Staging + marker from a prior embed-only crash → dropped on RECHUNK.
+
+    If a user runs a different rebuild mode after a crashed embed-only, the
+    staging tables are stale: the new mode recreates chunks from a
+    different source (e.g. the stored docling blob), so the staging copy is
+    not useful.
+    """
+    from haiku.rag.client.rebuild import (
+        _StagingChunkRecord,
+        _StagingMarkerRecord,
+    )
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.create_document(content=qa_corpus["document_extracted"][0])
+        assert doc.id is not None
+
+        await client.store.db.create_table(
+            "chunks_rebuild_staging", schema=_StagingChunkRecord
+        )
+        marker = await client.store.db.create_table(
+            "chunks_rebuild_marker", schema=_StagingMarkerRecord
+        )
+        await marker.add([_StagingMarkerRecord(id="phase1_complete")])
+
+        processed_ids = [
+            doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
+        ]
+        assert doc.id in processed_ids
+
+        tables = (await client.store.db.list_tables()).tables
+        assert "chunks_rebuild_staging" not in tables
+        assert "chunks_rebuild_marker" not in tables
 
 
 @pytest.mark.vcr()