remove redundant serialization
This commit is contained in:
parent
e73a4272f8
commit
e95ac2e25d
3 changed files with 59 additions and 4 deletions
|
|
@ -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 []:
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue