Merge pull request #550 from ggozad/feat/scalar-index-migration

Index every hot lookup key from one shared definition

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`. Existing databases need `haiku-rag migrate`.
+
 ### 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`.

app/backend/pyproject.toml

@@ -8,7 +8,7 @@ dependencies = [
     "uvicorn[standard]>=0.40.0",
     "pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0",
     "python-dotenv>=1.2.1",
-    "haiku.rag-slim>=0.74.0",
+    "haiku.rag-slim>=0.75.0",
     "logfire[pydantic-ai]>=3.17.0",
 ]
 

evaluations/pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag-evals"
 description = "Benchmarking and evaluation scripts for haiku.rag"
-version = "0.74.0"
+version = "0.75.0"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 requires-python = ">=3.12"

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,59 @@ 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 each table carries."""
+    match table_name:
+        case "documents":
+            return [("id", BTree())]
+        case "document_meta":
+            return [("id", BTree()), ("uri", BTree())]
+        case "chunks":
+            return [
+                # Positions and stop words are required for phrase queries.
+                ("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) -> list[str]:
+    """Create any declared index missing from a column. Returns the columns indexed.
+
+    Matches on index type, not column coverage, so a BTree does not satisfy a
+    declared Bitmap. Never drops or converts an index it did not declare.
+    Re-creating is not free: `create_index(replace=True)` rebuilds.
+    """
+    covering: dict[str, set[str]] = {}
+    for index in await table.list_indices():
+        for column in index.columns:
+            covering.setdefault(column, set()).add(index.index_type)
+
+    applied: list[str] = []
+    for column, config in index_specs(table_name):
+        declared = type(config).__name__
+        present = covering.get(column, set())
+        if declared in present:
+            continue
+        if present:
+            logger.info(
+                f"Adding {declared} index on {table_name}.{column}, which carries "
+                f"{', '.join(sorted(present))}"
+            )
+        await table.create_index(column, config=config, replace=True)
+        applied.append(column)
+    return applied
+
+
 class SettingsRecord(LanceModel):
     id: str = Field(default="settings")
     settings: str = Field(default="{}")
@@ -697,22 +750,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
+        # blob-bearing documents row).
-        # hot look-up keys (get_by_id, get_by_uri).
         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(
+            await ensure_indexes(self.document_meta_table, "document_meta")
-                "id", config=BTree(), replace=True
-            )
-            await self.document_meta_table.create_index(
-                "uri", config=BTree(), replace=True
-            )
 
         # Create or open chunks table
         if "chunks" in existing_tables:
@@ -721,12 +769,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 ensure_indexes(self.chunks_table, "chunks")
-            await self.chunks_table.create_index(
-                "content_fts",
-                config=FTS(with_position=True, remove_stop_words=False),
-                replace=True,
-            )
 
         # Create or open document_items table
         if "document_items" in existing_tables:
@@ -735,15 +778,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(
+            await ensure_indexes(self.document_items_table, "document_items")
-                "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
-            )
 
         # Create or open settings table
         if "settings" in existing_tables:
@@ -872,13 +907,7 @@ class Store:
         self.chunks_table = await self.db.create_table(
             "chunks", schema=self.ChunkRecord
         )
-
+        await ensure_indexes(self.chunks_table, "chunks")
-        # 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,
-        )
 
     def close(self):
         """Close the database connection."""

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 ensure_indexes(self.store.chunks_table, "chunks")
-        await self.store.chunks_table.create_index(
-            "content_fts",
-            config=FTS(with_position=True, remove_stop_words=False),
-            replace=True,
-        )
 
     async def delete_by_document_id(self, document_id: str) -> bool:
         """Delete all chunks for a document."""

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
+            "document_items", schema=get_document_items_arrow_schema()
-        )
-        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
         )
+        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(
+            await ensure_indexes(self.store.document_meta_table, "document_meta")
-                "id", config=BTree(), replace=True
-            )
-            await self.store.document_meta_table.create_index(
-                "uri", config=BTree(), replace=True
-            )

haiku_rag_slim/haiku/rag/store/upgrades/__init__.py

@@ -99,6 +99,9 @@ from haiku.rag.store.upgrades.v0_58_0 import (
 from haiku.rag.store.upgrades.v0_64_0 import (
     upgrade_rename_document_meta_id as upgrade_0_64_0_rename_document_meta_id,
 )
+from haiku.rag.store.upgrades.v0_75_0 import (
+    upgrade_index_hot_lookup_keys as upgrade_0_75_0_index_hot_lookup_keys,
+)
 
 upgrades.append(upgrade_0_20_0_docling)
 upgrades.append(upgrade_0_23_1_contextualize)
@@ -110,3 +113,4 @@ upgrades.append(upgrade_0_48_0_heading_hierarchy)
 upgrades.append(upgrade_0_50_0_canonical_metadata_keys)
 upgrades.append(upgrade_0_58_0_split_document_meta)
 upgrades.append(upgrade_0_64_0_rename_document_meta_id)
+upgrades.append(upgrade_0_75_0_index_hot_lookup_keys)

haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py

@@ -0,0 +1,24 @@
+import logging
+
+from haiku.rag.store.engine import Store, ensure_indexes
+from haiku.rag.store.upgrades import Upgrade
+
+logger = logging.getLogger(__name__)
+
+
+async def _apply_index_hot_lookup_keys(store: Store) -> None:
+    """Add the declared indexes to a database created before 0.75.0.
+
+    Rewrites no rows. Each index build reads the column it indexes.
+    """
+    for table_name, table in store._tables().items():
+        applied = await ensure_indexes(table, table_name)
+        if applied:
+            logger.info(f"Indexed {table_name}: {', '.join(sorted(applied))}")
+
+
+upgrade_index_hot_lookup_keys = Upgrade(
+    version="0.75.0",
+    apply=_apply_index_hot_lookup_keys,
+    description="Index documents.id, chunks.id, chunks.document_id and document_items.label",
+)

haiku_rag_slim/pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag-slim"
 description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
-version = "0.74.0"
+version = "0.75.0"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag"
 description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
-version = "0.74.0"
+version = "0.75.0"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }
@@ -30,7 +30,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.74.0",
+    "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.75.0",
 ]
 
 [project.scripts]
@@ -38,9 +38,9 @@ haiku-rag = "haiku.rag.cli:cli"
 
 [project.optional-dependencies]
 tui = ["textual>=8.2.4"]
-s3 = ["haiku.rag-slim[s3]==0.74.0"]
+s3 = ["haiku.rag-slim[s3]==0.75.0"]
-cross-encoder = ["haiku.rag-slim[cross-encoder]==0.74.0"]
+cross-encoder = ["haiku.rag-slim[cross-encoder]==0.75.0"]
-ingester = ["haiku.rag-slim[ingester]==0.74.0"]
+ingester = ["haiku.rag-slim[ingester]==0.75.0"]
 
 [build-system]
 requires = ["hatchling"]

tests/store/test_indexes.py

@@ -0,0 +1,131 @@
+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
+
+
+async def _covering(table, column: str) -> list[tuple[str, str]]:
+    """Every index over `column`, as (name, index_type)."""
+    return [
+        (index.name, index.index_type)
+        for index in await table.list_indices()
+        if column in index.columns
+    ]
+
+
+@pytest.mark.asyncio
+async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path):
+    """A new database carries the full index set."""
+    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):
+    """A second pass must not rebuild: replace=True writes a new version."""
+    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 wrong-typed index does not satisfy the declared one."""
+    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_ensure_indexes_adds_the_declared_type_beside_a_custom_index(
+    temp_db_path,
+):
+    """A custom-named index neither satisfies the check nor is destroyed."""
+    async with Store(temp_db_path, create=True) as store:
+        table = store.document_items_table
+        await table.drop_index("label_idx")
+        await table.create_index("label", config=BTree(), name="operator_label")
+
+        await ensure_indexes(table, "document_items")
+
+        covering = dict(await _covering(table, "label"))
+        assert covering["operator_label"] == "BTree"
+        assert "Bitmap" in covering.values()
+
+
+@pytest.mark.asyncio
+async def test_ensure_indexes_keeps_an_operator_index_on_a_declared_column(
+    temp_db_path,
+):
+    """An index we did not declare survives, even on a declared column."""
+    async with Store(temp_db_path, create=True) as store:
+        table = store.document_items_table
+        await table.create_index("label", config=BTree(), name="operator_label")
+
+        await ensure_indexes(table, "document_items")
+
+        covering = dict(await _covering(table, "label"))
+        assert covering == {"label_idx": "Bitmap", "operator_label": "BTree"}
+
+
+@pytest.mark.asyncio
+async def test_delete_all_restores_the_full_index_set(temp_db_path):
+    """Recreated tables come back with the full index set."""
+    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):
+    """picture_data must survive delete_all as large_binary, not binary."""
+    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()

