From c184a25d6875b9e2d9cefb9ece427f93f49233fd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Aug 2026 14:54:01 +0300 Subject: [PATCH] Index every hot lookup key from one shared definition `_init_tables` left `chunks.id`, `chunks.document_id` and `documents.id` unindexed, so those lookups scanned the column. On object storage that is network I/O per query, on paths that run per document: citation lookup, delete-by-document, the re-ingest merge, and every dedup probe. Declare the index set per table in `index_specs()` and apply it through `ensure_indexes()`, which skips a column only when it is already indexed with the declared type. Both halves of that are load-bearing. Skipping is required because `create_index(replace=True)` rebuilds an identical index, writing a new index and a new table version and orphaning the old files until the next vacuum. Comparing the type is required because column coverage alone would let a wrong-typed index stand, and a BTree on `label` silently loses the low-cardinality equality lookup the Bitmap is there for. Columns not declared for a table are left alone, so an externally created index such as a vector index on `chunks` survives. `_init_tables`, `recreate_embeddings_table`, `ChunkRepository.delete_all` and `DocumentRepository.delete_all` now all route through it instead of repeating their own subsets. Also recreate `document_items` from `get_document_items_arrow_schema()` in `DocumentRepository.delete_all`, which was using the LanceModel and so returned `picture_data` as 32-bit `binary`. Existing databases are unchanged; the migration follows separately. --- CHANGELOG.md | 6 ++ haiku_rag_slim/haiku/rag/store/engine.py | 98 +++++++++++++------ .../haiku/rag/store/repositories/chunk.py | 9 +- .../haiku/rag/store/repositories/document.py | 25 ++--- tests/store/test_indexes.py | 98 +++++++++++++++++++ 5 files changed, 179 insertions(+), 57 deletions(-) create mode 100644 tests/store/test_indexes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bcc7be9..2b4aa52a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,10 +6,16 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. +- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. + ### Changed - `import_documents` embeds chunks across the whole batch in one pass instead of per document. +### Fixed + +- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`. + ### Removed - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`. diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 50d4a73a..3ef990af 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -12,7 +12,7 @@ from uuid import uuid4 import lancedb import pyarrow as pa -from lancedb.index import FTS, BTree, IvfPq +from lancedb.index import FTS, Bitmap, BTree, IvfPq from lancedb.pydantic import LanceModel, Vector from packaging.version import parse from pydantic import BaseModel, Field @@ -176,6 +176,66 @@ def get_document_items_arrow_schema() -> pa.Schema: return pa.schema(fields) +def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: + """The index set a haiku.rag table is expected to carry. + + Single source of truth for every path that creates a table: initialization, + migration, and the drop-and-recreate paths in the repositories. + + `label` gets a Bitmap rather than a BTree because it holds around ten + distinct values across every item of every document, and Bitmap is the + low-cardinality equality case. The FTS options are load-bearing: + `with_position` enables phrase queries and keeping stop words lets them match. + """ + match table_name: + case "documents": + return [("id", BTree())] + case "document_meta": + return [("id", BTree()), ("uri", BTree())] + case "chunks": + return [ + ("content_fts", FTS(with_position=True, remove_stop_words=False)), + ("id", BTree()), + ("document_id", BTree()), + ] + case "document_items": + return [ + ("document_id", BTree()), + ("position", BTree()), + ("self_ref", BTree()), + ("label", Bitmap()), + ] + case _: + return [] + + +async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> None: + """Create the table's declared indexes, skipping those already correct. + + Correct means the column is indexed *and* carries the declared index type. + Matching on the column alone would let a wrong-typed index stand: a BTree on + `label` covers the column while losing the low-cardinality equality lookup a + Bitmap gives, and no amount of column coverage reveals that. + + Skipping matters as much as creating: `create_index(replace=True)` rebuilds + an identical index, writing a fresh index and a new table version rather + than no-oping, and leaves the previous index behind until the next vacuum. + On a large table over object storage that is a full column sort per pass. + + Columns not declared for the table are left untouched, so an externally + added index (a vector index on `chunks`, say) survives. + """ + indexed = { + column: index.index_type + for index in await table.list_indices() + for column in index.columns + } + for column, config in index_specs(table_name): + if indexed.get(column) == type(config).__name__: + continue + await table.create_index(column, config=config, replace=True) + + class SettingsRecord(LanceModel): id: str = Field(default="settings") settings: str = Field(default="{}") @@ -697,22 +757,17 @@ class Store: self.documents_table = await self.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.documents_table, "documents") # Create or open document_meta table (mutable attributes kept out of the - # blob-bearing documents row). Indexed by document_id and uri — both are - # hot look-up keys (get_by_id, get_by_uri). + # blob-bearing documents row). if "document_meta" in existing_tables: self.document_meta_table = await self.db.open_table("document_meta") else: self.document_meta_table = await self.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.document_meta_table, "document_meta") # Create or open chunks table if "chunks" in existing_tables: @@ -721,12 +776,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") # Create or open document_items table if "document_items" in existing_tables: @@ -735,15 +785,7 @@ class Store: self.document_items_table = await self.db.create_table( "document_items", schema=get_document_items_arrow_schema() ) - await self.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "self_ref", config=BTree(), replace=True - ) + await ensure_indexes(self.document_items_table, "document_items") # Create or open settings table if "settings" in existing_tables: @@ -872,13 +914,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") def close(self): """Close the database connection.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 09cc3498..ecfe9f49 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from lancedb.index import FTS from lancedb.rerankers import RRFReranker -from haiku.rag.store.engine import Store, query_to_pydantic +from haiku.rag.store.engine import Store, ensure_indexes, query_to_pydantic from haiku.rag.store.models.chunk import Chunk, SearchType from haiku.rag.utils import escape_sql_string @@ -187,12 +187,7 @@ class ChunkRepository: self.store.chunks_table = await self.store.db.create_table( "chunks", schema=self.store.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.store.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.store.chunks_table, "chunks") async def delete_by_document_id(self, document_id: str) -> bool: """Delete all chunks for a document.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 9bfd9b57..be242be3 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -3,12 +3,12 @@ from datetime import datetime from typing import overload from uuid import uuid4 -from lancedb.index import BTree - from haiku.rag.store.engine import ( DocumentMetaRecord, DocumentRecord, Store, + ensure_indexes, + get_document_items_arrow_schema, get_documents_arrow_schema, query_to_pydantic, ) @@ -407,23 +407,14 @@ class DocumentRepository: async def delete_all(self) -> None: """Delete all documents from the database.""" self.store._assert_writable() - from haiku.rag.store.engine import DocumentItemRecord # Delete all chunks and items first await self.chunk_repository.delete_all() await self.store.db.drop_table("document_items") self.store.document_items_table = await self.store.db.create_table( - "document_items", schema=DocumentItemRecord - ) - await self.store.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "self_ref", config=BTree(), replace=True + "document_items", schema=get_document_items_arrow_schema() ) + await ensure_indexes(self.store.document_items_table, "document_items") # Get count before deletion count = len( @@ -437,13 +428,9 @@ class DocumentRepository: self.store.documents_table = await self.store.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.store.documents_table, "documents") await self.store.db.drop_table("document_meta") self.store.document_meta_table = await self.store.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.store.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.store.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.store.document_meta_table, "document_meta") diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py new file mode 100644 index 00000000..a0f011c4 --- /dev/null +++ b/tests/store/test_indexes.py @@ -0,0 +1,98 @@ +import pyarrow as pa +import pytest +from lancedb.index import BTree + +from haiku.rag.store.engine import Store, ensure_indexes +from haiku.rag.store.models import Document +from haiku.rag.store.repositories.document import DocumentRepository + +EXPECTED_INDEXED_COLUMNS = { + "documents": {"id"}, + "document_meta": {"id", "uri"}, + "chunks": {"content_fts", "id", "document_id"}, + "document_items": {"document_id", "position", "self_ref", "label"}, +} + + +async def _indexed_columns(table) -> set[str]: + return {column for index in await table.list_indices() for column in index.columns} + + +async def _index_type(table, column: str) -> str | None: + for index in await table.list_indices(): + if column in index.columns: + return index.index_type + return None + + +@pytest.mark.asyncio +async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path): + """A new database carries the full index set, not a subset.""" + async with Store(temp_db_path, create=True) as store: + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path): + """`create_index(replace=True)` rebuilds an identical index and writes a new + table version, so a second pass must skip rather than replace.""" + async with Store(temp_db_path, create=True) as store: + table = store.chunks_table + version_before = await table.version() + + await ensure_indexes(table, "chunks") + + assert await table.version() == version_before + assert await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["chunks"] + + +@pytest.mark.asyncio +async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): + """A column indexed with the wrong type must be re-indexed. `label` is the + live case: a BTree over ~ten distinct values loses the low-cardinality + equality lookup a Bitmap gives, and column coverage alone cannot see it. + """ + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.create_index("label", config=BTree(), replace=True) + assert await _index_type(table, "label") == "BTree" + + await ensure_indexes(table, "document_items") + + assert await _index_type(table, "label") == "Bitmap" + assert ( + await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["document_items"] + ) + + +@pytest.mark.asyncio +async def test_delete_all_restores_the_full_index_set(temp_db_path): + """delete_all drops and recreates tables; the recreated tables must come + back with the same indexes a fresh database gets.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): + """document_items must be recreated from the Arrow schema, which declares + picture_data as large_binary. The 32-bit `binary` type overflows its offsets + once a fragment holds enough embedded pictures. + """ + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + schema = await store.document_items_table.schema() + assert schema.field("picture_data").type == pa.large_binary()