⚡️ Speed up function slot_into_containers by 631%

[codeflash-ai]

📄 631% (6.31x) speedup for slot_into_containers in unstructured_inference/models/table_postprocess.py

⏱️ Runtime : 475 milliseconds → 65.1 milliseconds (best of 53 runs)

📝 Explanation and details

The optimized code achieves a 630% speedup (475ms → 65.1ms) by eliminating expensive object allocations and redundant operations in the hot path of slot_into_containers.

Key Optimizations

1. Eliminated Rect Object Creation (Major Impact)

The original code created 3 Rect objects per package-container pair:

• One for the package (Rect(package["bbox"]))
• One for the container (Rect(container["bbox"]))
• One temporary for intersection calculation (Rect(package["bbox"]) again)

With 171,432 container-package pairs evaluated, this meant 514,296 object allocations. The line profiler shows Rect.__init__ taking 745ms and Rect.intersect() taking 1.89s—together accounting for ~35% of total runtime.

The optimized version uses direct tuple unpacking and arithmetic, avoiding all these allocations.

2. Removed Intermediate List Building and Sorting

The original code built a match_scores list for each package containing all container candidates, then sorted it to find the best match. With 1,378 packages and ~124 containers each, this meant:

• 171,430 dictionary allocations for match score entries
• 1,376 calls to sort_objects_by_score() (78.7ms total)

The optimized version uses a single-pass search to track the best score directly, eliminating both list construction and sorting overhead.

3. Pre-extracted Container Bboxes

By extracting container["bbox"] once upfront into container_bboxes, the optimization avoids 171,432 dictionary lookups in the inner loop—a simple change that reduces overhead from repeated dict access.

4. Inlined Area Calculations

Instead of calling get_area() multiple times (which involves method call overhead), areas are computed inline using direct arithmetic: (x_max - x_min) * (y_max - y_min).

5. Preserved Zero-Area Container Logic

The optimization correctly handles the edge case where containers with zero area should match fully with packages (score 1.0), maintaining behavioral compatibility with the original Rect.intersect() logic.

Test Results Analysis

The speedup is most dramatic for:

• Large-scale scenarios (test_large_scale_many_containers_many_packages: 666% faster) where the N×M loop dominates
• Medium workloads (test_multiple_packages_and_containers_large_scale: 567% faster)
• Even small inputs benefit (test_single_container_full_overlap: 147% faster) due to eliminated overhead

Impact Assessment

Based on function_references, slot_into_containers is called from nms_by_containment, which performs non-maxima suppression for table detection. This is likely in a hot path for document processing pipelines that extract tables from images. The 6-7x speedup directly translates to faster document processing throughput, especially when processing documents with many table candidates or complex layouts.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from collections import defaultdict

# imports
import pytest  # used for our unit tests
from unstructured_inference.models.table_postprocess import \
    slot_into_containers

def test_empty_inputs_return_empty_structures():
    # When both container_objects and package_objects are empty,
    # the function should return empty assignment lists and no scores.
    containers = []
    packages = []
    c_assign, p_assign, scores = slot_into_containers(containers, packages) # 2.23μs -> 1.78μs (25.2% faster)

def test_single_container_full_overlap_assigns_package():
    # One container and one package with identical bbox -> full overlap (score 1.0)
    containers = [{"bbox": (0, 0, 10, 10)}]
    packages = [{"bbox": (0, 0, 10, 10)}]
    c_assign, p_assign, scores = slot_into_containers(containers, packages) # 16.6μs -> 6.73μs (147% faster)

def test_no_overlap_below_threshold_still_reports_score_but_no_assignment():
    # Container and package do not overlap -> intersection 0 -> score 0.0
    containers = [{"bbox": (100, 100, 110, 110)}]
    packages = [{"bbox": (0, 0, 10, 10)}]
    c_assign, p_assign, scores = slot_into_containers(containers, packages, overlap_threshold=0.5) # 15.2μs -> 6.75μs (125% faster)

def test_forced_assignment_overrides_threshold_and_assigns_even_with_low_score():
    # If forced_assignment is True, even a package with 0.0 overlap is assigned
    containers = [{"bbox": (100, 100, 110, 110)}]
    packages = [{"bbox": (0, 0, 10, 10)}]
    c_assign, p_assign, scores = slot_into_containers(containers, packages, overlap_threshold=0.9, forced_assignment=True) # 15.2μs -> 6.90μs (121% faster)