tests/store/test_v0_75_0_migration.py

@@ -0,0 +1,105 @@
+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.repositories.document import DocumentRepository
+from haiku.rag.store.upgrades.v0_75_0 import _apply_index_hot_lookup_keys
+
+# The indexes a pre-0.75.0 database lacked.
+LEGACY_DROPPED = {
+    "documents": ["id_idx"],
+    "chunks": ["id_idx", "document_id_idx"],
+    "document_items": ["label_idx"],
+}
+
+
+async def _indexed(table) -> dict[str, str]:
+    return {
+        column: index.index_type
+        for index in await table.list_indices()
+        for column in index.columns
+    }
+
+
+async def _make_legacy(store: Store) -> None:
+    for table_name, indexes in LEGACY_DROPPED.items():
+        table = store._tables()[table_name]
+        for index in indexes:
+            await table.drop_index(index)
+
+
+@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 _make_legacy(store)
+
+        await _apply_index_hot_lookup_keys(store)
+
+        assert await _indexed(store.documents_table) == {"id": "BTree"}
+        assert await _indexed(store.chunks_table) == {
+            "content_fts": "FTS",
+            "id": "BTree",
+            "document_id": "BTree",
+        }
+        assert await _indexed(store.document_items_table) == {
+            "document_id": "BTree",
+            "position": "BTree",
+            "self_ref": "BTree",
+            "label": "Bitmap",
+        }
+
+
+@pytest.mark.asyncio
+async def test_keeps_every_row(temp_db_path):
+    """Indexing must not touch data."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        repo = DocumentRepository(store)
+        await repo.create(Document(content="Kept", uri="test://kept"))
+        await _make_legacy(store)
+
+        await _apply_index_hot_lookup_keys(store)
+
+        docs = await repo.list_all(include_content=True)
+        assert [d.content for d in docs] == ["Kept"]
+
+
+@pytest.mark.asyncio
+async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path):
+    """A second run must not re-index: replace=True rebuilds."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        await _apply_index_hot_lookup_keys(store)
+        versions = {name: await t.version() for name, t in store._tables().items()}
+
+        await _apply_index_hot_lookup_keys(store)
+
+        assert {name: await t.version() for name, t in store._tables().items()} == (
+            versions
+        )
+
+
+@pytest.mark.asyncio
+async def test_replaces_a_wrong_typed_legacy_index(temp_db_path):
+    """A BTree on `label` does not satisfy the declared Bitmap."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        await store.document_items_table.create_index(
+            "label", config=BTree(), replace=True
+        )
+
+        await _apply_index_hot_lookup_keys(store)
+
+        indexed = await _indexed(store.document_items_table)
+        assert indexed["label"] == "Bitmap"
+
+
+@pytest.mark.asyncio
+async def test_leaves_undeclared_indexes_alone(temp_db_path):
+    """Indexes haiku.rag never declared are not dropped."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        await store.document_meta_table.create_index(
+            "title", config=BTree(), replace=True
+        )
+
+        await _apply_index_hot_lookup_keys(store)
+
+        assert "title" in await _indexed(store.document_meta_table)

