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:
parent
61df6b65a6
commit
35f8d721b1
10 changed files with 216 additions and 7 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in a new issue