def test_package_with_zero_area_is_ignored_no_scores_or_assignments():
    # A package with zero area (x_min == x_max) should produce no match score entries
    containers = [{"bbox": (0, 0, 10, 10)}]
    packages = [{"bbox": (5, 5, 5, 15)}]  # zero width -> area 0
    c_assign, p_assign, scores = slot_into_containers(containers, packages) # 12.2μs -> 4.01μs (205% faster)

def test_zero_area_container_behaves_as_full_match_when_intersecting():
    # If a container has zero area, Rect.intersect logic sets the container rect to the package rect,
    # which leads to an intersection area equal to the package area and thus score 1.0.
    containers = [{"bbox": (0, 0, 0, 0)}]  # zero-area container
    packages = [{"bbox": (0, 0, 10, 10)}]
    c_assign, p_assign, scores = slot_into_containers(containers, packages) # 13.5μs -> 6.28μs (114% faster)

def test_tie_breaking_prefers_earlier_container_in_list():
    # When two containers have identical overlap score for a package,
    # the first container in the original list should be chosen due to Python sort stability.
    containers = [
        {"bbox": (0, 0, 5, 10)},   # left half of package
        {"bbox": (5, 0, 10, 10)},  # right half of package
    ]
    packages = [{"bbox": (0, 0, 10, 10)}]
    # Both containers overlap exactly 50% of package -> equal scores 0.5
    c_assign, p_assign, scores = slot_into_containers(containers, packages, overlap_threshold=0.9) # 19.1μs -> 7.86μs (142% faster)
    # Now force assignment and ensure the earlier container is chosen
    c_assign2, p_assign2, scores2 = slot_into_containers(containers, packages, forced_assignment=True) # 11.9μs -> 4.74μs (151% faster)

def test_multiple_packages_and_containers_large_scale_consistent_assignment():
    # Large-ish scenario under the 1000-element guidance:
    # Create 10 containers (vertical strips) and 50 packages, each package fully inside one container.
    num_containers = 10
    num_packages = 50  # total loops = 500 (packages * containers in worst-case), well under 1000
    containers = []
    # containers are vertical strips of width 10
    for i in range(num_containers):
        containers.append({"bbox": (i * 10, 0, (i + 1) * 10, 100)})

    packages = []
    # Each package will sit entirely inside container index (i % num_containers)
    for i in range(num_packages):
        container_idx = i % num_containers
        x0 = container_idx * 10 + 1
        x1 = (container_idx + 1) * 10 - 1
        packages.append({"bbox": (x0, 10, x1, 90)})

    c_assign, p_assign, scores = slot_into_containers(containers, packages) # 1.47ms -> 220μs (567% faster)
    # Each package should be assigned to exactly one container, and that container index should be consistent.
    # Check counts per container: each container should have either floor/ceil distribution of packages.
    counts = [len(lst) for lst in c_assign]
    # Verify that for each package the assigned container equals the expected one
    for pkg_idx, assigned in enumerate(p_assign):
        expected_container = pkg_idx % num_containers

def test_partial_overlaps_create_proportional_scores_and_assignments():
    # Create one container that covers half of a package and another that covers the other half partially.
    containers = [
        {"bbox": (0, 0, 5, 10)},    # covers left half
        {"bbox": (4, 0, 10, 10)},   # overlaps right side but also overlaps an area of width 1 with left container
    ]
    packages = [{"bbox": (0, 0, 10, 10)}]  # full package width 10 -> area 100
    # Overlaps:
    # container 0 intersects (0,0)-(5,10) -> area 50 -> score 0.5
    # container 1 intersects (4,0)-(10,10) -> area 60 -> score 0.6
    c_assign, p_assign, scores = slot_into_containers(containers, packages, overlap_threshold=0.55) # 18.7μs -> 8.26μs (126% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from collections import defaultdict

import pytest
from unstructured_inference.models.table_postprocess import \
    slot_into_containers

