diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 6e2b9ace..ed085bdb 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -8,6 +8,7 @@ from typing import Any from uuid import uuid4 import lancedb +import pyarrow as pa from lancedb.pydantic import LanceModel, Vector from pydantic import Field @@ -24,12 +25,33 @@ class DocumentRecord(LanceModel): uri: str | None = None title: str | None = None metadata: str = Field(default="{}") - docling_document_json: str | None = None + docling_document: bytes | None = None docling_version: str | None = None created_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): """Create a ChunkRecord model with the specified vector dimension. @@ -281,7 +303,7 @@ class Store: self.documents_table = self.db.open_table("documents") else: self.documents_table = self.db.create_table( - "documents", schema=DocumentRecord + "documents", schema=get_documents_arrow_schema() ) # Create or get chunks table diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 5dd6105a..0d98ce9d 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -4,6 +4,8 @@ from typing import TYPE_CHECKING from cachetools import LRUCache from pydantic import BaseModel, Field +from haiku.rag.store.compression import decompress_json + if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -11,13 +13,16 @@ if TYPE_CHECKING: _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.""" if document_id in _docling_document_cache: return _docling_document_cache[document_id] from docling_core.types.doc.document import DoclingDocument + json_str = decompress_json(compressed_data) doc = DoclingDocument.model_validate_json(json_str) _docling_document_cache[document_id] = doc return doc @@ -38,7 +43,7 @@ class Document(BaseModel): uri: str | None = None title: str | None = None metadata: dict = {} - docling_document_json: str | None = None + docling_document: bytes | None = None docling_version: str | None = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) @@ -51,13 +56,14 @@ class Document(BaseModel): Returns: 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 # No caching for documents without ID if self.id is None: 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) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 3e97042f..6daef6b0 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -2,7 +2,7 @@ import json from datetime import datetime 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 @@ -35,7 +35,7 @@ class DocumentRepository: uri=record.uri, title=record.title, metadata=json.loads(record.metadata), - docling_document_json=record.docling_document_json, + docling_document=record.docling_document, docling_version=record.docling_version, created_at=datetime.fromisoformat(record.created_at) if record.created_at @@ -61,7 +61,7 @@ class DocumentRepository: uri=entity.uri, title=entity.title, metadata=json.dumps(entity.metadata), - docling_document_json=entity.docling_document_json, + docling_document=entity.docling_document, docling_version=entity.docling_version, created_at=now, updated_at=now, @@ -111,7 +111,7 @@ class DocumentRepository: "uri": entity.uri, "title": entity.title, "metadata": json.dumps(entity.metadata), - "docling_document_json": entity.docling_document_json, + "docling_document": entity.docling_document, "docling_version": entity.docling_version, "updated_at": now, }, @@ -198,5 +198,5 @@ class DocumentRepository: # Drop and recreate table to clear all data self.store.db.drop_table("documents") self.store.documents_table = self.store.db.create_table( - "documents", schema=DocumentRecord + "documents", schema=get_documents_arrow_schema() ) diff --git a/tests/test_document.py b/tests/test_document.py index 1cf25678..e062ca41 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -70,9 +70,11 @@ def test_document_get_docling_document(): import json + from haiku.rag.store.compression import compress_json + document = Document( content="Test content", - docling_document_json=json.dumps(doc_json), + docling_document=compress_json(json.dumps(doc_json)), 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.""" 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 @@ -118,13 +120,15 @@ def test_document_get_docling_document_caching(): 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 _docling_document_cache.clear() 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 @@ -157,13 +161,15 @@ def test_document_get_docling_document_no_id_no_cache(): 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 _docling_document_cache.clear() # 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() doc2 = document.get_docling_document()