Merge pull request #372 from ggozad/fix/rebuild-embed-memory
Bound memory usage during embeddings rebuild
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
0e4c462885
7 changed files with 1312 additions and 31 deletions
|
|
@ -1,6 +1,11 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### 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. 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
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -4,8 +4,11 @@ from collections.abc import AsyncGenerator
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
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
|
||||
|
|
@ -17,6 +20,44 @@ if TYPE_CHECKING:
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
_REBUILD_BATCH_SIZE = 50
|
||||
_STAGING_TABLE_NAME = "chunks_rebuild_staging"
|
||||
_STAGING_MARKER_TABLE_NAME = "chunks_rebuild_marker"
|
||||
_STAGING_COPY_BATCH_SIZE = 1000
|
||||
|
||||
|
||||
class _StagingChunkRecord(LanceModel):
|
||||
"""Non-vector copy of a chunk row, used by ``_rebuild_embed_only``.
|
||||
|
||||
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 — 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
|
||||
document_id: str
|
||||
content: str
|
||||
metadata: str
|
||||
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(
|
||||
|
|
@ -31,6 +72,12 @@ async def rebuild_database(
|
|||
if mode is None:
|
||||
mode = RebuildMode.FULL
|
||||
|
||||
# Resolve any leftover staging/marker tables from a previously
|
||||
# interrupted rebuild. Returns True only when phase 1 was already
|
||||
# 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
|
||||
# recreates tables (and creates indices); a concurrent optimize on the
|
||||
|
|
@ -54,7 +101,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()
|
||||
|
|
@ -121,54 +170,260 @@ async def _rebuild_title_only(
|
|||
yield doc.id
|
||||
|
||||
|
||||
async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool:
|
||||
"""Resolve any partially-completed rebuild state from a previous crash.
|
||||
|
||||
Returns ``True`` if ``_rebuild_embed_only`` should resume from the
|
||||
existing staging table (phase 1 was already complete). In all other
|
||||
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)
|
||||
"""
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
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 staging tables from a prior embed-only rebuild — current "
|
||||
"mode (%s) does not consume them.",
|
||||
mode.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 the non-vector columns of the chunks table into staging.
|
||||
|
||||
Uses ``to_batches`` for a single streaming read (no offset/limit
|
||||
pagination drift), so peak memory stays bounded regardless of corpus
|
||||
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
|
||||
if _STAGING_TABLE_NAME in tables:
|
||||
await db.drop_table(_STAGING_TABLE_NAME)
|
||||
|
||||
staging = await db.create_table(_STAGING_TABLE_NAME, schema=_StagingChunkRecord)
|
||||
if "chunks" not in tables:
|
||||
return
|
||||
|
||||
stream = (
|
||||
await client.store.chunks_table.query()
|
||||
.select(["id", "document_id", "content", "metadata", "order"])
|
||||
.to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE)
|
||||
)
|
||||
async for batch in stream:
|
||||
rows = batch.to_pylist()
|
||||
if not rows:
|
||||
continue
|
||||
records = [
|
||||
_StagingChunkRecord(
|
||||
id=r["id"],
|
||||
document_id=r["document_id"],
|
||||
content=r["content"],
|
||||
metadata=r["metadata"],
|
||||
order=r["order"],
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
await staging.add(records)
|
||||
|
||||
|
||||
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]:
|
||||
"""Read chunks for one document from the staging table.
|
||||
|
||||
Only non-vector columns are selected: the staging table may have a
|
||||
different vector dimension than the new chunks table (during a dim
|
||||
migration), and we re-embed anyway.
|
||||
"""
|
||||
rows = (
|
||||
await staging_table.query()
|
||||
.where(f"document_id = '{document_id}'")
|
||||
.select(["id", "document_id", "content", "metadata", "order"])
|
||||
.to_arrow()
|
||||
).to_pylist()
|
||||
chunks: list[Chunk] = []
|
||||
for row in rows:
|
||||
chunks.append(
|
||||
Chunk(
|
||||
id=row["id"],
|
||||
document_id=row["document_id"],
|
||||
content=row["content"],
|
||||
metadata=json.loads(row["metadata"]),
|
||||
order=row["order"],
|
||||
)
|
||||
)
|
||||
chunks.sort(key=lambda c: c.order)
|
||||
return chunks
|
||||
|
||||
|
||||
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."""
|
||||
"""Re-embed all chunks without changing chunk boundaries.
|
||||
|
||||
Two-phase pattern that keeps peak memory bounded regardless of corpus
|
||||
size and is idempotent across crashes:
|
||||
|
||||
1. Stream-copy the chunks table's non-vector columns into a staging
|
||||
table, then write a marker row that designates staging as complete.
|
||||
LanceDB OSS does not support ``rename_table``, so the staging copy
|
||||
is the only safe way to preserve chunk identity while the live
|
||||
``chunks`` table is dropped and recreated.
|
||||
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
|
||||
|
||||
# Collect all chunks with new embeddings
|
||||
all_chunk_data: list[tuple[str, dict]] = []
|
||||
db = client.store.db
|
||||
batch_size = client._config.embeddings.batch_size
|
||||
|
||||
if not resume_from_staging:
|
||||
# Phase 1: copy chunks into staging, then mark it complete. After the
|
||||
# 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 (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[ChunkRecordBase] = []
|
||||
yielded_docs: set[str] = set()
|
||||
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
chunks = await _read_chunks_from_staging(staging_table, doc.id)
|
||||
if not chunks:
|
||||
continue
|
||||
|
||||
texts = contextualize(chunks)
|
||||
embeddings = await client.chunk_repository.embedder.embed_documents(texts)
|
||||
embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch_embeddings = await client.chunk_repository.embedder.embed_documents(
|
||||
texts[i : i + batch_size]
|
||||
)
|
||||
embeddings.extend(batch_embeddings)
|
||||
|
||||
for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
|
||||
all_chunk_data.append(
|
||||
(
|
||||
doc.id,
|
||||
{
|
||||
"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,
|
||||
},
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
# Recreate chunks table (handles dimension changes)
|
||||
await client.store.recreate_embeddings_table()
|
||||
yielded_docs.add(doc.id)
|
||||
# Yield per-doc for progress reporting; the actual write batches up
|
||||
# to _REBUILD_BATCH_SIZE docs. If the process is interrupted between
|
||||
# yield and the next flush, the next rebuild resumes phase 2 from
|
||||
# the staging table and redoes the batch (see _rebuild_rechunk for
|
||||
# the original comment on the yield/flush gap).
|
||||
yield doc.id
|
||||
|
||||
# Insert all chunks
|
||||
if all_chunk_data:
|
||||
records = [client.store.ChunkRecord(**data) for _, data in all_chunk_data]
|
||||
await client.store.chunks_table.add(records)
|
||||
if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
pending_records = []
|
||||
|
||||
# Yield all processed doc IDs
|
||||
yielded_docs: set[str] = set()
|
||||
for doc_id, _ in all_chunk_data:
|
||||
if doc_id not in yielded_docs:
|
||||
yielded_docs.add(doc_id)
|
||||
yield doc_id
|
||||
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:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
|
@ -83,6 +84,238 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
|
|||
assert chunk.content == chunk_contents_before[chunk.id]
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only_multi_doc_streams_via_staging(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
):
|
||||
"""Embed-only rebuild with multiple documents preserves chunks via staging.
|
||||
|
||||
Regression guard for the OOM bug: the previous implementation buffered
|
||||
all chunks across all documents in memory before flushing. The current
|
||||
streaming implementation copies chunks to a staging table, recreates
|
||||
the chunks table, then streams doc-by-doc. This test verifies:
|
||||
|
||||
- chunks survive across multiple documents (correctness),
|
||||
- the staging table is dropped at the end (no leak), and
|
||||
- the rebuild yields every document with chunks.
|
||||
"""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc1 = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
doc2 = await client.create_document(content=qa_corpus["document_extracted"][1])
|
||||
assert doc1.id is not None and doc2.id is not None
|
||||
|
||||
chunks_before_1 = await client.chunk_repository.get_by_document_id(doc1.id)
|
||||
chunks_before_2 = await client.chunk_repository.get_by_document_id(doc2.id)
|
||||
assert chunks_before_1 and chunks_before_2
|
||||
ids_before = {c.id for c in chunks_before_1} | {c.id for c in chunks_before_2}
|
||||
|
||||
processed_ids = [
|
||||
doc_id
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
|
||||
]
|
||||
|
||||
assert doc1.id in processed_ids
|
||||
assert doc2.id in processed_ids
|
||||
|
||||
chunks_after_1 = await client.chunk_repository.get_by_document_id(doc1.id)
|
||||
chunks_after_2 = await client.chunk_repository.get_by_document_id(doc2.id)
|
||||
ids_after = {c.id for c in chunks_after_1} | {c.id for c in chunks_after_2}
|
||||
|
||||
# Same chunk IDs survive; content unchanged.
|
||||
assert ids_before == ids_after
|
||||
contents_before = {c.id: c.content for c in chunks_before_1 + chunks_before_2}
|
||||
for chunk in chunks_after_1 + chunks_after_2:
|
||||
assert chunk.content == contents_before[chunk.id]
|
||||
|
||||
# Staging table was cleaned up.
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_staging" not in tables
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
|
||||
"""Staging table without marker is treated as partial phase 1 and dropped.
|
||||
|
||||
Simulates a phase-1 interruption by creating only the staging table (no
|
||||
marker). On the next rebuild ``_resolve_rebuild_recovery`` should drop
|
||||
the partial staging — the live chunks table is still authoritative.
|
||||
"""
|
||||
from haiku.rag.client.rebuild import _StagingChunkRecord
|
||||
|
||||
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
|
||||
|
||||
# Simulate a partial phase 1 (staging exists, marker absent).
|
||||
await client.store.db.create_table(
|
||||
"chunks_rebuild_staging", schema=_StagingChunkRecord
|
||||
)
|
||||
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
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
|
||||
]
|
||||
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()
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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()
|
||||
async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_path):
|
||||
"""Test embed-only rebuild skips chunks with unchanged embeddings."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue