Merge pull request #465 from bd-mkt/perf_redundant_serialization

perf: avoid redundant serialization in docling ingestion path

CHANGELOG.md

@@ -4,6 +4,7 @@
 ### Fixed
 
 - zstd compression uses a fresh `zstandard` compressor/decompressor per call instead of shared module-level singletons, fixing process segfaults when ingester workers compress documents concurrently (Python < 3.14).
+- Ingestion compresses the DoclingDocument structure directly from the in-memory dict, removing one full-size serialized copy from per-document peak memory.
 
 ## [0.61.0] - 2026-06-23
 

haiku_rag_slim/haiku/rag/client/rebuild.py

@@ -643,7 +643,7 @@ def _apply_descriptions_sync(
             pic.meta = PictureMeta()
         pic.meta.description = DescriptionMetaField(text=text)
 
-    structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json())
+    structure_bytes, _ = compress_docling_split(docling_doc.model_dump(mode="json"))
     doc.docling_document = structure_bytes
     doc.docling_version = docling_doc.version
     return len(descriptions)

haiku_rag_slim/haiku/rag/store/compression.py

@@ -34,8 +34,8 @@ def decompress_json(data: bytes) -> str:
     return _zstd_decompress(data).decode("utf-8")
 
 
-def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
+def compress_docling_split(data: dict) -> tuple[bytes, bytes | None]:
-    """Split a DoclingDocument JSON into structure and pages, compress both with zstd.
+    """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
     the corresponding ``document_items.picture_data`` rows and don't need to be
@@ -43,21 +43,21 @@ 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
     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:
         Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the
         document has no page images.
     """
-    data = json.loads(json_str)
     pages = data.pop("pages", None)
 
     for picture in data.get("pictures") or []:
         if isinstance(picture, dict):
             picture["image"] = None
 
-    structure_bytes = _zstd_compress(json.dumps(data).encode("utf-8"))
+    structure_bytes = compress_json(json.dumps(data))
-
+    pages_bytes = compress_json(json.dumps(pages)) if pages else None
-    pages_bytes = None
-    if pages:
-        pages_bytes = _zstd_compress(json.dumps(pages).encode("utf-8"))
 
     return structure_bytes, pages_bytes

haiku_rag_slim/haiku/rag/store/models/document.py

@@ -32,7 +32,7 @@ class Document(BaseModel):
         Sets docling_document (zstd-compressed structure without pages),
         docling_pages (zstd-compressed page images), and docling_version.
         """
-        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+        structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
         self.docling_document = structure
         self.docling_pages = pages
         self.docling_version = docling_doc.version

haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py

@@ -58,7 +58,7 @@ async def _apply_split_pages_zstd(store: Store) -> None:  # pragma: no cover
                 json_str = docling_blob.decode("utf-8")
 
             # Split structure and pages, re-compress with zstd
-            structure_bytes, pages_bytes = compress_docling_split(json_str)
+            structure_bytes, pages_bytes = compress_docling_split(json.loads(json_str))
 
         metadata_raw = row.get("metadata")
         metadata_str = (

tests/store/test_compression.py

@@ -54,8 +54,7 @@ class TestDoclingCompressionSplit:
             "texts": [{"text": "hello"}],
             "pages": {"1": {"image": "base64data"}, "2": {"image": "more"}},
         }
-        json_str = json.dumps(data)
+        structure_bytes, pages_bytes = compress_docling_split(data)
-        structure_bytes, pages_bytes = compress_docling_split(json_str)
 
         assert structure_bytes is not None
         assert pages_bytes is not None
@@ -73,8 +72,7 @@ class TestDoclingCompressionSplit:
 
     def test_split_without_pages(self):
         data = {"name": "test_doc", "texts": []}
-        json_str = json.dumps(data)
+        structure_bytes, pages_bytes = compress_docling_split(data)
-        structure_bytes, pages_bytes = compress_docling_split(json_str)
 
         assert structure_bytes is not None
         assert pages_bytes is None
@@ -84,8 +82,7 @@ class TestDoclingCompressionSplit:
 
     def test_split_with_empty_pages(self):
         data = {"name": "test_doc", "texts": [], "pages": {}}
-        json_str = json.dumps(data)
+        structure_bytes, pages_bytes = compress_docling_split(data)
-        structure_bytes, pages_bytes = compress_docling_split(json_str)
 
         assert structure_bytes is not None
         assert pages_bytes is None

tests/store/test_document_items.py

@@ -397,8 +397,7 @@ class TestDocumentItemMigration:
         from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items
 
         docling_doc = _make_docling_doc()
-        json_str = docling_doc.model_dump_json()
+        structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
-        structure, pages = compress_docling_split(json_str)
 
         # Create a database at a pre-migration version with a document
         async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
@@ -687,7 +686,7 @@ class TestCompressDoclingSplitStripsPictureUris:
             "pages": {},
         }
 
-        structure_bytes, pages_bytes = compress_docling_split(json.dumps(doc_json))
+        structure_bytes, pages_bytes = compress_docling_split(doc_json)
         decoded = json.loads(decompress_json(structure_bytes))
 
         for pic in decoded["pictures"]:

tests/store/test_v0_40_0_migration.py

@@ -35,7 +35,7 @@ async def test_populate_handles_extra_columns_on_items_table(temp_db_path):
     accepted even though it only writes the original 6 columns.
     """
     docling_doc = _simple_docling_doc()
-    structure, pages = compress_docling_split(docling_doc.model_dump_json())
+    structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
 
     async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
         # _init_tables already created document_items with the latest schema.

tests/store/test_v0_48_0_migration.py

@@ -26,7 +26,7 @@ class TestV0_48_0Migration:
 
     async def test_backfill_populates_levels(self, temp_db_path):
         docling_doc = _docling_with_levels()
-        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+        structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
 
         async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
             await store.set_haiku_version("0.45.0")
@@ -79,7 +79,7 @@ class TestV0_48_0Migration:
 
     async def test_backfill_idempotent(self, temp_db_path):
         docling_doc = _docling_with_levels()
-        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+        structure, pages = compress_docling_split(docling_doc.model_dump(mode="json"))
 
         async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
             await store.set_haiku_version("0.45.0")

tests/test_document.py

@@ -281,6 +281,36 @@ def test_set_docling_with_page_images():
     assert "1" in pages
 
 
+def test_compress_docling_split_dict_sources_match():
+    """Both ways of building the compression input dict must yield identical
+    stored bytes: ``model_dump(mode="json")`` (used at ingest) and
+    ``json.loads(model_dump_json())`` (used when migrating a stored string).
+    In particular int-keyed ``pages`` must serialize to string keys either way.
+    """
+    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_split
+
+    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_dump, pages_dump = compress_docling_split(
+        docling_doc.model_dump(mode="json")
+    )
+    struct_str, pages_str = compress_docling_split(
+        json.loads(docling_doc.model_dump_json())
+    )
+
+    assert struct_dump == struct_str
+    assert pages_dump is not None and pages_str is not None
+    assert pages_dump == pages_str
+
+
 def test_get_page_images():
     """get_page_images returns requested pages from docling_pages blob."""
     import json