Skip to content

perf(tables): Vectorize slot_into_containers token-to-cell scoring - #530

Open
KRRT7 wants to merge 1 commit into
Unstructured-IO:mainfrom
KRRT7:perf/vectorize-slot-into-containers
Open

KRRT7 wants to merge 1 commit into
Unstructured-IO:mainfrom
KRRT7:perf/vectorize-slot-into-containers

Conversation

@KRRT7

@KRRT7 KRRT7 commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

Summary

slot_into_containers scored every token against every cell in pure Python, building a new Rect and a dict for each pair and fully sorting each token's scores just to find the best one. structure_to_cells calls it twice per table, and refine_rows/refine_columns call it again through nms_by_containment. On large tables it accounts for ~99% of table postprocessing time.

This computes the overlaps with numpy in chunks of at most ~1M token-cell pairs, so temporary memory stays proportional to the number of cells rather than tokens × cells, the same way #521/#522 bound memory.

Overlaps with #527, which rewrites the same function. #527 keeps the per-pair Python loop and picks the first maximum in one pass instead of sorting. This PR vectorizes the pair scoring itself. Only one of the two should merge.

Performance

structure_to_cells on synthetic tables:

Table Before After
20×6, 500 tokens 0.06 s <0.01 s
60×10, 3k tokens 1.84 s 0.03 s
120×12, 8k tokens 11.9 s 0.15 s

Behavior preserved

For Python int/float and NumPy float64/int64 coordinates the output is unchanged: the same assignments, scores and score order. That includes the edge-case behavior of Rect.intersect / Rect.get_area:

  • A zero-area container takes in every package with score 1.0 (Rect.intersect replaces a zero-area self with the other rect).
  • Packages with no positive area are skipped and get no score.
  • Ties go to the first container (argmax matches the stable descending sort).
  • Scores are returned as Python floats.

Known difference: NumPy float32 coordinates are now scored in float64 and returned as Python floats instead of np.float32. In 2,000 random float32 cases no assignment changed, but a score that lands exactly on overlap_threshold could in principle resolve differently.

Verification

  • Randomized differential check against the original implementation: 6,000 cases identical, including integer-grid boxes (forcing ties and zero-area/inverted boxes) and chunk sizes 1, 3, 7 and the default.
  • New tests in test_unstructured_inference/models/test_slot_into_containers.py:
    • a differential test against a pure-Rect reference across chunk sizes
    • explicit tests for zero-area containers, zero-area packages and tie-breaking (these also pass on the old implementation)
  • make test: 402 passed, 1 skipped. ruff check / ruff format --check clean.

Version bumped to 1.6.17 with a CHANGELOG entry.

View guided diff

slot_into_containers scored every token against every cell in pure
Python, building a Rect and a dict for each pair and sorting each
token's scores to find the best one. Compute the overlaps with numpy in
chunks of at most ~1M pairs instead, so temporary memory stays linear
in the cell count. structure_to_cells on a 120x12 table with 8,000
tokens drops from 11.9s to 0.15s.

Assignments, scores and score order match the Rect implementation for
Python and float64/int64 coordinates, including first-container tie
breaking, zero-area containers holding every package, and zero-area
packages going unscored. float32 coordinates are now scored in float64
and returned as Python floats.

Co-authored-by: wtfyd-agent[bot] <328989958+wtfyd-agent[bot]@users.noreply.github.com>
@KRRT7

KRRT7 commented Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Combined effect of #530, #531, #532 and #533 merged together onto main (9a227e7), measured against main on the same machine with alternating runs:

main all four
Peak RSS, layout-parser-paper.pdf (16 pages) through YoloX 1014–1015 MiB 666–674 MiB
RSS after importing inference.layout + models.base 369–417 MiB 148–149 MiB
structure_to_cells, synthetic 120×12 table, 8k tokens 28.3 s 0.33 s
PageLayout._get_image, page 200 of the 400-page sample 24.5 s 0.11 s
  • The 16-page YoloX output hash (coordinates, probabilities, class ids per page) is identical to main.
  • make test passes on the combined branch: 415 passed, 1 skipped.
  • The code changes merge without conflicts. All four PRs bump to 1.6.17, so each one merged after the first needs the next patch version.
  • Timings came from a busy machine and read about 2× high in absolute terms; the memory numbers were stable across runs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant