From 35f8d721b1360f64e6a51736ae391d1a99993167 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 17:26:59 +0300 Subject: [PATCH] Build the chunks FTS index on the first write instead of at table creation An FTS index built over an empty table indexes nothing and lance never catches it up on add. ensure_indexes skips FTS while the table is empty; create, replace_for_document and embed-only rebuild ensure indexes after writing, so the index always covers at least its first rows. The first write into a fresh table writes one extra chunks table version for the index build; a failed build fails the write. --- CHANGELOG.md | 3 + docs/configuration/storage.md | 2 + haiku_rag_slim/haiku/rag/client/rebuild.py | 4 +- .../haiku/rag/store/repositories/chunk.py | 5 +- haiku_rag_slim/haiku/rag/store/schema.py | 5 + tests/store/test_indexes.py | 126 +++++++++++++++++- tests/store/test_v0_75_0_migration.py | 12 ++ tests/test_chunk.py | 37 +++++ tests/test_client.py | 14 +- tests/test_rebuild.py | 15 ++- 10 files changed, 216 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4eba9af..9aba1895 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +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. - 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/docs/configuration/storage.md b/docs/configuration/storage.md index e9e9df2b..a94b2b3c 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -64,6 +64,8 @@ If that number times six exceeds available RAM, use one of: - Set `generate_page_images: false` if visual grounding through `visualize_chunk()` is not needed. This removes page rasters entirely. - Set `auto_vacuum: false` and run `haiku-rag vacuum` manually when the machine is otherwise idle, so the peak does not land alongside ingestion. +Vacuum also folds new rows into the full-text index. Search stays correct without it but scans the uncovered rows on every query. `haiku-rag doctor` reports the coverage. + This is an upstream limitation rather than a `haiku.rag` setting. Compaction bounds itself by row count instead of bytes, and LanceDB's async API exposes no batch size or fragment target to override it. Tracked at [lancedb/lancedb#2325](https://github.com/lancedb/lancedb/issues/2325). The requirement above will drop once compaction batches by bytes. ### Changing the Default Database Path diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index be8dffe5..66962315 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -21,7 +21,7 @@ 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 from haiku.rag.store.repositories.settings import SettingsRepository -from haiku.rag.store.schema import ChunkRecordBase +from haiku.rag.store.schema import ChunkRecordBase, ensure_indexes if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -504,6 +504,8 @@ async def _rebuild_embed_only( if pending_records: await session.store.chunks_table.add(pending_records) + 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 # rebuild discards harmlessly. diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index d3a63a57..061823a8 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -104,8 +104,9 @@ class ChunkRepository: chunk_record = self._to_record(entity, chunk_id) await self.store.chunks_table.add([chunk_record]) - entity.id = chunk_id + await ensure_indexes(self.store.chunks_table, "chunks") + return entity chunks = entity @@ -126,6 +127,7 @@ class ChunkRepository: chunk.id = chunk_id await self.store.chunks_table.add(chunk_records) + await ensure_indexes(self.store.chunks_table, "chunks") return chunks @@ -159,6 +161,7 @@ class ChunkRepository: .when_not_matched_by_source_delete(f"document_id = '{safe_id}'") .execute(records) ) + await ensure_indexes(self.store.chunks_table, "chunks") return chunks async def get_by_id(self, entity_id: str) -> Chunk | None: diff --git a/haiku_rag_slim/haiku/rag/store/schema.py b/haiku_rag_slim/haiku/rag/store/schema.py index e9832ad4..0aae03e1 100644 --- a/haiku_rag_slim/haiku/rag/store/schema.py +++ b/haiku_rag_slim/haiku/rag/store/schema.py @@ -174,6 +174,11 @@ async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str present = covering.get(column, set()) if declared in present: 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 + # unsorted by score, with matching rows dropped. + if isinstance(config, FTS) and not await table.count_rows(): + continue if present: logger.info( f"Adding {declared} index on {table_name}.{column}, which carries " diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index 0f9451f6..91aadb1d 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -3,17 +3,40 @@ import pytest from lancedb.index import BTree from haiku.rag.store.engine import Store -from haiku.rag.store.models import Document +from haiku.rag.store.models import Chunk, Document +from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.schema import ensure_indexes EXPECTED_INDEXED_COLUMNS = { "documents": {"id"}, "document_meta": {"id", "uri"}, - "chunks": {"content_fts", "id", "document_id"}, + "chunks": {"id", "document_id"}, "document_items": {"document_id", "position", "self_ref", "label"}, } +# content_fts joins the set once the table holds rows to index. +EXPECTED_POPULATED_CHUNK_COLUMNS = {"content_fts", "id", "document_id"} + + +async def _fts_indexed_rows(table) -> int | None: + """Rows the FTS index covers, or None when there is no FTS index.""" + for index in await table.list_indices(): + if index.index_type == "FTS": + return (await table.index_stats(index.name)).num_indexed_rows + return None + + +async def _add_chunk(store, content: str = "a chunk about gardens") -> None: + await ChunkRepository(store).create( + Chunk( + document_id="doc-1", + content=content, + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ) + async def _indexed_columns(table) -> set[str]: return {column for index in await table.list_indices() for column in index.columns} @@ -130,3 +153,102 @@ async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): schema = await store.document_items_table.schema() assert schema.field("picture_data").type == pa.large_binary() + + +@pytest.mark.asyncio +async def test_empty_chunks_table_carries_no_fts_index(temp_db_path): + """An FTS index over no rows indexes nothing and lance never catches it up, + so it waits for rows rather than being built with the table.""" + async with Store(temp_db_path, create=True) as store: + assert await _fts_indexed_rows(store.chunks_table) is None + + +@pytest.mark.asyncio +async def test_ensure_indexes_skips_fts_while_the_table_is_empty(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await ensure_indexes(store.chunks_table, "chunks") + + assert await _fts_indexed_rows(store.chunks_table) is None + assert ( + await _indexed_columns(store.chunks_table) + == (EXPECTED_INDEXED_COLUMNS["chunks"]) + ) + + +@pytest.mark.asyncio +async def test_ensure_indexes_builds_fts_over_existing_rows(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-1", + content="a chunk", + content_fts="a chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) + + await ensure_indexes(store.chunks_table, "chunks") + + assert await _fts_indexed_rows(store.chunks_table) == 1 + assert await _indexed_columns(store.chunks_table) == ( + EXPECTED_POPULATED_CHUNK_COLUMNS + ) + + +@pytest.mark.asyncio +async def test_creating_chunks_builds_a_covering_fts_index(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_replacing_chunks_builds_a_covering_fts_index(temp_db_path): + """replace_for_document inserts where nothing matches, so it can be the + first write into a fresh table.""" + async with Store(temp_db_path, create=True) as store: + await ChunkRepository(store).replace_for_document( + "doc-1", + [ + Chunk( + document_id="doc-1", + content="a chunk about gardens", + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ], + ) + + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_a_second_write_leaves_the_covering_index_in_place(temp_db_path): + """One indexed row is enough: later rows merge as a scanned tail, so the + index is not rebuilt per write.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first") + version_after_first = await store.chunks_table.version() + + await _add_chunk(store, "second") + + assert await _fts_indexed_rows(store.chunks_table) == 1 + assert await store.chunks_table.version() == version_after_first + 1 + + +@pytest.mark.asyncio +async def test_delete_all_then_write_rebuilds_a_covering_fts_index(temp_db_path): + """delete_all recreates the table empty, so the next write owns the index.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + await ChunkRepository(store).delete_all() + assert await _fts_indexed_rows(store.chunks_table) is None + + await _add_chunk(store) + + assert await _fts_indexed_rows(store.chunks_table) == 1 diff --git a/tests/store/test_v0_75_0_migration.py b/tests/store/test_v0_75_0_migration.py index 3b996d51..c8fdca91 100644 --- a/tests/store/test_v0_75_0_migration.py +++ b/tests/store/test_v0_75_0_migration.py @@ -32,6 +32,18 @@ async def _make_legacy(store: Store) -> None: @pytest.mark.asyncio async def test_adds_the_missing_indexes(temp_db_path): async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-1", + content="a chunk", + content_fts="a chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) await _make_legacy(store) await _apply_index_hot_lookup_keys(store) diff --git a/tests/test_chunk.py b/tests/test_chunk.py index ce091e5c..53d12bb1 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -503,11 +503,16 @@ async def test_fts_search_warns_when_index_covers_no_rows(temp_db_path): searching that state warns, once.""" import logging + from lancedb.index import FTS + from haiku.rag.store.repositories import chunk as chunk_module async with HaikuRAG( db_path=temp_db_path, config=get_config(), create=True ) as client: + await client.store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) await _import_one(client) with capture_logs(chunk_module.logger, logging.WARNING) as records: @@ -588,6 +593,33 @@ async def test_fts_coverage_check_failure_does_not_break_search(temp_db_path): assert results +async def test_create_assigns_the_id_when_index_maintenance_fails(temp_db_path): + """The row is committed before the index is ensured, so the chunk carries + the id it was written with even when that ensure fails.""" + from haiku.rag.store.repositories import chunk as chunk_module + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + repo = client.chunk_repository + + async def boom(*_args, **_kwargs): + raise RuntimeError("index build failed") + + chunk = Chunk( + document_id="doc-1", + content="a chunk about gardens", + embedding=[0.1] * get_config().embeddings.model.vector_dim, + order=0, + ) + with patch.object(chunk_module, "ensure_indexes", boom): + with pytest.raises(RuntimeError, match="index build failed"): + await repo.create(chunk) + + assert chunk.id is not None + assert await client.store.chunks_table.count_rows() == 1 + + async def test_fts_search_on_an_empty_table_does_not_suppress_later_warnings( temp_db_path, ): @@ -595,11 +627,16 @@ async def test_fts_search_on_an_empty_table_does_not_suppress_later_warnings( spend the once-per-repository check.""" import logging + from lancedb.index import FTS + from haiku.rag.store.repositories import chunk as chunk_module async with HaikuRAG( db_path=temp_db_path, config=get_config(), create=True ) as client: + await client.store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) await client.chunk_repository.search("gardens", search_type="fts") await _import_one(client) diff --git a/tests/test_client.py b/tests/test_client.py index ef1bdfad..99fdd429 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -895,8 +895,20 @@ async def test_client_import_documents_single_version_per_table(temp_db_path): assert all(d.id is not None for d in docs) assert len({d.id for d in docs}) == 3 + assert after["documents"] - before["documents"] == 1 + assert after["document_items"] - before["document_items"] == 1 + # Two on chunks: the batch, plus building the FTS index over the first + # rows the table has ever held. + assert after["chunks"] - before["chunks"] == 2 + + # With the index in place, a further batch is one version per table. + again = await client.import_documents( + [_import("d", "Delta document body", uri="mem://d", title="Delta")] + ) + assert [d.title for d in again] == ["Delta"] + latest = await client.store.current_table_versions() for table in ("documents", "chunks", "document_items"): - assert after[table] - before[table] == 1, table + assert latest[table] - after[table] == 1, table for doc, expected in zip(docs, ("Alpha", "Beta", "Gamma")): assert doc.id is not None diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 1cc7f985..402db0f3 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -101,8 +101,9 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging( 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. + - the staging table is dropped at the end (no leak), + - the rebuild yields every document with chunks, and + - the recreated chunks table carries an FTS index covering its rows. """ async with HaikuRAG(temp_db_path, create=True) as client: doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"]) @@ -136,6 +137,16 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging( tables = (await client.store.db.list_tables()).tables assert "chunks_rebuild_staging" not in tables + # Phase 2 recreates the chunks table, so it owns the FTS index. + 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_drops_leftover_staging_table(