tests/test_client.py

@@ -993,10 +993,15 @@ async def test_client_import_documents_mixed_embeddings(temp_db_path):
 async def test_client_update_document_replaces_rows_with_bounded_versions(
     temp_db_path,
 ):
-    """Updating one document should replace stale rows with bounded versions."""
+    """Updating one document should replace stale rows with bounded versions.
-    dim = Config.embeddings.model.vector_dim
 
-    async with HaikuRAG(temp_db_path, create=True) as client:
+    auto_vacuum is off: its writes would land inside the measured window.
+    """
+    dim = Config.embeddings.model.vector_dim
+    config = Config.model_copy(deep=True)
+    config.storage.auto_vacuum = False
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
         created = await client.import_document(
             _docling_doc("original", "Original body"),
             [Chunk(content="Original body", embedding=[0.1] * dim, order=0)],

uv.lock

@@ -1577,7 +1577,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.74.0"
+version = "0.75.0"
 source = { editable = "." }
 dependencies = [
     { name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "tui", "voyageai", "zeroentropy"] },
@@ -1644,7 +1644,7 @@ dev = [
 
 [[package]]
 name = "haiku-rag-evals"
-version = "0.74.0"
+version = "0.75.0"
 source = { editable = "evaluations" }
 dependencies = [
     { name = "datasets" },
@@ -1667,7 +1667,7 @@ requires-dist = [
 
 [[package]]
 name = "haiku-rag-slim"
-version = "0.74.0"
+version = "0.75.0"
 source = { editable = "haiku_rag_slim" }
 dependencies = [
     { name = "docling-core" },