From dc16f74b583a6f44fba7e3419c78268fb3e9fdb9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 27 Apr 2026 15:13:44 +0300 Subject: [PATCH] add picture_data column to document_items --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/store/engine.py | 23 +++++- .../haiku/rag/store/models/document_item.py | 1 + .../rag/store/repositories/document_item.py | 59 ++++++++++++- .../haiku/rag/store/upgrades/__init__.py | 1 - tests/store/test_document_items.py | 82 +++++++++++++++++++ 6 files changed, 167 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d62bbb9..e3f88ab3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- **Storage column for embedded picture bytes.** `DocumentItemRecord` gains a `picture_data: bytes | None` column (Arrow `large_binary`) to hold per-`PictureItem` image bytes addressable by `(document_id, self_ref)`. New repository accessors `get_picture_bytes` and `get_pictures_for_chunk` expose them; the existing items read paths (`get_all_items`, `get_all_items_grouped`, `get_items_in_range`, `_record_to_item`) now project an explicit lightweight column set so context expansion and the analysis-sandbox `items.jsonl` build never pull picture bytes into memory. Existing databases pick up the column via the `0.45.0` migration alongside the picture-byte backfill (see below). Foundation for upcoming vision-in-context retrieval; not yet wired into ingestion or search. + ## [0.44.0] - 2026-04-29 ### Added diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 8de4815e..454ba9fa 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -140,6 +140,27 @@ class DocumentItemRecord(LanceModel): label: str = Field(default="") text: str = Field(default="") page_numbers: str = Field(default="[]") + picture_data: bytes | None = None + + +def get_document_items_arrow_schema() -> pa.Schema: + """Generate Arrow schema for document_items with large_binary for picture_data. + + LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets + and is limited to ~2GB per column in a fragment. Many embedded picture PNGs in + one fragment can exceed that limit. `large_binary` uses 64-bit offsets and has + no practical size limit — same reasoning as `docling_document` on the + documents table. + """ + base_schema = DocumentItemRecord.to_arrow_schema() + large_binary_columns = {"picture_data"} + fields = [] + for field in base_schema: + if field.name in large_binary_columns: + fields.append(pa.field(field.name, pa.large_binary())) + else: + fields.append(field) + return pa.schema(fields) class SettingsRecord(LanceModel): @@ -456,7 +477,7 @@ class Store: self.document_items_table = await self.db.open_table("document_items") else: self.document_items_table = await self.db.create_table( - "document_items", schema=DocumentItemRecord + "document_items", schema=get_document_items_arrow_schema() ) await self.document_items_table.create_index( "document_id", config=BTree(), replace=True diff --git a/haiku_rag_slim/haiku/rag/store/models/document_item.py b/haiku_rag_slim/haiku/rag/store/models/document_item.py index ee17bf0c..7c781b41 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/models/document_item.py @@ -13,6 +13,7 @@ class DocumentItem(BaseModel): label: str = "" text: str = "" page_numbers: list[int] = [] + picture_data: bytes | None = None def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None: diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index de093a0b..64c0d0e0 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -4,6 +4,19 @@ from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.utils import escape_sql_string +# Columns returned by the lightweight read path. `picture_data` is intentionally +# excluded so context expansion and the analysis-sandbox items.jsonl build don't +# pull MB-scale image bytes into memory. Use get_picture_bytes / +# get_pictures_for_chunk to fetch picture_data explicitly. +_LIGHT_COLUMNS = [ + "document_id", + "position", + "self_ref", + "label", + "text", + "page_numbers", +] + class DocumentItemRepository: """Repository for DocumentItem operations.""" @@ -35,6 +48,7 @@ class DocumentItemRepository: label=item.label, text=item.text, page_numbers=json.dumps(item.page_numbers), + picture_data=item.picture_data, ) for item in items ] @@ -45,6 +59,7 @@ class DocumentItemRepository: safe_id = escape_sql_string(document_id) rows = await ( self.store.document_items_table.query() + .select(_LIGHT_COLUMNS) .where(f"document_id = '{safe_id}'") .to_list() ) @@ -64,7 +79,7 @@ class DocumentItemRepository: Returns: Dict mapping document_id to sorted list of DocumentItem. """ - query = self.store.document_items_table.query() + query = self.store.document_items_table.query().select(_LIGHT_COLUMNS) if document_ids is not None: safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids) query = query.where(f"document_id IN ({safe_ids})") @@ -85,6 +100,7 @@ class DocumentItemRepository: safe_id = escape_sql_string(document_id) rows = await ( self.store.document_items_table.query() + .select(_LIGHT_COLUMNS) .where( f"document_id = '{safe_id}' " f"AND position >= {start} AND position <= {end}" @@ -122,3 +138,44 @@ class DocumentItemRepository: self.store._assert_writable() safe_id = escape_sql_string(document_id) await self.store.document_items_table.delete(f"document_id = '{safe_id}'") + + async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None: + """Fetch raw picture bytes for a single picture item by self_ref.""" + safe_id = escape_sql_string(document_id) + safe_ref = escape_sql_string(self_ref) + rows = await ( + self.store.document_items_table.query() + .select(["picture_data"]) + .where(f"document_id = '{safe_id}' AND self_ref = '{safe_ref}'") + .limit(1) + .to_list() + ) + if not rows: + return None + return rows[0].get("picture_data") + + async def get_pictures_for_chunk( + self, document_id: str, refs: list[str] + ) -> dict[str, bytes]: + """Fetch picture bytes for multiple self_refs within a single document. + + Returns a mapping of self_ref → bytes, including only refs that have + non-null picture_data. Refs without bytes (or unknown refs) are omitted. + """ + if not refs: + return {} + + safe_id = escape_sql_string(document_id) + refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) + rows = await ( + self.store.document_items_table.query() + .select(["self_ref", "picture_data"]) + .where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})") + .to_list() + ) + result: dict[str, bytes] = {} + for row in rows: + data = row.get("picture_data") + if data: + result[row["self_ref"]] = data + return result diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index dd80f8f8..82ef610b 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -84,7 +84,6 @@ from haiku.rag.store.upgrades.v0_38_0 import ( from haiku.rag.store.upgrades.v0_40_0 import ( upgrade_populate_document_items as upgrade_0_40_0_document_items, ) - upgrades.append(upgrade_0_20_0_docling) upgrades.append(upgrade_0_23_1_contextualize) upgrades.append(upgrade_0_25_0_compress) diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index b495353a..68cf4a3a 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -349,3 +349,85 @@ class TestDocumentItemMigration: # No items should have been created assert await store.document_items_table.count_rows() == 0 + + +@pytest.mark.asyncio +class TestPictureDataStorage: + async def test_create_and_get_picture_bytes(self, temp_db_path): + """Round-trip picture bytes through DocumentItem and the repository.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + repo = DocumentItemRepository(rag.store) + + png_bytes = b"\x89PNG\r\n\x1a\nfake-picture-bytes" + items = [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/texts/0", + label="paragraph", + text="Some text", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/pictures/0", + label="picture", + text="", + picture_data=png_bytes, + ), + ] + await repo.create_items("doc-1", items) + + # Single-ref lookup + got = await repo.get_picture_bytes("doc-1", "#/pictures/0") + assert got == png_bytes + # Missing ref returns None + assert await repo.get_picture_bytes("doc-1", "#/pictures/999") is None + # Non-picture row has no bytes + assert await repo.get_picture_bytes("doc-1", "#/texts/0") is None + + # Batch lookup omits refs without bytes + batch = await repo.get_pictures_for_chunk( + "doc-1", ["#/pictures/0", "#/texts/0", "#/pictures/999"] + ) + assert batch == {"#/pictures/0": png_bytes} + + # Empty refs returns empty dict + assert await repo.get_pictures_for_chunk("doc-1", []) == {} + + async def test_hot_paths_exclude_picture_data(self, temp_db_path): + """Light read paths must NOT pull picture_data into memory.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + repo = DocumentItemRepository(rag.store) + + heavy = b"x" * 1024 + await repo.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/pictures/0", + label="picture", + text="", + picture_data=heavy, + ), + ], + ) + + for item in await repo.get_all_items("doc-1"): + assert item.picture_data is None + for item in await repo.get_items_in_range("doc-1", 0, 10): + assert item.picture_data is None + grouped = await repo.get_all_items_grouped(["doc-1"]) + for item in grouped["doc-1"]: + assert item.picture_data is None + + # But the picture-byte accessors still work + assert (await repo.get_picture_bytes("doc-1", "#/pictures/0")) == heavy + + async def test_fresh_db_has_picture_data_column(self, temp_db_path): + """A newly-created DB has picture_data on document_items via _init_tables.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + schema = await rag.store.document_items_table.schema() + assert "picture_data" in {f.name for f in schema}