⚡️ Speed up function align_rows by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for align_rows in unstructured_inference/models/table_postprocess.py

⏱️ Runtime : 356 microseconds → 330 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 7% speedup by reducing redundant dictionary lookups in the hot loop.

Key optimization:
Instead of accessing row["bbox"] twice per iteration (once for index 0, once for index 2), the optimized version stores the reference in a local variable rb = row["bbox"] and reuses it. This eliminates one dictionary lookup per iteration.

Why this works:
In Python, dictionary lookups (row["bbox"]) involve hash table operations with computational overhead. By caching the bbox reference in a local variable, subsequent array index assignments (rb[0] and rb[2]) operate directly on the list object without repeated dictionary access. Local variable access is significantly faster than dictionary key lookups.

Performance impact by test case:

• Best gains (10-14% faster): Tests with small to medium row counts (single row, 100-500 rows) where the loop dominates runtime. Examples include test_align_rows_negative_coordinates (14% faster), test_align_rows_bbox_wrong_type (13.3% faster), and test_large_scale_many_rows_all_adjusted (11.3% faster with 500 rows).
• Moderate gains (7-10% faster): Most basic alignment tests show consistent 7-10% improvements, indicating the optimization applies broadly.
• Minimal/negative impact: Error-handling test cases where exceptions are raised early (e.g., immutable tuples, missing keys) see no benefit or slight slowdowns (4-8% slower) since the optimization's value is in repeated successful iterations, not exception paths.

Impact on workloads:
If this function is called frequently in table post-processing pipelines (common in document analysis), the 7% improvement compounds across thousands of table extractions. The optimization is most valuable when processing tables with many rows (100+), where the cumulative effect of eliminating dictionary lookups per row becomes significant.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 39 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_rows

def test_basic_alignment_single_row():
    # Basic: single row with mutable bbox list should be updated in-place.
    rows = [{"bbox": [0, 10, 50, 20]}]  # original left=0, right=50
    final_bbox = [5, 0, 95, 100]  # we expect left->5, right->95
    codeflash_output = align_rows(rows, final_bbox); returned = codeflash_output # 1.20μs -> 1.16μs (3.28% faster)

def test_basic_alignment_multiple_rows():
    # Basic: multiple rows with mutable bbox lists all get updated.
    rows = [
        {"bbox": [0, 0, 10, 10]},
        {"bbox": [1, 1, 11, 11]},
        {"bbox": [2, 2, 12, 12]},
    ]
    final_bbox = [100, 0, 200, 0]
    codeflash_output = align_rows(rows, final_bbox); out = codeflash_output # 1.47μs -> 1.46μs (0.549% faster)

    # Every row should have left and right aligned to final_bbox
    for row in rows:
        pass

def test_empty_rows_returns_empty_list():
    # Edge: empty rows should be handled gracefully and returned as empty list.
    rows = []
    final_bbox = [0, 0, 10, 10]
    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 663ns -> 640ns (3.59% faster)

def test_inplace_modification_returns_same_object():
    # Contract: function returns the same list object (modifying in-place).
    rows = [{"bbox": [0, 0, 0, 0]} for _ in range(3)]
    final_bbox = [7, 0, 13, 0]
    before_id = id(rows)
    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 1.52μs -> 1.46μs (4.19% faster)
    after_id = id(result)
    # All modified as expected
    for r in rows:
        pass

def test_immutable_row_bbox_causes_no_change_and_returns_original_rows():
    # Edge: if a row's bbox is an immutable tuple, assignment to indexes will raise TypeError.
    # The function catches exceptions and returns the original rows (unchanged).
    rows = [{"bbox": (0, 0, 10, 10)}]  # tuple is immutable
    original_copy = rows.copy()  # shallow copy for comparison (tuples are immutable so safe)
    final_bbox = [5, 0, 95, 100]

    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 5.74μs -> 6.01μs (4.52% slower)

def test_partial_modification_on_mixed_mutable_and_immutable_bboxes():
    # Edge: when some rows are mutable and a subsequent row causes exception,
    # changes applied before the exception should remain.
    # Prepare first row mutable (list), second row immutable (tuple) -> second will cause TypeError.
    rows = [
        {"bbox": [0, 0, 10, 10]},  # will be modified successfully
        {"bbox": (1, 1, 11, 11)},  # assignment to tuple will raise TypeError
        {"bbox": [2, 2, 12, 12]},  # should remain untouched because exception stops the loop
    ]
    final_bbox = [50, 0, 150, 0]

    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 5.81μs -> 6.24μs (6.78% slower)

def test_missing_bbox_key_triggers_exception_and_returns_unchanged():
    # Edge: if a row does not have 'bbox' key at all, a KeyError occurs and it's caught.
    rows = [{"no_bbox": [0, 0, 0, 0]}]
    # Keep a deep-ish copy for comparison (shallow copy is fine since inner structures are lists)
    original = [{"no_bbox": [0, 0, 0, 0]}]
    final_bbox = [1, 2, 3, 4]

    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 4.85μs -> 4.83μs (0.518% faster)

