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

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

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)

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()
             )

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()