⚡️ Speed up function align_columns by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for align_columns in unstructured_inference/models/table_postprocess.py

⏱️ Runtime : 295 microseconds → 260 microseconds (best of 250 runs)

📝 Explanation and details

The optimization achieves a 13% speedup by reducing redundant dictionary lookups through caching the column["bbox"] reference in a local variable.

Key Optimization

Dictionary Lookup Caching: Instead of accessing column["bbox"] twice per iteration, the optimized code stores it once in column_bbox:

# Original: 2 dict lookups per column
column["bbox"][1] = bbox[1]  # lookup 1
column["bbox"][3] = bbox[3]  # lookup 2

# Optimized: 1 dict lookup per column
column_bbox = column["bbox"]  # single lookup
column_bbox[1] = bbox[1]      # direct list access
column_bbox[3] = bbox[3]      # direct list access

Why This Works

In Python, dictionary lookups (__getitem__) have overhead compared to direct variable access. By caching the bbox list reference:

• We eliminate one __getitem__ call per column (50% reduction in dict lookups)
• List item assignment via a local variable is faster than going through the dictionary each time
• The reference points to the same mutable list object, so modifications are still in-place

Performance Evidence

Line profiler shows the optimization is most effective with many columns:

• Single column tests: 1-4% improvement (overhead of extra assignment is minimal)
• Large-scale tests (500 columns): 13-23% improvement
  • test_large_scale_alignment_performance_and_correctness: 67.3μs → 59.1μs (13.8% faster)
  • test_align_columns_performance_many_columns: 62.4μs → 50.9μs (22.7% faster)

The speedup scales with the number of columns because the saved dictionary lookups accumulate across iterations.

Impact on Error Handling

The optimization preserves the original exception behavior - when errors occur (malformed bbox, missing keys, etc.), the exception is caught at the same point in execution, ensuring partial modifications behave identically in both versions.

✅ Correctness verification report:

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

def test_basic_single_column_alignment():
    # Basic: a single column with a mutable bbox list should have its y-top (index 1)
    # and y-bottom (index 3) updated to match the provided bbox.
    col = {"bbox": [0, 10, 100, 20]}  # initial bbox: left, top, right, bottom
    columns = [col]
    target_bbox = [5, 1, 105, 99]  # want top=1, bottom=99

    codeflash_output = align_columns(columns, target_bbox); result = codeflash_output # 1.22μs -> 1.17μs (4.36% faster)

def test_multiple_columns_alignment_all_updated():
    # Basic: multiple columns should all get their top and bottom overwritten.
    columns = [
        {"bbox": [0, 10, 50, 20]},
        {"bbox": [50, 11, 100, 21]},
        {"bbox": [100, 12, 150, 22]},
    ]
    target_bbox = [0, 2, 200, 80]

    codeflash_output = align_columns(columns, target_bbox); returned = codeflash_output # 1.52μs -> 1.48μs (2.78% faster)

    # all columns updated
    for c in returned:
        pass

def test_inplace_modification_and_reference_sharing():
    # If multiple columns share the same bbox list object, updating one entry
    # will be visible across all references (function mutates in place).
    shared_bbox = [0, 10, 100, 20]
    columns = [
        {"bbox": shared_bbox},
        {"bbox": shared_bbox},  # same object reference
    ]
    target_bbox = [0, 0, 100, 999]

    align_columns(columns, target_bbox) # 1.33μs -> 1.34μs (1.27% slower)

def test_tuple_column_bbox_causes_graceful_handling_and_no_raise():
    # Edge: If a column's bbox is an immutable tuple, assignment will raise a
    # TypeError. The function should catch the exception and not propagate it.
    col = {"bbox": (0, 10, 100, 20)}  # tuple -> immutable, assignments will fail
    columns = [col]
    target_bbox = [0, 1, 100, 99]

    # Should not raise despite the internal TypeError.
    codeflash_output = align_columns(columns, target_bbox); returned = codeflash_output # 6.07μs -> 5.67μs (6.95% faster)

