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(