af6a6b0bbe
`_populate_image_data` ran its stages once per result document, so a result set spanning N documents cost 4N `document_items` queries. Measured on a remote object-store corpus, a limit=5 search with expansion was 18 queries, 16 of them against `document_items`. The stages now run once each across every document, and flat in document count: two queries for the dependent caption-to-picture mapping when results ranked on a caption, one for the picture bytes. Two queries for a picture-ref result set, three at most. Picture text comes back with the bytes rather than from a second query, since it is on the same rows. Predicates are per document, `(document_id = 'a' AND self_ref IN (…)) OR (…)`, rather than `self_ref IN (union)`. self_ref and position values repeat across documents, so a union predicate would return other documents' rows: for picture_data that fetches blobs nobody asked for, and it can hand one document another document's picture.
116 lines
3.7 KiB
Python
116 lines
3.7 KiB
Python
import lancedb
|
|
import pytest
|
|
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.models import DocumentItem
|
|
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
|
|
|
|
|
async def _seed(repo: DocumentItemRepository, document_id: str) -> None:
|
|
"""A picture at position 0 with its caption at position 1, the docling
|
|
layout. Every document uses the same self_refs."""
|
|
await repo.create_items(
|
|
document_id,
|
|
[
|
|
DocumentItem(
|
|
document_id=document_id,
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text=f"caption text {document_id}",
|
|
picture_data=f"bytes-{document_id}".encode(),
|
|
),
|
|
DocumentItem(
|
|
document_id=document_id,
|
|
position=1,
|
|
self_ref="#/texts/1",
|
|
label="caption",
|
|
text=f"figure 1 {document_id}",
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def item_queries(monkeypatch):
|
|
tally = {"n": 0}
|
|
query = lancedb.AsyncTable.query
|
|
|
|
def counted(self):
|
|
if self.name == "document_items":
|
|
tally["n"] += 1
|
|
return query(self)
|
|
|
|
monkeypatch.setattr(lancedb.AsyncTable, "query", counted)
|
|
return tally
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pictures_grouped_keeps_documents_apart(temp_db_path, item_queries):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentItemRepository(store)
|
|
await _seed(repo, "doc-a")
|
|
await _seed(repo, "doc-b")
|
|
|
|
item_queries["n"] = 0
|
|
blobs, texts = await repo.get_pictures_grouped(
|
|
{"doc-a": ["#/pictures/0"], "doc-b": ["#/pictures/0"]}
|
|
)
|
|
|
|
assert item_queries["n"] == 1
|
|
assert blobs == {
|
|
"doc-a": {"#/pictures/0": b"bytes-doc-a"},
|
|
"doc-b": {"#/pictures/0": b"bytes-doc-b"},
|
|
}
|
|
assert texts == {
|
|
"doc-a": {"#/pictures/0": "caption text doc-a"},
|
|
"doc-b": {"#/pictures/0": "caption text doc-b"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pictures_grouped_fetches_only_requested_documents(temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentItemRepository(store)
|
|
await _seed(repo, "doc-a")
|
|
await _seed(repo, "doc-b")
|
|
|
|
blobs, _ = await repo.get_pictures_grouped({"doc-a": ["#/pictures/0"]})
|
|
|
|
assert blobs == {"doc-a": {"#/pictures/0": b"bytes-doc-a"}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_caption_picture_refs_grouped_uses_two_queries(
|
|
temp_db_path, item_queries
|
|
):
|
|
"""The stages are dependent: the caption's position is what finds its
|
|
picture, so this is two queries however many documents are asked for."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentItemRepository(store)
|
|
for document_id in ("doc-a", "doc-b", "doc-c"):
|
|
await _seed(repo, document_id)
|
|
|
|
item_queries["n"] = 0
|
|
got = await repo.get_caption_picture_refs_grouped(
|
|
{did: ["#/texts/1"] for did in ("doc-a", "doc-b", "doc-c")}
|
|
)
|
|
|
|
assert item_queries["n"] == 2
|
|
for document_id in ("doc-a", "doc-b", "doc-c"):
|
|
assert got[document_id] == {"#/texts/1": "#/pictures/0"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_grouped_calls_with_nothing_asked_for_do_not_query(
|
|
temp_db_path, item_queries
|
|
):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
repo = DocumentItemRepository(store)
|
|
|
|
item_queries["n"] = 0
|
|
assert await repo.get_pictures_grouped({}) == ({}, {})
|
|
assert await repo.get_pictures_grouped({"doc-a": []}) == ({}, {})
|
|
assert await repo.get_caption_picture_refs_grouped({}) == {}
|
|
|
|
assert item_queries["n"] == 0
|