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:
parent
face6f95bc
commit
afc5b02285
4 changed files with 52 additions and 18 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue