diff --git a/CHANGELOG.md b/CHANGELOG.md index 9aba1895..68fcc72f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,9 +16,9 @@ ### Fixed -- FTS and hybrid search on a database whose FTS index covers no rows. The index - is now built on the first chunk write; `haiku-rag vacuum` repairs existing - databases. +- FTS and hybrid search on a database whose FTS index covers no rows. Chunk + writes now build the index and rebuild it if it covers none; `haiku-rag + vacuum` also repairs it. - Migration to 0.38.0 no longer fails with `UnicodeDecodeError` on a `docling_document` blob written as zstd. - Vector search sets `search.vector_index_metric` on every query. diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 66962315..4de95b14 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -499,12 +499,12 @@ async def _rebuild_embed_only( if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records: await session.store.chunks_table.add(pending_records) + await ensure_indexes(session.store.chunks_table, "chunks") pending_records = [] if pending_records: await session.store.chunks_table.add(pending_records) - - await ensure_indexes(session.store.chunks_table, "chunks") + await ensure_indexes(session.store.chunks_table, "chunks") # Phase 2 finished. Drop the recovery state — marker first so a crash # between the two drops leaves only staging behind, which the next diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 3b004395..570eceba 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -271,7 +271,7 @@ async def _check_fts_coverage(store: Store) -> CheckResult: return CheckResult( name="fts_index_coverage", severity=Severity.OK, - message="Full-text search indexes cover their tables.", + message="Full-text search indexes are present and cover rows.", ) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 061823a8..3cac1864 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -230,6 +230,7 @@ class ChunkRepository: return False await self.store.chunks_table.delete(f"document_id = '{document_id}'") + await ensure_indexes(self.store.chunks_table, "chunks") return True async def search( diff --git a/haiku_rag_slim/haiku/rag/store/schema.py b/haiku_rag_slim/haiku/rag/store/schema.py index 0aae03e1..a59935ae 100644 --- a/haiku_rag_slim/haiku/rag/store/schema.py +++ b/haiku_rag_slim/haiku/rag/store/schema.py @@ -161,18 +161,32 @@ async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str Matches on index type, not column coverage, so a BTree does not satisfy a declared Bitmap. Never drops or converts an index it did not declare. - Re-creating is not free: `create_index(replace=True)` rebuilds. + Re-creating is not free: `create_index(replace=True)` rebuilds. A declared + FTS index covering no rows of a populated table is rebuilt in place: it + serves the same broken scan path as a missing one. """ covering: dict[str, set[str]] = {} + fts_names: dict[str, str] = {} for index in await table.list_indices(): for column in index.columns: covering.setdefault(column, set()).add(index.index_type) + if index.index_type == "FTS": + fts_names.setdefault(column, index.name) applied: list[str] = [] for column, config in index_specs(table_name): declared = type(config).__name__ present = covering.get(column, set()) if declared in present: + if isinstance(config, FTS): + stats = await table.index_stats(fts_names[column]) + if ( + stats is None or stats.num_indexed_rows == 0 + ) and await table.count_rows(): + await table.create_index( + column, config=config, replace=True, name=fts_names[column] + ) + applied.append(column) continue # lance indexes nothing when the table is empty and never catches an # FTS index up on add: the scan path it then serves returns results diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index 91aadb1d..c059e78b 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -1,6 +1,6 @@ import pyarrow as pa import pytest -from lancedb.index import BTree +from lancedb.index import FTS, BTree from haiku.rag.store.engine import Store from haiku.rag.store.models import Chunk, Document @@ -27,10 +27,12 @@ async def _fts_indexed_rows(table) -> int | None: return None -async def _add_chunk(store, content: str = "a chunk about gardens") -> None: +async def _add_chunk( + store, content: str = "a chunk about gardens", document_id: str = "doc-1" +) -> None: await ChunkRepository(store).create( Chunk( - document_id="doc-1", + document_id=document_id, content=content, embedding=[0.1] * store.embedder.vector_dim, order=0, @@ -252,3 +254,89 @@ async def test_delete_all_then_write_rebuilds_a_covering_fts_index(temp_db_path) await _add_chunk(store) assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_deleting_the_indexed_rows_rebuilds_the_fts_index(temp_db_path): + """Deleting every row the index covers, while unindexed rows remain, + reaches the zero-coverage scan path; the delete repairs it.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first", document_id="doc-a") + await _add_chunk(store, "second", document_id="doc-b") + assert await _fts_indexed_rows(store.chunks_table) == 1 + + await ChunkRepository(store).delete_by_document_id("doc-a") + + assert await store.chunks_table.count_rows() == 1 + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_replacing_the_indexed_rows_rebuilds_the_fts_index(temp_db_path): + """Replacement rewrites rows, and rewritten rows are unindexed.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first", document_id="doc-a") + await _add_chunk(store, "second", document_id="doc-b") + + await ChunkRepository(store).replace_for_document( + "doc-a", + [ + Chunk( + document_id="doc-a", + content="rewritten", + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ], + ) + + assert await _fts_indexed_rows(store.chunks_table) == 2 + + +@pytest.mark.asyncio +async def test_a_write_repairs_a_legacy_index_that_covers_no_rows(temp_db_path): + """A database whose FTS index predates its rows is repaired by the first + write that runs index maintenance.""" + async with Store(temp_db_path, create=True) as store: + await store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) + await store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-a", + content="a legacy chunk", + content_fts="a legacy chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) + assert await _fts_indexed_rows(store.chunks_table) == 0 + + await _add_chunk(store, "second", document_id="doc-b") + + assert await _fts_indexed_rows(store.chunks_table) == 2 + + +@pytest.mark.asyncio +async def test_unavailable_index_stats_repair_matches_doctor(temp_db_path, monkeypatch): + """index_stats may return None; doctor treats that as uncovered, so the + write-path repair does too.""" + from lancedb.table import AsyncTable + + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + + original = AsyncTable.index_stats + + async def no_stats(self, name): + if self.name == "chunks": + return None + return await original(self, name) + + monkeypatch.setattr(AsyncTable, "index_stats", no_stats) + applied = await ensure_indexes(store.chunks_table, "chunks") + + assert applied == ["content_fts"] diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 53d12bb1..be15b96d 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -498,6 +498,22 @@ async def _import_one(client) -> None: ) +async def _add_legacy_row(client) -> None: + """A row written without index maintenance, as older releases wrote them.""" + await client.store.chunks_table.add( + [ + client.store.ChunkRecord( + document_id="doc-legacy", + content="a document about gardens", + content_fts="a document about gardens", + metadata="{}", + order=0, + vector=[0.1] * get_config().embeddings.model.vector_dim, + ) + ] + ) + + async def test_fts_search_warns_when_index_covers_no_rows(temp_db_path): """A database whose FTS index predates its rows covers none of them; searching that state warns, once.""" @@ -513,7 +529,7 @@ async def test_fts_search_warns_when_index_covers_no_rows(temp_db_path): await client.store.chunks_table.create_index( "content_fts", config=FTS(with_position=True, remove_stop_words=False) ) - await _import_one(client) + await _add_legacy_row(client) with capture_logs(chunk_module.logger, logging.WARNING) as records: await client.chunk_repository.search("gardens", search_type="fts") @@ -638,7 +654,7 @@ async def test_fts_search_on_an_empty_table_does_not_suppress_later_warnings( "content_fts", config=FTS(with_position=True, remove_stop_words=False) ) await client.chunk_repository.search("gardens", search_type="fts") - await _import_one(client) + await _add_legacy_row(client) with capture_logs(chunk_module.logger, logging.WARNING) as records: await client.chunk_repository.search("gardens", search_type="fts") diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 402db0f3..7a8d55ef 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -341,6 +341,55 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state( assert "chunks_rebuild_marker" not in tables +async def test_rebuild_embed_only_covers_the_index_from_the_first_flush( + temp_db_path, monkeypatch +): + """A rebuild abandoned after a flushed batch leaves the FTS index + covering the flushed rows.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + from haiku.rag import embeddings as embeddings_module + from haiku.rag.client import rebuild as rebuild_module + from haiku.rag.config import get_config + from haiku.rag.store.models.chunk import Chunk + + monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 1) + + async def keep_embeddings(chunks, embedder, config=None): + for chunk in chunks: + if chunk.embedding is None: + chunk.embedding = [0.1] * embedder.vector_dim + return chunks + + monkeypatch.setattr(embeddings_module, "embed_chunks", keep_embeddings) + + async with HaikuRAG(temp_db_path, create=True) as client: + dim = get_config().embeddings.model.vector_dim + for name in ("one", "two"): + doc = DoclingDocument(name=name) + doc.add_text(label=DocItemLabel.TEXT, text=f"document {name}") + await client.import_document( + doc, + [Chunk(content=f"document {name}", embedding=[0.1] * dim, order=0)], + uri=f"test://{name}", + ) + + rebuild = client.rebuild_database(mode=RebuildMode.EMBED_ONLY) + await anext(rebuild) + await anext(rebuild) + await rebuild.aclose() + + fts = [ + index + for index in await client.store.chunks_table.list_indices() + if index.index_type == "FTS" + ] + assert len(fts) == 1 + stats = await client.store.chunks_table.index_stats(fts[0].name) + assert stats.num_indexed_rows > 0 + + @pytest.mark.vcr() async def test_rebuild_embed_only_skips_unchanged( qa_corpus: list[dict[str, str]], temp_db_path