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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-31 17:26:59 +03:00
parent 61df6b65a6
commit 35f8d721b1
No known key found for this signature in database
10 changed files with 216 additions and 7 deletions

View file

@ -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.

View file

@ -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

View file

@ -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.

View file

@ -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:

View file

@ -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 "

View file

@ -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

View file

@ -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)

View file

@ -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)

View file

@ -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

View file

@ -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(