Update schema to store docling_document as compressed bytes

- Change DocumentRecord.docling_document_json (str) to docling_document (bytes)
- Use large_binary Arrow type (64-bit offsets) to avoid 2GB column limit
- Decompress in Document.get_docling_document() using gzip
- Update DocumentRepository field mappings
This commit is contained in:
Yiorgis Gozadinos 2026-01-10 15:17:27 +02:00
parent face6f95bc
commit afc5b02285
No known key found for this signature in database
4 changed files with 52 additions and 18 deletions

View file

@ -8,6 +8,7 @@ from typing import Any
from uuid import uuid4 from uuid import uuid4
import lancedb import lancedb
import pyarrow as pa
from lancedb.pydantic import LanceModel, Vector from lancedb.pydantic import LanceModel, Vector
from pydantic import Field from pydantic import Field
@ -24,12 +25,33 @@ class DocumentRecord(LanceModel):
uri: str | None = None uri: str | None = None
title: str | None = None title: str | None = None
metadata: str = Field(default="{}") metadata: str = Field(default="{}")
docling_document_json: str | None = None docling_document: bytes | None = None
docling_version: str | None = None docling_version: str | None = None
created_at: str = Field(default_factory=lambda: "") created_at: str = Field(default_factory=lambda: "")
updated_at: str = Field(default_factory=lambda: "") updated_at: str = Field(default_factory=lambda: "")
def get_documents_arrow_schema() -> pa.Schema:
"""Generate Arrow schema for documents table with large_binary for docling_document.
LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets
and is limited to ~2GB per column in a fragment. When many large documents
(with embedded page images) are grouped in a single fragment, this limit is
exceeded, causing "byte array offset overflow" panics.
This function overrides the default mapping to use `large_binary` instead,
which has 64-bit offsets and no practical size limit.
"""
base_schema = DocumentRecord.to_arrow_schema()
fields = []
for field in base_schema:
if field.name == "docling_document":
fields.append(pa.field("docling_document", pa.large_binary()))
else:
fields.append(field)
return pa.schema(fields)
def create_chunk_model(vector_dim: int): def create_chunk_model(vector_dim: int):
"""Create a ChunkRecord model with the specified vector dimension. """Create a ChunkRecord model with the specified vector dimension.
@ -281,7 +303,7 @@ class Store:
self.documents_table = self.db.open_table("documents") self.documents_table = self.db.open_table("documents")
else: else:
self.documents_table = self.db.create_table( self.documents_table = self.db.create_table(
"documents", schema=DocumentRecord "documents", schema=get_documents_arrow_schema()
) )
# Create or get chunks table # Create or get chunks table

View file

@ -4,6 +4,8 @@ from typing import TYPE_CHECKING
from cachetools import LRUCache from cachetools import LRUCache
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from haiku.rag.store.compression import decompress_json
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
@ -11,13 +13,16 @@ if TYPE_CHECKING:
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
def _get_cached_docling_document(document_id: str, json_str: str) -> "DoclingDocument": def _get_cached_docling_document(
document_id: str, compressed_data: bytes
) -> "DoclingDocument":
"""Get or parse DoclingDocument with LRU caching by document ID.""" """Get or parse DoclingDocument with LRU caching by document ID."""
if document_id in _docling_document_cache: if document_id in _docling_document_cache:
return _docling_document_cache[document_id] return _docling_document_cache[document_id]
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data)
doc = DoclingDocument.model_validate_json(json_str) doc = DoclingDocument.model_validate_json(json_str)
_docling_document_cache[document_id] = doc _docling_document_cache[document_id] = doc
return doc return doc
@ -38,7 +43,7 @@ class Document(BaseModel):
uri: str | None = None uri: str | None = None
title: str | None = None title: str | None = None
metadata: dict = {} metadata: dict = {}
docling_document_json: str | None = None docling_document: bytes | None = None
docling_version: str | None = None docling_version: str | None = None
created_at: datetime = Field(default_factory=datetime.now) created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now)
@ -51,13 +56,14 @@ class Document(BaseModel):
Returns: Returns:
The parsed DoclingDocument, or None if not stored or no ID. The parsed DoclingDocument, or None if not stored or no ID.
""" """
if self.docling_document_json is None: if self.docling_document is None:
return None return None
# No caching for documents without ID # No caching for documents without ID
if self.id is None: if self.id is None:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
return DoclingDocument.model_validate_json(self.docling_document_json) json_str = decompress_json(self.docling_document)
return DoclingDocument.model_validate_json(json_str)
return _get_cached_docling_document(self.id, self.docling_document_json) return _get_cached_docling_document(self.id, self.docling_document)

