add picture_data column to document_items
This commit is contained in:
parent
3391335e3a
commit
dc16f74b58
6 changed files with 167 additions and 3 deletions
|
|
@ -1,6 +1,10 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
## [0.44.0] - 2026-04-29
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
|
||||||
|
|
@ -140,6 +140,27 @@ class DocumentItemRecord(LanceModel):
|
||||||
label: str = Field(default="")
|
label: str = Field(default="")
|
||||||
text: str = Field(default="")
|
text: str = Field(default="")
|
||||||
page_numbers: 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):
|
class SettingsRecord(LanceModel):
|
||||||
|
|
@ -456,7 +477,7 @@ class Store:
|
||||||
self.document_items_table = await self.db.open_table("document_items")
|
self.document_items_table = await self.db.open_table("document_items")
|
||||||
else:
|
else:
|
||||||
self.document_items_table = await self.db.create_table(
|
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(
|
await self.document_items_table.create_index(
|
||||||
"document_id", config=BTree(), replace=True
|
"document_id", config=BTree(), replace=True
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ class DocumentItem(BaseModel):
|
||||||
label: str = ""
|
label: str = ""
|
||||||
text: str = ""
|
text: str = ""
|
||||||
page_numbers: list[int] = []
|
page_numbers: list[int] = []
|
||||||
|
picture_data: bytes | None = None
|
||||||
|
|
||||||
|
|
||||||
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None:
|
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None:
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,19 @@ from haiku.rag.store.engine import DocumentItemRecord, Store
|
||||||
from haiku.rag.store.models.document_item import DocumentItem
|
from haiku.rag.store.models.document_item import DocumentItem
|
||||||
from haiku.rag.utils import escape_sql_string
|
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:
|
class DocumentItemRepository:
|
||||||
"""Repository for DocumentItem operations."""
|
"""Repository for DocumentItem operations."""
|
||||||
|
|
@ -35,6 +48,7 @@ class DocumentItemRepository:
|
||||||
label=item.label,
|
label=item.label,
|
||||||
text=item.text,
|
text=item.text,
|
||||||
page_numbers=json.dumps(item.page_numbers),
|
page_numbers=json.dumps(item.page_numbers),
|
||||||
|
picture_data=item.picture_data,
|
||||||
)
|
)
|
||||||
for item in items
|
for item in items
|
||||||
]
|
]
|
||||||
|
|
@ -45,6 +59,7 @@ class DocumentItemRepository:
|
||||||
safe_id = escape_sql_string(document_id)
|
safe_id = escape_sql_string(document_id)
|
||||||
rows = await (
|
rows = await (
|
||||||
self.store.document_items_table.query()
|
self.store.document_items_table.query()
|
||||||
|
.select(_LIGHT_COLUMNS)
|
||||||
.where(f"document_id = '{safe_id}'")
|
.where(f"document_id = '{safe_id}'")
|
||||||
.to_list()
|
.to_list()
|
||||||
)
|
)
|
||||||
|
|
@ -64,7 +79,7 @@ class DocumentItemRepository:
|
||||||
Returns:
|
Returns:
|
||||||
Dict mapping document_id to sorted list of DocumentItem.
|
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:
|
if document_ids is not None:
|
||||||
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
|
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
|
||||||
query = query.where(f"document_id IN ({safe_ids})")
|
query = query.where(f"document_id IN ({safe_ids})")
|
||||||
|
|
@ -85,6 +100,7 @@ class DocumentItemRepository:
|
||||||
safe_id = escape_sql_string(document_id)
|
safe_id = escape_sql_string(document_id)
|
||||||
rows = await (
|
rows = await (
|
||||||
self.store.document_items_table.query()
|
self.store.document_items_table.query()
|
||||||
|
.select(_LIGHT_COLUMNS)
|
||||||
.where(
|
.where(
|
||||||
f"document_id = '{safe_id}' "
|
f"document_id = '{safe_id}' "
|
||||||
f"AND position >= {start} AND position <= {end}"
|
f"AND position >= {start} AND position <= {end}"
|
||||||
|
|
@ -122,3 +138,44 @@ class DocumentItemRepository:
|
||||||
self.store._assert_writable()
|
self.store._assert_writable()
|
||||||
safe_id = escape_sql_string(document_id)
|
safe_id = escape_sql_string(document_id)
|
||||||
await self.store.document_items_table.delete(f"document_id = '{safe_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
|
||||||
|
|
|
||||||
|
|
@ -84,7 +84,6 @@ from haiku.rag.store.upgrades.v0_38_0 import (
|
||||||
from haiku.rag.store.upgrades.v0_40_0 import (
|
from haiku.rag.store.upgrades.v0_40_0 import (
|
||||||
upgrade_populate_document_items as upgrade_0_40_0_document_items,
|
upgrade_populate_document_items as upgrade_0_40_0_document_items,
|
||||||
)
|
)
|
||||||
|
|
||||||
upgrades.append(upgrade_0_20_0_docling)
|
upgrades.append(upgrade_0_20_0_docling)
|
||||||
upgrades.append(upgrade_0_23_1_contextualize)
|
upgrades.append(upgrade_0_23_1_contextualize)
|
||||||
upgrades.append(upgrade_0_25_0_compress)
|
upgrades.append(upgrade_0_25_0_compress)
|
||||||
|
|
|
||||||
|
|
@ -349,3 +349,85 @@ class TestDocumentItemMigration:
|
||||||
|
|
||||||
# No items should have been created
|
# No items should have been created
|
||||||
assert await store.document_items_table.count_rows() == 0
|
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}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue