⚡️ Speed up function align_supercells by 177%

[codeflash-ai]

📄 177% (1.77x) speedup for align_supercells in unstructured_inference/models/table_postprocess.py

⏱️ Runtime : 7.48 milliseconds → 2.70 milliseconds (best of 88 runs)

📝 Explanation and details

The optimized code achieves a 177% speedup (7.48ms → 2.70ms) by eliminating expensive object allocations and method calls in hot loops through two key strategies:

What Was Optimized

1. Precomputed Row/Column Metrics: The original code repeatedly accessed row["bbox"] and col["bbox"] dictionaries inside nested loops (3,309 row iterations × 3,117 column iterations per supercell). The optimized version extracts these once into tuples (ymin, ymax, height, row_dict) and (xmin, xmax, width, col_dict), avoiding ~10,000+ dictionary lookups per function call.

2. Eliminated Rect Object Creation: The original code created 1,263 Rect objects and called their methods (955 include_rect calls, 106 intersect calls). The optimized version replaces this with direct numeric aggregation using simple min/max comparisons on primitive values, removing all object allocation overhead.

3. Replaced Conditional Expressions with Ternary Operations: Changed max(a, b) and min(a, b) calls to inline ternary expressions like a if a > b else b, which are faster in Python's bytecode execution.

Why This is Faster

• Memory Allocation: Creating objects in Python involves significant overhead (memory allocation, initialization, method dispatch). Eliminating 1,263 Rect allocations removes this entirely.
• Method Call Overhead: Each include_rect() call involved multiple attribute accesses and comparisons. Direct numeric operations on local variables are much faster.
• Cache Locality: Tuples of primitives are more cache-friendly than dictionary lookups and object attributes.
• Reduced Bytecode: Ternary expressions generate fewer bytecode operations than function calls.

Performance by Test Case

The optimization excels when:

• Many supercells with large row/column grids: The test_large_scale_multiple_supercells_efficiency_and_correctness shows 86.4% speedup (112μs → 60.5μs) because precomputation amortizes across all iterations
• Supercells spanning multiple rows/columns: Tests like test_single_supercell_single_row_single_column (41.7% faster) and test_supercell_spanning_two_rows (38.8% faster) benefit from eliminated Rect operations in the aggregation phase

The optimization is slightly slower (~6-13%) on trivial cases with no intersecting rows/columns (early exits), where the upfront precomputation cost exceeds savings. However, real workloads typically involve substantial overlap detection, making this trade-off worthwhile.

Impact on Workloads

This function processes table structure extraction results. If called frequently during document parsing (e.g., processing multiple tables per document), the 2.77× speedup directly reduces end-to-end latency. The optimization is most beneficial when:

• Documents contain complex tables with many rows/columns
• Multiple supercells need alignment per table
• High-throughput batch processing is required

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 50 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
from unstructured_inference.models.table_postprocess import align_supercells

# unit tests

# Helper utilities for creating regular grid rows and columns and bboxes
def make_rows(count, row_height=10, header_count=0, x_min=0, x_max=100):
    """
    Create `count` rows each with height row_height.
    The first `header_count` rows will have 'header': True.
    Each row's bbox is [x_min, y, x_max, y+row_height].
    """
    rows = []
    for i in range(count):
        bbox = [x_min, i * row_height, x_max, (i + 1) * row_height]
        row = {"bbox": bbox}
        if i < header_count:
            row["header"] = True
        rows.append(row)
    return rows

def make_columns(count, col_width=10, y_min=0, y_max=100):
    """
    Create `count` columns each with width col_width.
    Each column's bbox is [x, y_min, x+col_width, y_max].
    """
    columns = []
    for i in range(count):
        bbox = [i * col_width, y_min, (i + 1) * col_width, y_max]
        columns.append({"bbox": bbox})
    return columns

def test_basic_alignment_single_row_spanning_multiple_columns():
    # Basic scenario: one data row and multiple columns. A supercell spans two columns.
    rows = make_rows(count=1, row_height=10, header_count=0, x_min=0, x_max=60)
    # Three columns each 20 width across x range 0..60
    columns = make_columns(count=3, col_width=20, y_min=0, y_max=10)

    # Supercell spans first two columns exactly and matches row vertically.
    supercells = [
        {"bbox": [0, 0, 40, 10], "score": 0.9}  # overlays columns 0 and 1 fully and row 0 fully
    ]

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 20.0μs -> 12.8μs (56.0% faster)

    out = aligned[0]

def test_span_supercell_must_be_in_header_skipped_if_not_header():
    # Edge case: a 'span' supercell that intersects only data rows should be skipped.
    rows = make_rows(count=2, row_height=10, header_count=0, x_min=0, x_max=60)
    columns = make_columns(count=2, col_width=30, y_min=0, y_max=20)

    # Span supercell present but not in header area (rows are all data rows)
    supercells = [{"bbox": [0, 0, 60, 10], "score": 0.7, "span": True}]

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 7.05μs -> 7.97μs (11.5% slower)

def test_header_boundary_resolution_prefers_larger_group():
    # When a supercell overlaps both header and data rows, the smaller group is removed.
    # Setup: 1 header row and 2 data rows below. Supercell overlaps header row partly and both data rows.
    rows = make_rows(count=3, row_height=5, header_count=1, x_min=0, x_max=100)
    # Single wide column so the column condition won't block the result.
    columns = make_columns(count=1, col_width=100, y_min=0, y_max=20)

    # Create a bbox that overlaps header row (row 0) and both data rows (rows 1 and 2).
    # Overlap fractions are arranged so both data rows and header are >=0.5 overlap.
    supercells = [{"bbox": [0, 2, 100, 16], "score": 0.5}]

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 20.1μs -> 13.5μs (49.0% faster)
    out = aligned[0]

def test_header_span_propagation_creates_propagated_supercells():
    # A 'span' supercell in the header that spans multiple columns should produce propagated
    # supercells above it in the header (for rows above the supercell's rows).
    # Setup two header rows (indices 0 and 1). The supercell will cover row 1 only but span multiple columns.
    rows = make_rows(count=3, row_height=10, header_count=2, x_min=0, x_max=100)
    # Three columns; the supercell spans columns 0 and 1
    columns = make_columns(count=3, col_width=33, y_min=0, y_max=30)

    # Supercell sits in header region overlapping the second header row (index 1)
    supercells = [
        {
            "bbox": [0, 10, 66, 20],  # overlaps header row index 1 vertically
            "score": 0.95,
            "span": True,
        }
    ]

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 29.8μs -> 22.5μs (32.4% faster)

    # Find the propagated one: it has 'propagated': True
    propagated = [s for s in aligned if s.get("propagated")]

    p = propagated[0]
    # Propagated supercell should carry the same column_numbers as the span
    original = next(s for s in aligned if "propagated" not in s)

def test_no_overlap_returns_empty_list():
    # Edge case: supercell that does not meet 50% overlap in either rows or columns
    rows = make_rows(count=2, row_height=10, header_count=0, x_min=0, x_max=40)
    columns = make_columns(count=2, col_width=20, y_min=0, y_max=20)

    # Supercell is placed far to the right with minimal overlap with existing columns and rows.
    supercells = [{"bbox": [100, 100, 120, 120], "score": 0.1}]

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 6.84μs -> 7.32μs (6.45% slower)

def test_large_scale_multiple_supercells_efficiency_and_correctness():
    # Large scale-ish: generate a grid of rows and columns and multiple supercells.
    # This test keeps the dataset modest to avoid long runtimes but exercizes many iterations.
    num_rows = 20  # kept under 1000 to respect the instruction
    num_cols = 20
    rows = make_rows(count=num_rows, row_height=10, header_count=2, x_min=0, x_max=200)
    columns = make_columns(count=num_cols, col_width=10, y_min=0, y_max=200)

    supercells = []
    # Create 2 supercells that each span multiple adjacent columns but only one row,
    # which is a valid supercell provided they span multiple columns.
    # Supercell 0 spans columns 2..6 on row index 5
    sc0_bbox = [2 * 10, 5 * 10, 7 * 10, 6 * 10]  # covers columns 2-6 (5 columns) and row 5
    supercells.append({"bbox": sc0_bbox, "score": 0.8})

    # Supercell 1 spans columns 10..15 on header row index 1 (a header scenario if span provided)
    sc1_bbox = [10 * 10, 1 * 10, 16 * 10, 2 * 10]  # overlaps columns 10-15 and header row 1
    supercells.append({"bbox": sc1_bbox, "score": 0.9, "span": True})

    codeflash_output = align_supercells(supercells, rows, columns); aligned = codeflash_output # 112μs -> 60.5μs (86.4% faster)

    # The first supercell should align to multiple columns (so length >=1 for aligned)
    # Because sc0 spans multiple columns, it should appear in results
    matched_sc0 = [
        s for s in aligned if set(s.get("row_numbers", [])) == {5} or s.get("score") == 0.8
    ]

    # The span supercell sc1 is a header span; because it is in header (row 1 is header),
    # it should be present and may also generate propagated supercells. Ensure no exceptions
    # and at least one of the aligned supercells corresponds to the sc1 score.
    matched_sc1 = [s for s in aligned if s.get("score") == 0.9]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from unstructured_inference.models.table_postprocess import (Rect,
                                                             align_supercells)