def test_basic_single_package_in_single_container():
    """Test a single package perfectly contained within a single container"""
    # Container from (0,0) to (100,100) with area 10000
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package from (10,10) to (90,90) with area 6400, fully contained
    packages = [{"bbox": (10, 10, 90, 90)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5
    ) # 15.1μs -> 7.30μs (107% faster)

def test_basic_multiple_packages_single_container():
    """Test multiple packages in a single container"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    packages = [
        {"bbox": (10, 10, 40, 40)},
        {"bbox": (50, 50, 80, 80)},
    ]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5
    ) # 21.7μs -> 9.15μs (137% faster)

def test_basic_single_package_multiple_containers():
    """Test a single package overlapping with multiple containers"""
    # Package from (40,40) to (60,60) overlaps with two containers
    containers = [
        {"bbox": (0, 0, 50, 50)},      # Overlaps on left side
        {"bbox": (50, 50, 100, 100)},  # Overlaps on right side
    ]
    packages = [{"bbox": (40, 40, 60, 60)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5
    ) # 18.9μs -> 7.95μs (137% faster)

def test_basic_no_overlap_threshold_not_met():
    """Test that package is not assigned when overlap is below threshold"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package only 20% overlapping
    packages = [{"bbox": (80, 80, 95, 95)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5
    ) # 15.1μs -> 7.07μs (113% faster)

def test_basic_forced_assignment_true():
    """Test that forced_assignment=True assigns package regardless of threshold"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package with very small overlap
    packages = [{"bbox": (95, 95, 105, 105)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5, forced_assignment=True
    ) # 15.1μs -> 7.23μs (108% faster)

def test_basic_exact_threshold_match():
    """Test assignment when overlap exactly equals threshold"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package positioned so overlap is exactly 50% of package area
    packages = [{"bbox": (50, 50, 100, 100)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.5
    ) # 15.5μs -> 7.17μs (115% faster)

def test_edge_empty_containers_list():
    """Test with empty containers list"""
    containers = []
    packages = [{"bbox": (10, 10, 20, 20)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 2.01μs -> 2.01μs (0.249% faster)

def test_edge_empty_packages_list():
    """Test with empty packages list"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    packages = []
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 2.22μs -> 2.17μs (2.59% faster)

def test_edge_both_empty():
    """Test with both empty containers and packages"""
    containers = []
    packages = []
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 1.91μs -> 1.93μs (0.830% slower)

def test_edge_zero_area_package():
    """Test with a package that has zero area (point or line)"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package with zero width (same x_min and x_max)
    packages = [{"bbox": (50, 50, 50, 100)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 12.5μs -> 4.18μs (198% faster)

def test_edge_zero_area_container():
    """Test with a container that has zero area"""
    # Container with zero area (point)
    containers = [{"bbox": (50, 50, 50, 50)}]
    packages = [{"bbox": (0, 0, 100, 100)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 13.5μs -> 6.47μs (109% faster)

def test_edge_no_overlap_packages():
    """Test packages completely outside containers"""
    containers = [{"bbox": (0, 0, 50, 50)}]
    packages = [{"bbox": (100, 100, 150, 150)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.0
    ) # 15.6μs -> 7.05μs (122% faster)

def test_edge_package_larger_than_container():
    """Test package that is larger than container"""
    containers = [{"bbox": (40, 40, 60, 60)}]
    # Large package encompassing container
    packages = [{"bbox": (0, 0, 100, 100)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.1
    ) # 15.1μs -> 7.03μs (114% faster)

def test_edge_negative_coordinates():
    """Test with negative coordinate values"""
    containers = [{"bbox": (-100, -100, 0, 0)}]
    packages = [{"bbox": (-50, -50, 50, 50)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.1
    ) # 15.0μs -> 7.16μs (109% faster)

def test_edge_very_small_overlap():
    """Test with very small overlap fractions"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    # Package mostly outside container
    packages = [{"bbox": (99, 99, 200, 200)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.01
    ) # 15.2μs -> 6.93μs (119% faster)

def test_edge_identical_containers_and_packages():
    """Test when container and package have identical bounds"""
    containers = [{"bbox": (10, 10, 90, 90)}]
    packages = [{"bbox": (10, 10, 90, 90)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 14.6μs -> 6.84μs (114% faster)

def test_edge_touching_but_not_overlapping():
    """Test containers and packages that touch at edges but don't overlap"""
    containers = [{"bbox": (0, 0, 50, 50)}]
    # Package starts where container ends
    packages = [{"bbox": (50, 50, 100, 100)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.0
    ) # 16.2μs -> 7.22μs (124% faster)

def test_edge_multiple_containers_same_package():
    """Test a package with equal overlap in multiple containers"""
    containers = [
        {"bbox": (0, 0, 50, 100)},
        {"bbox": (50, 0, 100, 100)},
    ]
    # Package centered between containers with equal overlap
    packages = [{"bbox": (40, 40, 60, 60)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 19.6μs -> 7.94μs (146% faster)

def test_edge_threshold_zero():
    """Test with overlap_threshold set to 0"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    packages = [{"bbox": (99.5, 99.5, 150, 150)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.0
    ) # 16.5μs -> 8.32μs (97.8% faster)

def test_edge_threshold_one():
    """Test with overlap_threshold set to 1.0"""
    containers = [{"bbox": (0, 0, 100, 100)}]
    packages = [{"bbox": (10, 10, 90, 90)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=1.0
    ) # 15.5μs -> 7.20μs (115% faster)

def test_edge_floating_point_coordinates():
    """Test with floating point coordinate values"""
    containers = [{"bbox": (0.5, 0.5, 99.5, 99.5)}]
    packages = [{"bbox": (25.25, 25.25, 75.75, 75.75)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 15.0μs -> 6.82μs (120% faster)

def test_large_scale_many_containers_single_package():
    """Test with many containers and a single package"""
    # Create 500 containers in a grid
    containers = []
    for i in range(25):
        for j in range(20):
            containers.append({
                "bbox": (i * 40, j * 40, (i + 1) * 40, (j + 1) * 40)
            })
    
    # Single package in the middle
    packages = [{"bbox": (480, 380, 520, 420)}]
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 1.50ms -> 236μs (534% faster)
    # Exactly one container should have the package
    assigned_count = sum(1 for c in container_assigns if len(c) > 0)

def test_large_scale_many_packages_single_container():
    """Test with many packages in a single container"""
    containers = [{"bbox": (0, 0, 1000, 1000)}]
    
    # Create 400 packages in a grid
    packages = []
    for i in range(20):
        for j in range(20):
            packages.append({
                "bbox": (i * 50, j * 50, (i + 1) * 50, (j + 1) * 50)
            })
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 1.88ms -> 457μs (311% faster)

def test_large_scale_many_containers_many_packages():
    """Test with many containers and many packages"""
    # Create 200 containers in a grid
    containers = []
    for i in range(20):
        for j in range(10):
            containers.append({
                "bbox": (i * 50, j * 100, (i + 1) * 50, (j + 1) * 100)
            })
    
    # Create 200 packages
    packages = []
    for i in range(20):
        for j in range(10):
            packages.append({
                "bbox": (i * 50 + 10, j * 100 + 20, (i + 1) * 50 - 10, (j + 1) * 100 - 20)
            })
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 110ms -> 14.4ms (666% faster)
    
    # Most packages should be assigned (with good threshold)
    assigned_packages = sum(1 for p in package_assigns if len(p) > 0)

def test_large_scale_scattered_packages():
    """Test with containers arranged normally but packages scattered"""
    # Create 100 containers in a grid
    containers = []
    for i in range(10):
        for j in range(10):
            containers.append({
                "bbox": (i * 100, j * 100, (i + 1) * 100, (j + 1) * 100)
            })
    
    # Create 300 packages scattered with various overlap amounts
    packages = []
    for idx in range(300):
        start_x = (idx * 17) % 1000
        start_y = (idx * 29) % 1000
        packages.append({
            "bbox": (start_x, start_y, start_x + 40, start_y + 40)
        })
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.3
    ) # 82.8ms -> 11.1ms (646% faster)
    
    # Verify consistency
    total_package_assignments = sum(len(p) for p in package_assigns)
    total_container_assignments = sum(len(c) for c in container_assigns)

def test_large_scale_dense_overlapping():
    """Test with dense, heavily overlapping objects"""
    # Create 100 overlapping containers
    containers = []
    for i in range(100):
        x_offset = (i * 5) % 500
        y_offset = (i // 10) * 5
        containers.append({
            "bbox": (x_offset, y_offset, x_offset + 100, y_offset + 100)
        })
    
    # Create 100 packages
    packages = []
    for i in range(100):
        x_offset = (i * 3.5) % 500
        y_offset = (i // 10) * 3.5
        packages.append({
            "bbox": (x_offset, y_offset, x_offset + 50, y_offset + 50)
        })
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages, overlap_threshold=0.1
    ) # 29.1ms -> 4.63ms (529% faster)
    
    # With high overlap, most should be assigned
    assigned_packages = sum(1 for p in package_assigns if len(p) > 0)

def test_large_scale_performance_non_overlapping():
    """Test performance with many non-overlapping objects"""
    # Create 300 non-overlapping containers
    containers = []
    for i in range(30):
        for j in range(10):
            containers.append({
                "bbox": (i * 100, j * 100, (i + 1) * 100, (j + 1) * 100)
            })
    
    # Create 300 non-overlapping packages in different locations
    packages = []
    for i in range(30):
        for j in range(10):
            packages.append({
                "bbox": (i * 100 + 10000, j * 100 + 10000, (i + 1) * 100 + 10000, (j + 1) * 100 + 10000)
            })
    
    container_assigns, package_assigns, scores = slot_into_containers(
        containers, packages
    ) # 247ms -> 33.8ms (633% faster)
    
    # No packages should be assigned (no overlap)
    assigned_packages = sum(1 for p in package_assigns if len(p) > 0)
# 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-slot_into_containers-mkos4jpd and push.