diff --git a/CHANGELOG.md b/CHANGELOG.md index d5e45fb..166b946 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,39 @@ the project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html ## [Unreleased] +### Fixed + +- **Cell bounding boxes were dropped on multi-page merge** + (`adapters/docling.py`). When N per-page fragments merged into one logical + table, every `TableCell` in the injected result came back with `bbox=None`, + breaking word-level grounding for multi-page tables. `_reemit_body_row` + rebuilt each untouched body row's cells without copying the source cell's + `bbox`; it is now carried across, so untouched body rows — the large + majority on a clean multi-page table — keep their geometry. Header rows + reused from the anchor already preserved theirs; rows the merger transformed + (stitched continuations, folded overflow) genuinely have none and stay + `bbox=None`. The FIFO alignment between repeated identical-text rows and + their source occurrences is now a documented invariant: it is what guarantees + a repeated row gets its own page's boxes, not a sibling's. + +### Added + +- **Row-level page association for restored geometry** (`models.py`, + `adapters/docling.py`). A merged table's `prov` lists all N source pages, but + pages share one coordinate space, so a restored `bbox` alone cannot say which + page it belongs to. `LogicalTable.row_pages` now maps grid row index (header + rows included) to the resolved `page_no` the row's cell boxes are valid on — + resolved page numbers rather than prov indices, so the map survives + downstream prov manipulation. A missing key means the row has no single + source page (transformed rows, which carry no geometry). Recorded during + injection, the only point where each row's source fragment is known; a + consumer cannot reconstruct the mapping after the fact. +- **`TableStitcher.last_logical_tables`** exposes the `LogicalTable` results of + the most recent `stitch()` call — previously discarded after injection, + which would have left `row_pages` unreachable through the public API. + Consumers needing the map should instantiate `TableStitcher` directly; the + `stitch_tables()` convenience function does not expose the stitcher instance. + ## [0.4.4] — 2026-08-13 ### Fixed diff --git a/src/table_stitcher/__init__.py b/src/table_stitcher/__init__.py index 70e7037..29829c8 100644 --- a/src/table_stitcher/__init__.py +++ b/src/table_stitcher/__init__.py @@ -77,6 +77,12 @@ def __init__( self.logger = logging.getLogger("table_stitcher") self.adapter = adapter self.config = config or MultiPageConfig() + self.last_logical_tables: list[LogicalTable] = [] + """ + The ``LogicalTable`` results of the most recent ``stitch()`` call, + including per-row page associations (``LogicalTable.row_pages``) + populated during injection. Reset at the start of every ``stitch()``. + """ self._validate_config() # ------------------------------------------------------------------------- @@ -164,6 +170,8 @@ def stitch( Returns: The document with merged tables. """ + self.last_logical_tables = [] + if doc is None: if raise_on_error: raise StitchingError("Input document is None") @@ -213,6 +221,10 @@ def stitch( raise StitchingError(f"Failed to merge tables: {e}") from e return doc + # Exposed for consumers needing per-table merge results — notably + # LogicalTable.row_pages, which injection fills in below. + self.last_logical_tables = logical_tables + multi_page_tables = [lt for lt in logical_tables if len(lt.pages) > 1] if not multi_page_tables: diff --git a/src/table_stitcher/adapters/docling.py b/src/table_stitcher/adapters/docling.py index 608d358..278a9c6 100644 --- a/src/table_stitcher/adapters/docling.py +++ b/src/table_stitcher/adapters/docling.py @@ -223,6 +223,18 @@ def _extract_y_bounds_from_prov(prov_list: list[Any]) -> Optional[tuple[float, f return None +def _first_page_no(prov: Any) -> Optional[int]: + """Resolved page number of the first prov entry carrying one, else None.""" + if prov is None: + return None + entries = prov if isinstance(prov, list) else [prov] + for p in entries: + page_no = getattr(p, "page_no", None) + if page_no is not None: + return page_no + return None + + def _resolve_page_height(prov_list: list[Any], doc: Any, fallback: float = 842.0) -> float: """ Look up the actual page height for the first prov entry from the document. @@ -421,7 +433,7 @@ def _extract_original_header_rows( def _index_member_rows( member_data: list[Optional[TableData]], -) -> dict[tuple, list[list[TableCell]]]: +) -> dict[tuple, list[tuple[int, list[TableCell]]]]: """ Index every member fragment's original grid rows by their expanded text-vector, for col_span reconstruction during injection. @@ -440,17 +452,24 @@ def _index_member_rows( across the value columns. (The anchor's own header rows are reconstructed separately and never looked up here.) Values are buckets because a row can legitimately repeat; identical text-vectors imply identical span structure, - so any occurrence is interchangeable. + but NOT identical geometry — re-emitted cells carry their source bbox, so + occurrences are not interchangeable. Buckets are filled in fragment (page) + order and popped FIFO, which aligns the k-th DataFrame occurrence with the + k-th source occurrence because the merged DataFrame preserves fragment + concatenation order. + + Bucket entries are ``(fragment_idx, row)`` pairs so the re-emitted row's + page association can be recorded (``fragment_idx`` indexes ``member_data``). """ - index: dict[tuple, list[list[TableCell]]] = {} - for data in member_data: + index: dict[tuple, list[tuple[int, list[TableCell]]]] = {} + for frag_idx, data in enumerate(member_data): if not data or not data.grid: continue for row in data.grid: if not row: continue key = tuple((getattr(c, "text", "") or "") if c else "" for c in row) - index.setdefault(key, []).append(row) + index.setdefault(key, []).append((frag_idx, row)) return index @@ -458,7 +477,9 @@ def _reemit_body_row( orig_row: list[TableCell], table_row_idx: int, has_row_headers: bool ) -> tuple[list[TableCell], list[TableCell]]: """ - Re-emit an original grid body row at a new row offset, preserving col_span. + Re-emit an original grid body row at a new row offset, preserving col_span + and the source cell's bbox (page-local coordinates; see the merged table's + prov for page association). Returns ``(grid_row, distinct_cells)`` where ``grid_row`` repeats each spanning cell across the columns it covers (Docling grid convention) and @@ -480,6 +501,7 @@ def _reemit_body_row( col_span = getattr(cell, "col_span", 1) or 1 new_cell = TableCell( text=getattr(cell, "text", "") or "", + bbox=getattr(cell, "bbox", None), row_span=1, col_span=col_span, column_header=False, @@ -539,6 +561,8 @@ def _dataframe_to_docling_data( df: pd.DataFrame, original_data: Optional[TableData] = None, member_data: Optional[list[Optional[TableData]]] = None, + member_pages: Optional[list[Optional[int]]] = None, + row_pages_out: Optional[dict[int, int]] = None, ) -> TableData: """ Converts a pandas DataFrame back into Docling's TableData structure. @@ -554,6 +578,13 @@ def _dataframe_to_docling_data( re-emitted from their original grid cells, preserving col_span. Rows the merger transformed (stitched continuations, folded overflow) fall back to a flat 1x1 rebuild from the DataFrame. + + When ``row_pages_out`` is provided (with ``member_pages`` aligned + index-for-index with ``member_data``), it is filled with the page + association for rows that carry geometry: grid row index -> resolved + page_no. Header rows reused from the anchor map to the anchor's page + (``member_pages[0]``); re-emitted body rows map to their source fragment's + page; flat-rebuilt rows get no entry (see ``LogicalTable.row_pages``). """ if df.empty: cols = list(df.columns) if len(df.columns) > 0 else ["Column_0"] @@ -588,6 +619,13 @@ def _dataframe_to_docling_data( num_header_rows = len(orig_header_rows) grid: list[list[TableCell]] = list(orig_header_rows) table_cells: list[TableCell] = list(orig_header_cells) + # Reused header cells keep the anchor's geometry, so they map to the + # anchor's page (member_pages[0] — the anchor is always members[0]). + if row_pages_out is not None and member_pages: + anchor_page = member_pages[0] + if anchor_page is not None: + for h_idx in range(num_header_rows): + row_pages_out[h_idx] = anchor_page else: # Fall back to building flat 1x1 header from DataFrame columns num_header_rows = 1 @@ -639,7 +677,7 @@ def _dataframe_to_docling_data( # adjacent values are never fused. bucket = body_index.get(tuple(row_vals)) if bucket: - orig_row = bucket.pop(0) + frag_idx, orig_row = bucket.pop(0) if header_sigs and _is_reprinted_header(orig_row, header_sigs): # Reprinted header from a continuation page — already present as # the header block; drop it instead of duplicating into the body. @@ -648,6 +686,10 @@ def _dataframe_to_docling_data( grid_row, distinct = _reemit_body_row(orig_row, table_row_idx, has_row_headers) grid.append(grid_row) table_cells.extend(distinct) + if row_pages_out is not None and member_pages: + page = member_pages[frag_idx] if frag_idx < len(member_pages) else None + if page is not None: + row_pages_out[table_row_idx] = page emitted += 1 continue @@ -1051,12 +1093,25 @@ def restore_snapshots(): member_data = [ table_snapshots[m]["data"] for m in lt.members if m in table_snapshots ] + # Page per fragment, same filter as member_data so indices stay + # aligned. Feeds LogicalTable.row_pages: resolved page_no rather + # than a prov index, so the map survives downstream prov + # manipulation by consumers. + member_pages = [ + _first_page_no(table_snapshots[m]["prov"]) + for m in lt.members + if m in table_snapshots + ] + row_pages: dict[int, int] = {} anchor_table.data = _dataframe_to_docling_data( lt.df, original_data=original_data, member_data=member_data, + member_pages=member_pages, + row_pages_out=row_pages, ) + lt.row_pages = row_pages for satellite_idx in lt.members[1:]: satellite_table = doc.tables[satellite_idx] diff --git a/src/table_stitcher/models.py b/src/table_stitcher/models.py index 4a427be..874c8c3 100644 --- a/src/table_stitcher/models.py +++ b/src/table_stitcher/models.py @@ -163,3 +163,15 @@ class LogicalTable: merge_reason: str = "" merge_traces: list[MergeTrace] = field(default_factory=list) warnings: list[str] = field(default_factory=list) + row_pages: dict[int, int] = field(default_factory=dict) + """ + Page association for the merged table's cell geometry, keyed by grid row + index (an index into the injected ``TableData.grid``, header rows + included). The value is the resolved ``page_no`` the row's cell bboxes + (when present) are valid on. A missing key means the row has no single + source page: it was transformed by the merger (stitched continuation, + folded overflow — which can legitimately span two pages) and carries no + geometry, or its source fragment exposed no page number. Header rows + preserved from the anchor map to the anchor's page. Populated during + injection; empty for tables that were not merged. + """ diff --git a/tests/test_docling_adapter.py b/tests/test_docling_adapter.py index 5c7229d..841d66b 100644 --- a/tests/test_docling_adapter.py +++ b/tests/test_docling_adapter.py @@ -3,6 +3,7 @@ """ from types import SimpleNamespace +from typing import Optional import pandas as pd import pytest @@ -388,6 +389,158 @@ def test_without_member_data_falls_back_to_flat(self): assert all(cell.col_span == 1 for cell in td.grid[1]) +class TestBboxPreservation: + """Re-emitted body rows must carry their source cells' bounding boxes. + + The merge rebuilds the anchor's TableData from the merged DataFrame; rows + the merger left untouched are re-emitted from the original grid cells. + Dropping bbox there loses word-level grounding for every multi-page table. + """ + + @staticmethod + def _cell(text, r, c, *, header=False, bbox=None): + return TableCell( + text=text, + bbox=bbox, + row_span=1, + col_span=1, + column_header=header, + row_header=False, + start_row_offset_idx=r, + end_row_offset_idx=r + 1, + start_col_offset_idx=c, + end_col_offset_idx=c + 1, + ) + + def _fragment_with_bboxes(self) -> TableData: + from docling_core.types.doc import BoundingBox + + c = self._cell + h = [c("Name", 0, 0, header=True), c("Value", 0, 1, header=True)] + r1 = [ + c("alpha", 1, 0, bbox=BoundingBox(l=10, t=100, r=60, b=110)), + c("1", 1, 1, bbox=BoundingBox(l=70, t=100, r=90, b=110)), + ] + return TableData(num_rows=2, num_cols=2, table_cells=h + r1, grid=[h, r1]) + + def test_reemitted_rows_keep_bbox(self): + original = self._fragment_with_bboxes() + merged_df = pd.DataFrame([["alpha", "1"]], columns=["Name", "Value"]) + + td = _dataframe_to_docling_data(merged_df, original_data=original, member_data=[original]) + + body_row = td.grid[1] + assert body_row[0].bbox is not None + assert body_row[0].bbox.l == 10 and body_row[0].bbox.t == 100 + assert body_row[1].bbox is not None + assert body_row[1].bbox.l == 70 + + def test_flat_rebuild_rows_have_no_bbox(self): + """Transformed rows (no member-row match) have no geometry to carry.""" + original = self._fragment_with_bboxes() + merged_df = pd.DataFrame([["alpha (stitched)", "1"]], columns=["Name", "Value"]) + + td = _dataframe_to_docling_data(merged_df, original_data=original, member_data=[original]) + + assert all(cell.bbox is None for cell in td.grid[1]) + + +class TestRowPageMap: + """row_pages: grid row index -> resolved page_no for restored geometry. + + A merged table's prov lists all N source pages, but pages share one + coordinate space, so a restored bbox is ambiguous without a per-row page + association. The map is recorded at merge time (the only point where each + row's source fragment is known) and keyed by grid row index — resolved + page_no, not prov index, so it survives downstream prov manipulation. + A missing key is the honest signal for a transformed row. + """ + + @staticmethod + def _cell(text, r, c, *, header=False): + return TableCell( + text=text, + row_span=1, + col_span=1, + column_header=header, + row_header=False, + start_row_offset_idx=r, + end_row_offset_idx=r + 1, + start_col_offset_idx=c, + end_col_offset_idx=c + 1, + ) + + def _one_col_fragment(self, texts: list[str], *, header: Optional[str] = None) -> TableData: + rows = [] + r = 0 + if header is not None: + rows.append([self._cell(header, 0, 0, header=True)]) + r = 1 + for i, t in enumerate(texts): + rows.append([self._cell(t, r + i, 0)]) + flat = [c for row in rows for c in row] + return TableData(num_rows=len(rows), num_cols=1, table_cells=flat, grid=rows) + + def test_rebuild_maps_rows_to_source_pages(self): + anchor = self._one_col_fragment(["alpha"], header="Name") + satellite = self._one_col_fragment(["beta"]) + merged_df = pd.DataFrame([["alpha"], ["beta"], ["gamma (stitched)"]], columns=["Name"]) + + row_pages: dict[int, int] = {} + _dataframe_to_docling_data( + merged_df, + original_data=anchor, + member_data=[anchor, satellite], + member_pages=[1, 2], + row_pages_out=row_pages, + ) + + # Grid: row 0 = anchor header, row 1 = alpha (page 1), + # row 2 = beta (page 2), row 3 = transformed -> no key. + assert row_pages == {0: 1, 1: 1, 2: 2} + + def test_unknown_member_page_yields_no_key(self): + anchor = self._one_col_fragment(["alpha"], header="Name") + satellite = self._one_col_fragment(["beta"]) + merged_df = pd.DataFrame([["alpha"], ["beta"]], columns=["Name"]) + + row_pages: dict[int, int] = {} + _dataframe_to_docling_data( + merged_df, + original_data=anchor, + member_data=[anchor, satellite], + member_pages=[1, None], + row_pages_out=row_pages, + ) + + assert row_pages == {0: 1, 1: 1} # beta's fragment has no page -> no key + + def test_inject_populates_logical_table_row_pages(self): + doc = _build_doc_with_tables(2) + doc.tables[0].prov = [SimpleNamespace(page_no=1, bbox=None)] + doc.tables[1].prov = [SimpleNamespace(page_no=2, bbox=None)] + + # Rows V0 / V1 match the fragments' original body rows untouched. + merged_df = pd.DataFrame({"H0": ["V0", "V1"]}) + lt = LogicalTable(0, [0, 1], [1, 2], merged_df) + + adapter = DoclingAdapter() + adapter.inject(doc, [lt]) + + # Grid row 0 = anchor header (page 1), row 1 = V0 (page 1), + # row 2 = V1 (page 2). + assert lt.row_pages == {0: 1, 1: 1, 2: 2} + + def test_single_member_table_has_empty_row_pages(self): + doc = _build_doc_with_tables(1) + lt = LogicalTable(0, [0], [1], pd.DataFrame({"H0": ["V0"]})) + + adapter = DoclingAdapter() + adapter.inject(doc, [lt]) + + assert lt.row_pages == {} + + class TestReprintedHeaderDedup: """Reprinted continuation-page headers are dropped from the injected body, but column_header-flagged rows that don't match the header (Docling diff --git a/tests/test_public_api.py b/tests/test_public_api.py index ad0fbb2..576dacd 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -81,3 +81,56 @@ def test_table_stitcher_class_signature(): def test_stitching_error_is_exception(): assert issubclass(table_stitcher.StitchingError, Exception) + + +def test_last_logical_tables_exposed_after_stitch(): + """`TableStitcher.last_logical_tables` surfaces the merge results. + + `stitch()` returns only the document, but grounding consumers need the + per-table merge results — notably `LogicalTable.row_pages`, filled in + during injection. The stitcher keeps its most recent results reachable. + """ + import pandas as pd + + from table_stitcher.merger import is_numeric_like_colnames, tokenize + from table_stitcher.models import TableMeta + + def _meta(idx: int, start_page: int, *, headerless: bool) -> TableMeta: + df = pd.DataFrame({"Name": ["Alice"], "Age": ["30"]}) + return TableMeta( + idx=idx, + df=df, + start_page=start_page, + pages=[start_page], + width=df.shape[1], + header_tokens=tokenize(" ".join(str(c) for c in df.columns)), + first_row_tokens=tokenize(" ".join(str(x) for x in df.iloc[0].tolist())), + raw_columns=[str(c) for c in df.columns], + vert_center=None, + vert_top=None, + vert_bottom=None, + is_header_orphan=False, + is_data_orphan=False, + numeric_like_cols=is_numeric_like_colnames([str(c) for c in df.columns]), + row_count=df.shape[0], + is_headerless=headerless, + ) + + class _StubAdapter: + def extract(self, doc, config): + return [_meta(0, 1, headerless=False), _meta(1, 2, headerless=True)] + + def inject(self, doc, logical_tables): + for lt in logical_tables: + lt.row_pages = {0: 1} + return doc + + stitcher = table_stitcher.TableStitcher(_StubAdapter()) + assert stitcher.last_logical_tables == [] + + stitcher.stitch(object()) + + assert len(stitcher.last_logical_tables) == 1 + lt = stitcher.last_logical_tables[0] + assert lt.members == [0, 1] + assert lt.row_pages == {0: 1}