def test_missing_bbox_key_is_handled_gracefully():
    # Edge: If a column dict lacks the "bbox" key, attempting to access it will
    # raise a KeyError which should be caught. The function should not raise.
    columns = [{"not_bbox": [1, 2, 3, 4]}, {"bbox": [0, 10, 50, 20]}]
    target_bbox = [0, 0, 100, 100]

    # Should not raise even though first dict misses "bbox".
    codeflash_output = align_columns(columns, target_bbox); result = codeflash_output # 4.93μs -> 4.70μs (5.04% faster)

def test_short_target_bbox_list_handled_without_exception():
    # Edge: If the provided bbox (target) is too short (missing index 1 or 3),
    # an IndexError occurs and must be caught; function should not raise.
    columns = [{"bbox": [0, 10, 50, 20]}]
    short_target = [999]  # no index 1 or 3

    # Ensure no exception escapes.
    codeflash_output = align_columns(columns, short_target); result = codeflash_output # 4.88μs -> 4.75μs (2.70% faster)

def test_non_dict_element_in_columns_handled():
    # Edge: If elements in columns are not dicts (e.g., integers), attempting to
    # do column["bbox"] will raise a TypeError. The function should catch it.
    columns = [{"bbox": [0, 10, 50, 20]}, 123, {"bbox": [10, 11, 12, 13]}]
    target_bbox = [0, 2, 100, 88]

    # Should not raise even though one element is invalid.
    codeflash_output = align_columns(columns, target_bbox); returned = codeflash_output # 5.66μs -> 5.45μs (3.91% faster)

def test_empty_columns_returns_empty_list():
    # Edge: Empty list should be handled and returned as is (no errors).
    columns = []
    target_bbox = [0, 1, 2, 3]

    codeflash_output = align_columns(columns, target_bbox); result = codeflash_output # 748ns -> 706ns (5.95% faster)

def test_partial_modification_on_exception_stops_processing():
    # Edge: Verify that if an exception occurs in the middle of the loop,
    # previous columns have been modified but subsequent ones are untouched.
    # First column OK, second column has tuple (immutable) causing TypeError,
    # third column should remain as original because processing halts on exception.
    columns = [
        {"bbox": [0, 10, 50, 20]},  # will be modified before exception
        {"bbox": (0, 11, 50, 21)},  # tuple -> causes TypeError when assigning
        {"bbox": [0, 12, 50, 22]},  # should remain unchanged
    ]
    target_bbox = [0, 2, 100, 99]

    # Execute; should not raise
    codeflash_output = align_columns(columns, target_bbox); result = codeflash_output # 5.91μs -> 5.84μs (1.13% faster)

def test_large_scale_alignment_performance_and_correctness():
    # Large Scale: create a large number (but under 1000) of columns to test
    # scalability and correctness. We use 500 columns.
    n = 500
    columns = [{"bbox": [i, i + 10, i + 100, i + 20]} for i in range(n)]
    target_bbox = [0, 1, 9999, 77]

    codeflash_output = align_columns(columns, target_bbox); returned = codeflash_output # 67.3μs -> 59.1μs (13.8% faster)
    # Verify that every column's top and bottom are updated correctly
    for idx, c in enumerate(returned):
        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_columns

def test_align_columns_single_column():
    """Test aligning a single column with a simple bbox."""
    columns = [{"bbox": [0, 0, 100, 100], "text": "Header"}]
    bbox = [0, 10, 500, 510]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.16μs -> 1.15μs (1.04% faster)

def test_align_columns_multiple_columns():
    """Test aligning multiple columns with different original boundaries."""
    columns = [
        {"bbox": [0, 5, 100, 95], "text": "Col1"},
        {"bbox": [100, 15, 200, 85], "text": "Col2"},
        {"bbox": [200, 25, 300, 75], "text": "Col3"}
    ]
    bbox = [0, 50, 300, 450]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.48μs -> 1.49μs (0.670% slower)
    
    # Verify all columns are aligned to the same top and bottom
    for i, column in enumerate(result):
        pass

def test_align_columns_empty_list():
    """Test aligning an empty list of columns."""
    columns = []
    bbox = [0, 10, 500, 510]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 675ns -> 706ns (4.39% slower)