def test_short_bbox_list_partial_modification_when_index_error():
    # Edge: if a row bbox is a list but too short (e.g., length 1), setting index 0 can succeed
    # but setting index 2 will raise IndexError; the function should leave whatever was changed
    # before the error.
    rows = [{"bbox": [7]}, {"bbox": [1, 2, 3, 4]}]
    final_bbox = [100, 0, 200, 0]

    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 4.84μs -> 5.30μs (8.75% slower)

def test_invalid_bbox_parameter_type_raises_and_returns_unchanged():
    # Edge: if bbox parameter is not subscriptable (e.g., None), accessing bbox[0] will raise TypeError.
    # The function should catch the exception and return the rows unchanged.
    rows = [{"bbox": [0, 0, 10, 10]}]
    original = [{"bbox": [0, 0, 10, 10]}]
    final_bbox = None  # invalid type for indexing

    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 5.02μs -> 5.36μs (6.34% slower)

def test_large_scale_many_rows_all_adjusted():
    # Large Scale: create a large but bounded number of rows (500) to test scalability.
    # Each row has a mutable bbox list; after align_rows all should reflect new left/right.
    N = 500  # within the allowed limit (<1000)
    rows = [{"bbox": [i, i + 1, i + 2, i + 3]} for i in range(N)]
    final_bbox = [999, 0, 1001, 0]

    # Run the function
    codeflash_output = align_rows(rows, final_bbox); result = codeflash_output # 59.4μs -> 53.4μs (11.3% faster)

    # Verify every row was updated correctly.
    for idx, row in enumerate(rows):
        pass
# 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 align_rows

def test_align_rows_single_row_basic():
    """Test aligning a single row with basic numeric values."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.18μs -> 1.09μs (8.24% faster)

def test_align_rows_multiple_rows_basic():
    """Test aligning multiple rows with basic numeric values."""
    rows = [
        {"bbox": [10, 20, 30, 40]},
        {"bbox": [15, 25, 35, 45]},
        {"bbox": [12, 22, 32, 42]}
    ]
    bbox = [5, 50, 40, 60]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.53μs -> 1.46μs (4.45% faster)
    for row in result:
        pass

def test_align_rows_preserves_middle_values():
    """Test that align_rows preserves the middle y-coordinate values (indices 1 and 3)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [5, 100, 35, 200]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.20μs -> 1.07μs (12.6% faster)

def test_align_rows_with_zero_bbox():
    """Test align_rows with bbox values of zero."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [0, 20, 0, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.17μs -> 1.10μs (5.90% faster)

def test_align_rows_with_negative_bbox():
    """Test align_rows with negative bbox values."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [-5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.16μs -> 1.06μs (9.23% faster)

def test_align_rows_with_float_values():
    """Test align_rows with floating point bbox values."""
    rows = [{"bbox": [10.5, 20.3, 30.7, 40.1]}]
    bbox = [5.2, 20.3, 35.8, 40.1]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.18μs -> 1.10μs (6.89% faster)

def test_align_rows_returns_same_rows_object():
    """Test that align_rows modifies and returns the same rows list."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.19μs -> 1.10μs (7.71% faster)

def test_align_rows_with_additional_keys_in_row():
    """Test that align_rows preserves other keys in row dictionaries."""
    rows = [{"bbox": [10, 20, 30, 40], "text": "hello", "id": 1}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.21μs -> 1.16μs (4.24% faster)

def test_align_rows_empty_list():
    """Test align_rows with an empty rows list."""
    rows = []
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 695ns -> 639ns (8.76% faster)

def test_align_rows_very_large_bbox_values():
    """Test align_rows with very large numeric values."""
    rows = [{"bbox": [100, 200, 300, 400]}]
    bbox = [1000000, 200, 2000000, 400]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.23μs -> 1.11μs (11.3% faster)

def test_align_rows_bbox_left_equals_right():
    """Test align_rows where bbox left and right coordinates are equal."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [15, 20, 15, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.20μs -> 1.09μs (9.89% faster)

