danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Index every hot lookup key from one shared definition

Browse source 
`_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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-17 14:54:01 +03:00
parent 980b1985ab
commit c184a25d68
No known key found for this signature in database
5 changed files with 179 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

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

98
haiku_rag_slim/haiku/rag/store/engine.py
View file

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

9
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

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

25
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

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

98
tests/store/test_indexes.py Normal file
View file

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