def test_align_columns_returns_same_object():
    """Test that the function returns the same list object (modified in-place)."""
    columns = [{"bbox": [0, 0, 100, 100], "text": "Header"}]
    bbox = [0, 10, 500, 510]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.15μs -> 1.20μs (3.68% slower)

def test_align_columns_preserves_other_attributes():
    """Test that other column attributes are preserved during alignment."""
    columns = [
        {
            "bbox": [0, 0, 100, 100],
            "text": "Header",
            "confidence": 0.95,
            "id": 1,
            "color": "blue"
        }
    ]
    bbox = [0, 10, 500, 510]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.14μs -> 1.16μs (1.89% slower)

def test_align_columns_zero_boundaries():
    """Test aligning columns with zero-valued boundaries."""
    columns = [{"bbox": [0, 0, 0, 0]}]
    bbox = [0, 0, 0, 0]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.19μs -> 1.14μs (4.11% faster)

def test_align_columns_negative_coordinates():
    """Test aligning columns with negative coordinate values."""
    columns = [{"bbox": [-50, -100, 50, 100]}]
    bbox = [-100, -200, 100, 200]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.15μs -> 1.15μs (0.437% faster)

def test_align_columns_very_large_coordinates():
    """Test aligning columns with very large coordinate values."""
    columns = [{"bbox": [1000000, 1000000, 2000000, 2000000]}]
    bbox = [0, 500000, 3000000, 2500000]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.17μs -> 1.20μs (2.42% slower)

def test_align_columns_float_coordinates():
    """Test aligning columns with floating-point coordinate values."""
    columns = [{"bbox": [10.5, 20.3, 100.7, 150.2]}]
    bbox = [0, 15.5, 500, 505.8]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.20μs -> 1.15μs (4.52% faster)

def test_align_columns_inverted_bbox_coordinates():
    """Test aligning columns when bbox has inverted coordinates (top > bottom)."""
    columns = [{"bbox": [0, 100, 100, 200]}]
    bbox = [0, 300, 500, 50]  # top (300) > bottom (50)
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.18μs -> 1.20μs (1.76% slower)

def test_align_columns_single_column_in_large_bbox():
    """Test aligning a single small column within a much larger bbox."""
    columns = [{"bbox": [10, 10, 20, 20]}]
    bbox = [0, 0, 1000, 1000]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.18μs -> 1.20μs (1.59% slower)

def test_align_columns_column_larger_than_bbox():
    """Test aligning columns that are larger than the bbox itself."""
    columns = [{"bbox": [0, 0, 1000, 1000]}]
    bbox = [100, 200, 300, 400]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.08μs -> 1.20μs (10.3% slower)

def test_align_columns_bbox_with_zero_height():
    """Test aligning columns with a bbox that has zero height (top == bottom)."""
    columns = [{"bbox": [0, 0, 100, 100]}]
    bbox = [0, 50, 500, 50]  # bbox[1] == bbox[3] (zero height)
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.16μs -> 1.16μs (0.692% faster)

def test_align_columns_many_columns_same_bbox():
    """Test aligning many columns all to the same bbox."""
    columns = [
        {"bbox": [i * 50, i, i * 50 + 40, 100 - i], "id": i}
        for i in range(20)
    ]
    bbox = [0, 10, 1000, 990]
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 3.60μs -> 3.26μs (10.6% faster)
    
    # All columns should have same top and bottom
    for i, column in enumerate(result):
        pass

def test_align_columns_with_minimal_bbox():
    """Test aligning columns to a minimal bbox (single point essentially)."""
    columns = [
        {"bbox": [100, 200, 300, 400]},
        {"bbox": [500, 600, 700, 800]}
    ]
    bbox = [0, 100, 1000, 100]  # Zero height
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.40μs -> 1.37μs (1.82% faster)

def test_align_columns_with_malformed_bbox():
    """Test that function handles malformed bbox gracefully."""
    columns = [{"bbox": [0, 0, 100, 100]}]
    bbox = [0, 10, 500]  # Missing bbox[3]
    
    # Function should catch the error and return columns
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 5.27μs -> 4.79μs (10.0% faster)