class TestAlignSupercellsBasic:
    """Basic test cases for align_supercells function under normal conditions."""

    def test_empty_supercells_list(self):
        """Test that an empty supercells list returns an empty aligned list."""
        supercells = []
        rows = [{"bbox": [0, 0, 100, 50]}]
        columns = [{"bbox": [0, 0, 50, 50]}]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 959ns -> 2.51μs (61.8% slower)

    def test_single_supercell_single_row_single_column(self):
        """Test alignment of a single supercell with single row and column."""
        supercells = [{"bbox": [10, 10, 40, 40], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 50]}]
        columns = [{"bbox": [0, 0, 50, 100]}]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 14.7μs -> 10.4μs (41.7% faster)

    def test_supercell_spanning_two_rows(self):
        """Test supercell that intersects two rows at 50% threshold."""
        supercells = [{"bbox": [10, 25, 40, 75], "score": 0.9}]
        rows = [
            {"bbox": [0, 0, 100, 50]},
            {"bbox": [0, 50, 100, 100]}
        ]
        columns = [{"bbox": [0, 0, 100, 100]}]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 14.0μs -> 10.1μs (38.8% faster)

    def test_supercell_spanning_two_columns(self):
        """Test supercell that intersects two columns at 50% threshold."""
        supercells = [{"bbox": [25, 10, 75, 40], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 100]}]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 5.78μs -> 6.45μs (10.3% slower)

    def test_supercell_with_header_row(self):
        """Test supercell alignment with header row marked."""
        supercells = [{"bbox": [10, 5, 40, 35], "score": 0.9}]
        rows = [
            {"bbox": [0, 0, 100, 50], "header": True},
            {"bbox": [0, 50, 100, 100]}
        ]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 16.9μs -> 12.5μs (35.2% faster)

    def test_supercell_below_50_percent_overlap_excluded(self):
        """Test that supercells with less than 50% overlap are excluded."""
        supercells = [{"bbox": [10, 10, 40, 30], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 50]}]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 5.69μs -> 6.47μs (12.1% slower)

    def test_multiple_supercells_mixed_results(self):
        """Test multiple supercells with varied intersection patterns."""
        supercells = [
            {"bbox": [10, 10, 40, 60], "score": 0.9},  # Spans 2 rows
            {"bbox": [60, 10, 90, 40], "score": 0.8},  # Doesn't span
        ]
        rows = [
            {"bbox": [0, 0, 100, 50]},
            {"bbox": [0, 50, 100, 100]}
        ]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 26.7μs -> 17.7μs (50.8% faster)

    def test_supercell_with_span_property_header_required(self):
        """Test that supercells with 'span' property must be in header."""
        supercells = [{"bbox": [10, 60, 40, 90], "score": 0.9, "span": True}]
        rows = [
            {"bbox": [0, 0, 100, 50], "header": True},
            {"bbox": [0, 50, 100, 100]}
        ]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 7.34μs -> 8.44μs (13.0% slower)

    def test_supercell_no_row_intersection(self):
        """Test that supercells with no row intersection are excluded."""
        supercells = [{"bbox": [10, 200, 40, 250], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 100]}]
        columns = [{"bbox": [0, 0, 100, 100]}]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 5.85μs -> 6.20μs (5.58% slower)

    def test_supercell_no_column_intersection(self):
        """Test that supercells with no column intersection are excluded."""
        supercells = [{"bbox": [200, 10, 250, 40], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 100]}]
        columns = [{"bbox": [0, 0, 100, 100]}]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 5.84μs -> 6.07μs (3.66% slower)

class TestAlignSupercellsEdgeCases:
    """Edge case tests for align_supercells function."""

    def test_supercell_exact_50_percent_overlap(self):
        """Test supercell with exactly 50% overlap on height."""
        # Supercell height is 50, row height is 100, overlap is 50 = 50%
        supercells = [{"bbox": [10, 25, 40, 75], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 100]}]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 15.8μs -> 11.4μs (38.9% faster)

    def test_supercell_just_below_50_percent(self):
        """Test supercell with just below 50% overlap."""
        # Carefully construct to get 49.9% overlap
        supercells = [{"bbox": [10, 25, 40, 74.9], "score": 0.9}]
        rows = [{"bbox": [0, 0, 100, 100]}]
        columns = [
            {"bbox": [0, 0, 50, 100]},
            {"bbox": [50, 0, 100, 100]}
        ]
        codeflash_output = align_supercells(supercells, rows, columns); result = codeflash_output # 6.09μs -> 6.95μs (12.4% slower)

    

To edit these changes git checkout codeflash/optimize-align_supercells-mkou96v9 and push.