View file

@ -2,7 +2,7 @@ import json
from datetime import datetime from datetime import datetime
from uuid import uuid4 from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
@ -35,7 +35,7 @@ class DocumentRepository:
uri=record.uri, uri=record.uri,
title=record.title, title=record.title,
metadata=json.loads(record.metadata), metadata=json.loads(record.metadata),
docling_document_json=record.docling_document_json, docling_document=record.docling_document,
docling_version=record.docling_version, docling_version=record.docling_version,
created_at=datetime.fromisoformat(record.created_at) created_at=datetime.fromisoformat(record.created_at)
if record.created_at if record.created_at
@ -61,7 +61,7 @@ class DocumentRepository:
uri=entity.uri, uri=entity.uri,
title=entity.title, title=entity.title,
metadata=json.dumps(entity.metadata), metadata=json.dumps(entity.metadata),
docling_document_json=entity.docling_document_json, docling_document=entity.docling_document,
docling_version=entity.docling_version, docling_version=entity.docling_version,
created_at=now, created_at=now,
updated_at=now, updated_at=now,
@ -111,7 +111,7 @@ class DocumentRepository:
"uri": entity.uri, "uri": entity.uri,
"title": entity.title, "title": entity.title,
"metadata": json.dumps(entity.metadata), "metadata": json.dumps(entity.metadata),
"docling_document_json": entity.docling_document_json, "docling_document": entity.docling_document,
"docling_version": entity.docling_version, "docling_version": entity.docling_version,
"updated_at": now, "updated_at": now,
}, },
@ -198,5 +198,5 @@ class DocumentRepository:
# Drop and recreate table to clear all data # Drop and recreate table to clear all data
self.store.db.drop_table("documents") self.store.db.drop_table("documents")
self.store.documents_table = self.store.db.create_table( self.store.documents_table = self.store.db.create_table(
"documents", schema=DocumentRecord "documents", schema=get_documents_arrow_schema()
) )

View file

@ -70,9 +70,11 @@ def test_document_get_docling_document():
import json import json
from haiku.rag.store.compression import compress_json
document = Document( document = Document(
content="Test content", content="Test content",
docling_document_json=json.dumps(doc_json), docling_document=compress_json(json.dumps(doc_json)),
docling_version="1.3.0", docling_version="1.3.0",
) )
@ -88,7 +90,7 @@ def test_document_get_docling_document_none():
"""Test get_docling_document returns None when not stored.""" """Test get_docling_document returns None when not stored."""
document = Document(content="Test content") document = Document(content="Test content")
assert document.docling_document_json is None assert document.docling_document is None
assert document.get_docling_document() is None assert document.get_docling_document() is None
@ -118,13 +120,15 @@ def test_document_get_docling_document_caching():
import json import json
json_str = json.dumps(doc_json) from haiku.rag.store.compression import compress_json
compressed = compress_json(json.dumps(doc_json))
# Clear cache to get clean state # Clear cache to get clean state
_docling_document_cache.clear() _docling_document_cache.clear()
document = Document( document = Document(
id="test-doc-id", content="Test content", docling_document_json=json_str id="test-doc-id", content="Test content", docling_document=compressed
) )
# First call - not in cache # First call - not in cache
@ -157,13 +161,15 @@ def test_document_get_docling_document_no_id_no_cache():
import json import json
json_str = json.dumps(doc_json) from haiku.rag.store.compression import compress_json
compressed = compress_json(json.dumps(doc_json))
# Clear cache # Clear cache
_docling_document_cache.clear() _docling_document_cache.clear()
# Document without ID # Document without ID
document = Document(content="Test content", docling_document_json=json_str) document = Document(content="Test content", docling_document=compressed)
doc1 = document.get_docling_document() doc1 = document.get_docling_document()
doc2 = document.get_docling_document() doc2 = document.get_docling_document()