remove redundant serialization

This commit is contained in:
bryan davis 2026-06-23 15:55:24 -05:00 committed by Yiorgis Gozadinos
parent e73a4272f8
commit e95ac2e25d
No known key found for this signature in database
3 changed files with 59 additions and 4 deletions

View file

@ -35,7 +35,20 @@ def decompress_json(data: bytes) -> str:
def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]: def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
"""Split a DoclingDocument JSON into structure and pages, compress both with zstd. """Parse a DoclingDocument JSON string and compress it.
Thin wrapper over :func:`compress_docling_data` for callers that only hold
the serialized string store migrations and rebuild-from-blob, neither of
which is speed-sensitive. The ingestion hot path should call
``compress_docling_data`` with ``DoclingDocument.model_dump(mode="json")``
instead, to avoid serializing the document to a full JSON string only to
parse it straight back into a dict.
"""
return compress_docling_data(json.loads(json_str))
def compress_docling_data(data: dict) -> tuple[bytes, bytes | None]:
"""Split a DoclingDocument dict into structure and pages, compress both with zstd.
Picture image URIs are stripped from the structure blob they are stored on Picture image URIs are stripped from the structure blob they are stored on
the corresponding ``document_items.picture_data`` rows and don't need to be the corresponding ``document_items.picture_data`` rows and don't need to be
@ -43,11 +56,14 @@ def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
field is present, so each picture's ``image`` is set to ``None`` rather than field is present, so each picture's ``image`` is set to ``None`` rather than
partially mutated to keep the JSON re-validating cleanly. partially mutated to keep the JSON re-validating cleanly.
Mutates ``data`` in place (pops ``pages``, nulls picture images); callers
pass a freshly built dict (``model_dump`` / ``json.loads`` output), so this
never touches a live DoclingDocument.
Returns: Returns:
Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the
document has no page images. document has no page images.
""" """
data = json.loads(json_str)
pages = data.pop("pages", None) pages = data.pop("pages", None)
for picture in data.get("pictures") or []: for picture in data.get("pictures") or []:

View file

@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from haiku.rag.store.compression import compress_docling_split, decompress_json from haiku.rag.store.compression import compress_docling_data, decompress_json
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, PageItem from docling_core.types.doc.document import DoclingDocument, PageItem
@ -32,7 +32,7 @@ class Document(BaseModel):
Sets docling_document (zstd-compressed structure without pages), Sets docling_document (zstd-compressed structure without pages),
docling_pages (zstd-compressed page images), and docling_version. docling_pages (zstd-compressed page images), and docling_version.
""" """
structure, pages = compress_docling_split(docling_doc.model_dump_json()) structure, pages = compress_docling_data(docling_doc.model_dump(mode="json"))
self.docling_document = structure self.docling_document = structure
self.docling_pages = pages self.docling_pages = pages
self.docling_version = docling_doc.version self.docling_version = docling_doc.version

View file

@ -281,6 +281,45 @@ def test_set_docling_with_page_images():
assert "1" in pages assert "1" in pages
def test_set_docling_dict_path_matches_string_path():
"""The ingestion dict path must produce the same stored bytes as the
legacy string path.
``set_docling`` feeds ``compress_docling_data`` the dict from
``model_dump(mode="json")`` instead of round-tripping through
``model_dump_json()`` + ``json.loads``. The on-disk format must not change
in particular int-keyed ``pages`` must still serialize to string keys.
"""
import json
from docling_core.types.doc.base import Size
from docling_core.types.doc.document import DoclingDocument, PageItem
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.compression import (
compress_docling_data,
compress_docling_split,
decompress_json,
)
docling_doc = DoclingDocument(name="equivalence_test")
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
docling_doc.pages[1] = PageItem(size=Size(width=612, height=792), page_no=1)
struct_dict, pages_dict = compress_docling_data(
docling_doc.model_dump(mode="json")
)
struct_str, pages_str = compress_docling_split(docling_doc.model_dump_json())
assert json.loads(decompress_json(struct_dict)) == json.loads(
decompress_json(struct_str)
)
assert pages_dict is not None and pages_str is not None
assert json.loads(decompress_json(pages_dict)) == json.loads(
decompress_json(pages_str)
)
def test_get_page_images(): def test_get_page_images():
"""get_page_images returns requested pages from docling_pages blob.""" """get_page_images returns requested pages from docling_pages blob."""
import json import json