陈辅元
Cursor
@@ -0,0 +1,177 @@
|
||||
"""Cross-page table detection and merging."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType
|
||||
from rag_cut.parsers.pdf.tables import (
|
||||
detect_header_row_count,
|
||||
header_signature,
|
||||
normalize_rows,
|
||||
rows_to_markdown,
|
||||
split_body_and_footnotes,
|
||||
)
|
||||
|
||||
|
||||
def _page(block: Block) -> int | None:
|
||||
page = block.meta.get("page")
|
||||
return int(page) if page is not None else None
|
||||
|
||||
|
||||
def _header_rows(block: Block) -> int:
|
||||
rows = block.meta.get("rows") or []
|
||||
return int(block.meta.get("header_rows") or detect_header_row_count(normalize_rows(rows)) or 1)
|
||||
|
||||
|
||||
def _header_sig(block: Block) -> tuple[tuple[str, ...], ...]:
|
||||
if block.meta.get("header_signature"):
|
||||
raw = block.meta["header_signature"]
|
||||
return tuple(tuple(r) for r in raw)
|
||||
rows = normalize_rows(block.meta.get("rows") or [])
|
||||
return header_signature(rows, _header_rows(block))
|
||||
|
||||
|
||||
def _column_count(block: Block) -> int:
|
||||
rows = normalize_rows(block.meta.get("rows") or [])
|
||||
return max((len(r) for r in rows), default=0)
|
||||
|
||||
|
||||
def _similar_columns(a: Block, b: Block) -> bool:
|
||||
ca, cb = _column_count(a), _column_count(b)
|
||||
if ca == 0 or cb == 0:
|
||||
return False
|
||||
return ca == cb or abs(ca - cb) <= 1
|
||||
|
||||
|
||||
def _repeated_header(rows: list[list[str]], sig: tuple[tuple[str, ...], ...]) -> int:
|
||||
"""Return number of leading rows in `rows` that repeat the header signature."""
|
||||
if not sig:
|
||||
return 0
|
||||
n = len(sig)
|
||||
if len(rows) < n:
|
||||
return 0
|
||||
if header_signature(rows, n) == sig:
|
||||
return n
|
||||
if n == 1 and rows and tuple(rows[0]) == sig[0]:
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
def _can_merge_continuation(prev: Block, nxt: Block) -> bool:
|
||||
if prev.type != BlockType.TABLE or nxt.type != BlockType.TABLE:
|
||||
return False
|
||||
|
||||
prev_page = _page(prev)
|
||||
nxt_page = _page(nxt)
|
||||
if prev_page is None or nxt_page is None or nxt_page != prev_page + 1:
|
||||
return False
|
||||
if not _similar_columns(prev, nxt):
|
||||
return False
|
||||
|
||||
sig = _header_sig(prev)
|
||||
if not sig:
|
||||
return False
|
||||
|
||||
rows_b = normalize_rows(nxt.meta.get("rows") or [])
|
||||
if not rows_b:
|
||||
return False
|
||||
|
||||
if _repeated_header(rows_b, sig) > 0:
|
||||
return True
|
||||
|
||||
# Continuation without repeated header: similar width and no title on next table
|
||||
if nxt.meta.get("table_title"):
|
||||
return False
|
||||
|
||||
prev_bbox = prev.meta.get("bbox") or []
|
||||
nxt_bbox = nxt.meta.get("bbox") or []
|
||||
if len(prev_bbox) == 4 and len(nxt_bbox) == 4:
|
||||
prev_width = prev_bbox[2] - prev_bbox[0]
|
||||
nxt_width = nxt_bbox[2] - nxt_bbox[0]
|
||||
if prev_width > 0 and abs(prev_width - nxt_width) / prev_width <= 0.15:
|
||||
return True
|
||||
|
||||
return _column_count(prev) == _column_count(nxt)
|
||||
|
||||
|
||||
def _merge_two_tables(prev: Block, nxt: Block) -> Block:
|
||||
rows_a = normalize_rows(prev.meta.get("rows") or [])
|
||||
rows_b = normalize_rows(nxt.meta.get("rows") or [])
|
||||
header_rows = _header_rows(prev)
|
||||
sig = _header_sig(prev)
|
||||
|
||||
skip = _repeated_header(rows_b, sig)
|
||||
merged_rows = rows_a + rows_b[skip:]
|
||||
|
||||
body_rows, _, foot_a = split_body_and_footnotes(rows_a, header_rows)
|
||||
_, _, foot_b = split_body_and_footnotes(rows_b, skip or header_rows)
|
||||
footnotes = " ".join(x for x in (prev.meta.get("footnotes") or foot_a, foot_b) if x).strip()
|
||||
|
||||
pages = sorted(set((prev.meta.get("pages") or [_page(prev)]) + [_page(nxt)]))
|
||||
pages = [p for p in pages if p is not None]
|
||||
bboxes = list(prev.meta.get("bboxes") or [])
|
||||
if prev.meta.get("bbox"):
|
||||
bboxes.append({"page": _page(prev), "bbox": prev.meta["bbox"]})
|
||||
if nxt.meta.get("bbox"):
|
||||
bboxes.append({"page": _page(nxt), "bbox": nxt.meta["bbox"]})
|
||||
|
||||
prev_bbox = prev.meta.get("bbox") or [0, 0, 0, 0]
|
||||
nxt_bbox = nxt.meta.get("bbox") or prev_bbox
|
||||
merged_bbox = [
|
||||
min(prev_bbox[0], nxt_bbox[0]),
|
||||
min(prev_bbox[1], nxt_bbox[1]),
|
||||
max(prev_bbox[2], nxt_bbox[2]),
|
||||
max(prev_bbox[3], nxt_bbox[3]),
|
||||
]
|
||||
|
||||
md = rows_to_markdown(merged_rows, header_rows=header_rows, include_footnotes=footnotes)
|
||||
meta = dict(prev.meta)
|
||||
meta.update(
|
||||
{
|
||||
"rows": merged_rows,
|
||||
"row_count": len(merged_rows),
|
||||
"col_count": max((len(r) for r in merged_rows), default=0),
|
||||
"header_rows": header_rows,
|
||||
"header_signature": [list(r) for r in sig],
|
||||
"footnotes": footnotes,
|
||||
"pages": pages,
|
||||
"page": pages[0] if pages else prev.meta.get("page"),
|
||||
"bbox": merged_bbox,
|
||||
"bboxes": bboxes,
|
||||
"cross_page": len(pages) > 1,
|
||||
"merged_table_count": int(prev.meta.get("merged_table_count") or 1) + 1,
|
||||
}
|
||||
)
|
||||
if nxt.meta.get("following_text") and not meta.get("following_text"):
|
||||
meta["following_text"] = nxt.meta.get("following_text")
|
||||
|
||||
return prev.model_copy(update={"markdown": md, "meta": meta})
|
||||
|
||||
|
||||
def merge_cross_page_tables(blocks: list[Block]) -> list[Block]:
|
||||
"""Merge consecutive cross-page table blocks that share headers/structure."""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
result: list[Block] = []
|
||||
i = 0
|
||||
while i < len(blocks):
|
||||
current = blocks[i]
|
||||
if current.type != BlockType.TABLE:
|
||||
result.append(current)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
merged = current
|
||||
j = i + 1
|
||||
while j < len(blocks):
|
||||
nxt = blocks[j]
|
||||
if nxt.type == BlockType.TABLE and _can_merge_continuation(merged, nxt):
|
||||
merged = _merge_two_tables(merged, nxt)
|
||||
j += 1
|
||||
continue
|
||||
break
|
||||
|
||||
result.append(merged)
|
||||
i = j
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user