def test_align_rows_row_bbox_left_equals_right():
    """Test align_rows with row bbox left and right coordinates equal."""
    rows = [{"bbox": [25, 20, 25, 40]}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.19μs -> 1.12μs (6.52% faster)

def test_align_rows_single_row_empty_bbox_list():
    """Test align_rows with bbox as empty list (edge case for bbox structure)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = []
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 4.47μs -> 4.90μs (8.77% slower)

def test_align_rows_row_missing_bbox_key():
    """Test align_rows when a row is missing the 'bbox' key (error handling)."""
    rows = [{"text": "hello"}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 5.11μs -> 4.87μs (4.85% faster)

def test_align_rows_row_bbox_wrong_type():
    """Test align_rows when bbox is not a list/indexable (error handling)."""
    rows = [{"bbox": "not_a_list"}]
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 5.52μs -> 5.73μs (3.70% slower)

def test_align_rows_bbox_wrong_type():
    """Test align_rows when bbox parameter is not indexable (error handling)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = "not_a_bbox"
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.38μs -> 1.22μs (13.3% faster)

def test_align_rows_bbox_too_short():
    """Test align_rows when bbox has fewer than 3 elements (error handling)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [5, 20]  # Only 2 elements, needs at least 3
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 4.60μs -> 4.57μs (0.679% faster)

def test_align_rows_row_bbox_too_short():
    """Test align_rows when row bbox has fewer than 3 elements (error handling)."""
    rows = [{"bbox": [10, 20]}]  # Only 2 elements, needs at least 3
    bbox = [5, 20, 35, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 4.98μs -> 4.94μs (0.911% faster)

def test_align_rows_with_none_values():
    """Test align_rows when bbox contains None values (error handling)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [None, 20, None, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.22μs -> 1.12μs (8.54% faster)

def test_align_rows_mixed_row_types():
    """Test align_rows with rows that have different structures."""
    rows = [
        {"bbox": [10, 20, 30, 40], "id": 1},
        {"bbox": [15, 25, 35, 45], "id": 2, "extra": "data"},
        {"bbox": [12, 22, 32, 42]}
    ]
    bbox = [5, 50, 40, 60]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.52μs -> 1.51μs (0.862% faster)
    for i, row in enumerate(result):
        pass

def test_align_rows_with_scientific_notation():
    """Test align_rows with scientific notation numbers."""
    rows = [{"bbox": [1e1, 2e1, 3e1, 4e1]}]
    bbox = [5e0, 2e1, 3.5e1, 4e1]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.20μs -> 1.09μs (10.3% faster)

def test_align_rows_bbox_index_out_of_range():
    """Test align_rows when bbox index access would be out of range (error handling)."""
    rows = [{"bbox": [10, 20, 30, 40]}]
    bbox = [5]  # Only index 0 exists
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 4.69μs -> 4.58μs (2.45% faster)

def test_align_rows_negative_coordinates():
    """Test align_rows with all negative coordinates."""
    rows = [{"bbox": [-50, -100, -30, -40]}]
    bbox = [-60, -100, -20, -40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 1.22μs -> 1.07μs (14.0% faster)

def test_align_rows_large_number_of_rows():
    """Test align_rows with a large number of rows (500 rows)."""
    # Create 500 rows with different bbox values
    rows = [
        {"bbox": [10 + i, 20 + i, 30 + i, 40 + i], "row_id": i}
        for i in range(500)
    ]
    bbox = [5, 1000, 1005, 2000]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 60.3μs -> 54.5μs (10.5% faster)

def test_align_rows_large_bbox_values_performance():
    """Test align_rows with large numeric bbox values on 200 rows."""
    rows = [
        {"bbox": [1000000 + i*100, 2000000, 3000000 + i*100, 4000000]}
        for i in range(200)
    ]
    bbox = [999999, 2000000, 3000001, 4000000]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 25.4μs -> 23.6μs (7.77% faster)
    for row in result:
        pass

def test_align_rows_high_precision_floats():
    """Test align_rows with high precision floating point values on 100 rows."""
    rows = [
        {"bbox": [10.123456789 + i*0.1, 20.987654321, 30.555555555 + i*0.1, 40.111111111]}
        for i in range(100)
    ]
    bbox = [5.999999999, 20.987654321, 35.888888888, 40.111111111]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 13.2μs -> 12.1μs (9.48% faster)
    for row in result:
        pass

def test_align_rows_many_additional_keys():
    """Test align_rows on rows with many additional keys (400 rows with 10 extra keys each)."""
    rows = [
        {
            "bbox": [10 + i, 20, 30 + i, 40],
            "key1": f"value1_{i}",
            "key2": f"value2_{i}",
            "key3": f"value3_{i}",
            "key4": f"value4_{i}",
            "key5": f"value5_{i}",
            "key6": f"value6_{i}",
            "key7": f"value7_{i}",
            "key8": f"value8_{i}",
            "key9": f"value9_{i}",
            "key10": f"value10_{i}"
        }
        for i in range(400)
    ]
    bbox = [5, 20, 1005, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 49.3μs -> 44.6μs (10.6% faster)

def test_align_rows_alternating_structure():
    """Test align_rows with 300 rows alternating between different structures."""
    rows = []
    for i in range(300):
        if i % 2 == 0:
            rows.append({"bbox": [10 + i, 20, 30 + i, 40], "text": f"row_{i}"})
        else:
            rows.append({"bbox": [10 + i, 20, 30 + i, 40], "text": f"row_{i}", "extra": i})
    
    bbox = [5, 20, 1005, 40]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 35.4μs -> 31.7μs (11.9% faster)
    for i, row in enumerate(result):
        pass

def test_align_rows_extreme_value_ranges():
    """Test align_rows with extreme value ranges across 250 rows."""
    rows = [
        {"bbox": [-1e6 + i, -1e5, 1e6 - i, 1e5]}
        for i in range(250)
    ]
    bbox = [-2e6, -1e5, 2e6, 1e5]
    codeflash_output = align_rows(rows, bbox); result = codeflash_output # 31.3μs -> 29.1μs (7.37% faster)
    for row in result:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-align_rows-mkotouia and push.