def test_align_columns_with_missing_bbox_key():
    """Test that function handles missing 'bbox' key in column."""
    columns = [{"text": "Header", "id": 1}]  # No 'bbox' key
    bbox = [0, 10, 500, 510]
    
    # Function should catch the error and return columns
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 4.91μs -> 4.71μs (4.29% faster)

def test_align_columns_with_none_bbox():
    """Test that function handles None values in bbox."""
    columns = [{"bbox": None}]
    bbox = [0, 10, 500, 510]
    
    # Function should catch the error and return columns
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 5.62μs -> 5.20μs (8.05% faster)

def test_align_columns_with_non_numeric_coordinates():
    """Test that function handles non-numeric coordinates."""
    columns = [{"bbox": ["a", "b", "c", "d"]}]
    bbox = [0, 10, 500, 510]
    
    # Function should catch the error and return columns
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.20μs -> 1.20μs (0.418% faster)

def test_align_columns_with_string_bbox_instead_of_list():
    """Test that function handles string bbox instead of list."""
    columns = [{"bbox": [0, 0, 100, 100]}]
    bbox = "not a bbox"
    
    # Function should catch the error and return columns
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 1.33μs -> 1.22μs (9.08% faster)

def test_align_columns_performance_many_columns():
    """Test aligning a large number of columns for performance."""
    # Create 500 columns with varying properties
    columns = [
        {
            "bbox": [i * 2, i % 100, i * 2 + 50, (i % 100) + 100],
            "id": i,
            "text": f"Column_{i}",
            "confidence": 0.5 + (i % 50) / 100
        }
        for i in range(500)
    ]
    bbox = [0, 250, 2000, 2750]
    
    # Perform alignment
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 62.4μs -> 50.9μs (22.7% faster)
    
    # Verify all columns have correct alignment
    for i, column in enumerate(result):
        pass

def test_align_columns_with_complex_column_structures():
    """Test aligning columns with complex nested structures."""
    columns = [
        {
            "bbox": [i * 10, i, i * 10 + 50, i + 100],
            "metadata": {
                "nested_id": i,
                "nested_list": [1, 2, 3],
                "nested_dict": {"a": 1, "b": 2}
            },
            "extra_field": "preserved",
            "data": [0, 100, 200, 300, 400]
        }
        for i in range(100)
    ]
    bbox = [0, 50, 1500, 1050]
    
    # Perform alignment
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 13.5μs -> 11.3μs (19.4% faster)
    
    # Verify all columns preserved their complex structures
    for i, column in enumerate(result):
        pass

def test_align_columns_preserves_list_identity_large_scale():
    """Test that with many columns, the function still returns the same list object."""
    columns = [
        {"bbox": [i * 5, i, i * 5 + 20, i + 50]}
        for i in range(300)
    ]
    bbox = [0, 0, 2000, 2000]
    
    # Store original list id
    original_id = id(columns)
    
    # Perform alignment
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 38.2μs -> 33.0μs (15.9% faster)

def test_align_columns_with_extreme_coordinate_ranges():
    """Test aligning columns with very wide ranges of coordinates."""
    columns = [
        {
            "bbox": [i * 1000, i * 100, i * 1000 + 500, i * 100 + 500],
            "id": i
        }
        for i in range(100)
    ]
    bbox = [0, 5000, 100000, 15000]
    
    # Perform alignment
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 15.0μs -> 14.2μs (5.08% faster)
    
    # Verify alignment across all columns
    for i, column in enumerate(result):
        pass

def test_align_columns_mixed_coordinate_types_large():
    """Test aligning 200 columns with mixed int and float coordinates."""
    columns = []
    for i in range(200):
        if i % 2 == 0:
            # Integer coordinates
            columns.append({"bbox": [i * 10, i, i * 10 + 50, i + 100]})
        else:
            # Float coordinates
            columns.append({"bbox": [i * 10.5, i * 1.5, i * 10.5 + 50.5, i * 1.5 + 100.5]})
    
    bbox = [0, 99.5, 2500, 299.8]
    
    # Perform alignment
    codeflash_output = align_columns(columns, bbox); result = codeflash_output # 27.7μs -> 23.1μs (19.9% faster)
    
    # Verify all columns aligned correctly
    for i, column in enumerate(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_columns-mkotkdqk and push.