re-constitute split pdfs on the fly rather than holding extra copies

haiku_rag_slim/haiku/rag/converters/pdf_split.py

@@ -71,6 +71,63 @@ def iter_pdf_slices(
             src.close()
 
 
+def _incremental_merge_available() -> bool:
+    """Whether docling-core exposes the internal merge buffer this module
+    drives incrementally. False on versions that renamed/removed it, in which
+    case ``_SliceMerger`` falls back to the public ``concatenate``."""
+    from docling_core.types.doc.document import DoclingDocument
+
+    return hasattr(DoclingDocument, "_DocIndex") and hasattr(
+        DoclingDocument, "_update_from_index"
+    )
+
+
+class _SliceMerger:
+    """Merges converted PDF-slice documents into one ``DoclingDocument``.
+
+    ``DoclingDocument.concatenate`` indexes each input into an internal
+    ``_DocIndex`` merge buffer (deep-copying its items), then materialises the
+    result. Collecting every slice up front and concatenating at the end thus
+    holds all slices *and* the merged copy at once — roughly twice the
+    document. This merger drives the same ``_DocIndex`` incrementally: each
+    slice is folded in and released before the next is converted, so peak
+    working set is ~one merged document plus one slice. The sequence of
+    ``index()`` calls is identical to ``concatenate``'s, so the merged result
+    is byte-for-byte the same.
+
+    ``_DocIndex``/``_update_from_index`` are docling-core internals; when they
+    are unavailable (version drift, per ``_incremental_merge_available``) the
+    merger collects slices and calls the public ``concatenate`` instead.
+    """
+
+    def __init__(self) -> None:
+        from docling_core.types.doc.document import DoclingDocument
+
+        self._incremental = _incremental_merge_available()
+        self._index = DoclingDocument._DocIndex() if self._incremental else None
+        self._collected: list[DoclingDocument] = []
+
+    def add(self, slice_doc: "DoclingDocument") -> None:
+        """Fold one slice into the merge. CPU-bound (deep-copies the slice's
+        items), so call via ``asyncio.to_thread`` to keep it off the loop."""
+        if self._incremental:
+            assert self._index is not None
+            self._index.index(slice_doc)
+        else:
+            self._collected.append(slice_doc)
+
+    def result(self) -> "DoclingDocument":
+        """Materialise the merged document. Call via ``asyncio.to_thread``."""
+        from docling_core.types.doc.document import DoclingDocument
+
+        if self._incremental:
+            assert self._index is not None
+            merged = DoclingDocument(name="")
+            merged._update_from_index(self._index)
+            return merged
+        return DoclingDocument.concatenate(self._collected)
+
+
 async def convert_pdf_with_splitting(
     converter: "DocumentConverter",
     path: Path,
@@ -81,19 +138,21 @@ async def convert_pdf_with_splitting(
     merged ``DoclingDocument``.
 
     Slices are produced lazily — only one slice's bytes live in memory at a
-    time, so peak working set stays bounded regardless of page count. A
-    single slice failure aborts the whole job (raised as ``ValueError`` so
-    the ingester pipeline classifies it as ``TransientError`` and the queue
-    retries the entire document). Per-slice retry would risk interleaved
-    partial state with subsequent runs and is not worth the complexity.
+    time — and each converted slice is folded into the running merge and
+    released before the next is converted (see ``_SliceMerger``), so peak
+    working set is ~one merged document plus one slice rather than every slice
+    at once. A single slice failure aborts the whole job (raised as
+    ``ValueError`` so the ingester pipeline classifies it as ``TransientError``
+    and the queue retries the entire document). Per-slice retry would risk
+    interleaved partial state with subsequent runs and is not worth the
+    complexity.
     """
-    from docling_core.types.doc.document import DoclingDocument
 
     def _next(it):
         return next(it, _SENTINEL)
 
     it = iter_pdf_slices(path, slice_size)
-    converted: list[DoclingDocument] = []
+    merger = _SliceMerger()
     try:
         while True:
             slice_item = await asyncio.to_thread(_next, it)
@@ -122,7 +181,12 @@ async def convert_pdf_with_splitting(
                         raise ValueError(
                             f"Failed to convert slice pages {start}-{end} of {path}: {exc}"
                         ) from exc
-                converted.append(slice_doc)
+                # Fold the slice in and drop it before the next slice converts,
+                # so peak memory holds ~one merged doc + one slice rather than
+                # every slice at once. index() deep-copies, so run it off the
+                # event loop.
+                await asyncio.to_thread(merger.add, slice_doc)
+                del slice_doc
             finally:
                 # `delete=False` is required so the converter (which opens
                 # tmp_path itself) sees a fully written, closed file. Unlink
@@ -136,7 +200,7 @@ async def convert_pdf_with_splitting(
         # Off the event loop because the close path acquires the lock.
         await asyncio.to_thread(it.close)
 
-    # Merge off the event loop: concatenating slice documents that carry
-    # inlined base64 page/picture images is CPU-heavy and proportional to the
-    # total document size, so running it inline would block other coroutines.
-    return await asyncio.to_thread(DoclingDocument.concatenate, converted)
+    # Materialise off the event loop: the merge buffer holds inlined base64
+    # page/picture images, so building the result is CPU-heavy and proportional
+    # to total document size — running it inline would block other coroutines.
+    return await asyncio.to_thread(merger.result)

tests/test_pdf_split.py

@@ -158,16 +158,18 @@ async def test_convert_aborts_and_cleans_up_on_mid_stream_slice_failure(
 
 
 @pytest.mark.asyncio
-async def test_concatenate_runs_off_event_loop_thread(tmp_path, monkeypatch):
-    """DoclingDocument.concatenate merges slice documents that carry inlined
-    base64 page/picture images — CPU-heavy and proportional to total document
-    size. It must run off the event-loop thread so it doesn't stall other
-    workers' coroutines. Capture the thread it runs on and assert it is not the
-    event-loop thread."""
+async def test_slice_merge_runs_off_event_loop_thread(tmp_path, monkeypatch):
+    """Folding each converted slice into the merge buffer deep-copies items
+    carrying inlined base64 page/picture images — CPU-heavy and proportional to
+    document size. It must run off the event-loop thread so it doesn't stall
+    other workers' coroutines. Capture the thread each fold runs on and assert
+    none is the event-loop thread."""
     import threading
 
     from docling_core.types.doc.document import DoclingDocument
 
+    from haiku.rag.converters import pdf_split
+
     src = _make_pdf(4, tmp_path)
 
     class _Converter:
@@ -177,14 +179,13 @@ async def test_concatenate_runs_off_event_loop_thread(tmp_path, monkeypatch):
     event_loop_thread = threading.current_thread()
     called_from: list[threading.Thread] = []
 
-    def spy(docs):
-        called_from.append(threading.current_thread())
-        # Return a slice doc rather than exercising the real concatenate —
-        # this test only asserts the dispatch thread, not merge correctness
-        # (covered by test_concatenate_shifts_page_nos_and_unique_self_refs).
-        return docs[0]
+    real_add = pdf_split._SliceMerger.add
 
-    monkeypatch.setattr(DoclingDocument, "concatenate", staticmethod(spy))
+    def spy_add(self, slice_doc):
+        called_from.append(threading.current_thread())
+        return real_add(self, slice_doc)
+
+    monkeypatch.setattr(pdf_split._SliceMerger, "add", spy_add)
 
     await convert_pdf_with_splitting(
         _Converter(),  # ty: ignore[invalid-argument-type]
@@ -193,13 +194,91 @@ async def test_concatenate_runs_off_event_loop_thread(tmp_path, monkeypatch):
         slice_size=2,
     )
 
-    assert called_from, "concatenate was never called"
-    assert called_from[0] is not event_loop_thread, (
-        "DoclingDocument.concatenate ran on the event-loop thread; it must be "
-        "dispatched via asyncio.to_thread"
+    assert called_from, "slice merge was never invoked"
+    assert all(t is not event_loop_thread for t in called_from), (
+        "slice merge ran on the event-loop thread; it must be dispatched via "
+        "asyncio.to_thread"
     )
 
 
+def _one_page_slice(name: str, text: str):
+    """A single-page DoclingDocument with one provenanced text item, mimicking
+    a converted PDF slice (docling numbers each slice's pages from 1)."""
+    from docling_core.types.doc.base import BoundingBox, CoordOrigin
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        PageItem,
+        ProvenanceItem,
+        Size,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+
+    d = DoclingDocument(name=name)
+    d.pages[1] = PageItem(page_no=1, size=Size(width=595.0, height=842.0))
+    d.add_text(
+        label=DocItemLabel.TEXT,
+        text=text,
+        prov=ProvenanceItem(
+            page_no=1,
+            bbox=BoundingBox(
+                l=0.0, t=0.0, r=100.0, b=20.0, coord_origin=CoordOrigin.TOPLEFT
+            ),
+            charspan=(0, len(text)),
+        ),
+    )
+    return d
+
+
+def test_incremental_merge_matches_concatenate():
+    """The incremental _SliceMerger must produce a byte-identical document to
+    the public DoclingDocument.concatenate over the same slices — it issues the
+    same _DocIndex.index() sequence, just interleaved with conversion. Also the
+    canary for docling-core internals drift: if the incremental path silently
+    disappears, the _incremental assertion fails loudly."""
+    pytest.importorskip("docling_core")
+    from docling_core.types.doc.document import DoclingDocument
+
+    from haiku.rag.converters.pdf_split import _SliceMerger
+
+    slices = [("s1", "alpha"), ("s2", "beta"), ("s3", "gamma")]
+
+    batch = DoclingDocument.concatenate([_one_page_slice(n, t) for n, t in slices])
+
+    merger = _SliceMerger()
+    assert merger._incremental, "expected the incremental path on this docling-core"
+    for n, t in slices:
+        merger.add(_one_page_slice(n, t))
+    incremental = merger.result()
+
+    assert incremental.model_dump_json() == batch.model_dump_json()
+
+
+def test_slice_merger_falls_back_to_concatenate(monkeypatch):
+    """When the docling-core merge internals are unavailable, the merger falls
+    back to the public concatenate and still merges correctly."""
+    pytest.importorskip("docling_core")
+    from docling_core.types.doc.document import DoclingDocument
+
+    from haiku.rag.converters import pdf_split
+
+    # Simulate version drift without breaking concatenate itself (which uses
+    # _DocIndex internally): force the capability probe to report unavailable.
+    monkeypatch.setattr(pdf_split, "_incremental_merge_available", lambda: False)
+
+    slices = [("s1", "alpha"), ("s2", "beta")]
+    expected = DoclingDocument.concatenate(
+        [_one_page_slice(n, t) for n, t in slices]
+    )
+
+    merger = pdf_split._SliceMerger()
+    assert not merger._incremental
+    for n, t in slices:
+        merger.add(_one_page_slice(n, t))
+    merged = merger.result()
+
+    assert merged.model_dump_json() == expected.model_dump_json()
+
+
 def test_concatenate_shifts_page_nos_and_unique_self_refs():
     """Pins the docling-core contract we rely on: when two docs (each with
     items on page 1) are concatenated, the second doc's items move to page 2