diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py index 02d6a04b..4f43daa1 100644 --- a/haiku_rag_slim/haiku/rag/converters/pdf_split.py +++ b/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) diff --git a/tests/test_pdf_split.py b/tests/test_pdf_split.py index 09b3f274..423ffe0d 100644 --- a/tests/test_pdf_split.py +++ b/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