⚡️ Speed up function _fix_metadata_field_precision by 304%

[codeflash-ai]

📄 304% (3.04x) speedup for _fix_metadata_field_precision in unstructured/staging/base.py

⏱️ Runtime : 164 milliseconds → 40.6 milliseconds (best of 57 runs)

📝 Explanation and details

The optimized code achieves a 304% speedup by replacing expensive deepcopy() operations with selective shallow copying. Here's what changed and why it matters:

Key Optimization: Selective Shallow Copying

Original approach: Called deepcopy(element) on every element, which recursively copies the entire object graph—accounting for 94.8% of the function's runtime.

Optimized approach: Uses copy(element) for shallow copying, then selectively copies only the metadata objects that need modification:

• Performs shallow copy() of element.metadata and element.metadata.coordinates only when coordinates exist
• Uses a generator expression with tuple() instead of building a list for rounded points, which is more memory-efficient
• Only copies element.metadata when detection_class_prob needs rounding AND metadata hasn't already been copied (via if el.metadata is element.metadata check)

Why This Works

The optimization exploits the fact that most element attributes don't need modification—only specific metadata fields require precision adjustment. By sharing unmodified attributes between the original and new elements, we avoid the overhead of deep copying nested objects like coordinate systems, detection origins, and other metadata fields.

Performance breakdown:

• Shallow copy() is ~32x faster than deepcopy() for these element objects
• Generator expression for point rounding eliminates intermediate list allocation
• The identity check (is) prevents redundant metadata copying when both coordinate and probability rounding are needed

Impact on Workloads

Based on function_references, this function is called in serialization hot paths:

• elements_to_base64_gzipped_json() - Used for compressing elements in HTTP responses
• elements_to_json() - Direct JSON serialization
• elements_to_ndjson() - Newline-delimited JSON serialization

All three functions call _fix_metadata_field_precision() before serialization, making this optimization critical for document processing pipelines that serialize large volumes of elements.

Test Results

The optimization excels across all test scenarios:

• Simple elements (no coordinates): 305-721% faster
• Single element with coordinates: 151-157% faster
• Large-scale (500 elements): 293-310% faster
• Many points per element (200 points): 240% faster

The speedup is consistent regardless of coordinate system type (PixelSpace vs. CoordinateSystem) or whether detection probabilities are present, demonstrating robust performance gains across diverse workloads.

✅ 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

from copy import deepcopy  # used by the function under test (kept for exact function copy)

# imports
import pytest  # used for our unit tests

from unstructured.documents.coordinates import CoordinateSystem, PixelSpace

# Import real domain classes from their modules (do not stub or fake them).
from unstructured.documents.elements import CoordinatesMetadata, Element, ElementMetadata
from unstructured.staging.base import _fix_metadata_field_precision


def test_pixel_space_rounding_and_probability_rounding_basic():
    # Basic test: PixelSpace should round coordinates to 1 decimal place,
    # and a truthy detection_class_prob should be rounded to 5 decimals.
    system = PixelSpace(100, 100)  # Pixel coordinate system => precision 1
    coords = ((1.234, 5.6789), (2.3456, 3.4567))  # points to round
    el = Element(coordinates=coords, coordinate_system=system)
    # Set a detection class probability that will be rounded to 5 decimals.
    el.metadata.detection_class_prob = 0.123456789

    codeflash_output = _fix_metadata_field_precision([el])
    results = codeflash_output  # 105μs -> 41.0μs (157% faster)

    fixed = results[0]
    # PixelSpace precision 1 -> round to 1 decimal place.
    expected_points = ((round(1.234, 1), round(5.6789, 1)), (round(2.3456, 1), round(3.4567, 1)))


def test_non_pixel_space_precision_two_and_probability_edge_cases():
    # Non-PixelSpace (CoordinateSystem instance) should use precision 2.
    system = CoordinateSystem(200, 200)
    coords = ((1.23456, -2.34567), (0.0, 4.9999))
    el = Element(coordinates=coords, coordinate_system=system)

    # detection_class_prob very small but truthy -> will be rounded to 5 places (likely 0.0)
    el.metadata.detection_class_prob = 1e-7

    codeflash_output = _fix_metadata_field_precision([el])
    results = codeflash_output  # 106μs -> 41.9μs (154% faster)
    fixed = results[0]

    # Coordinates rounded to 2 decimals
    expected_points = tuple((round(x, 2), round(y, 2)) for (x, y) in coords)

    # Now test that detection_class_prob == 0.0 (falsy) is not rounded/changed by the function.
    el_zero = Element(coordinates=coords, coordinate_system=system)
    el_zero.metadata.detection_class_prob = 0.0  # falsy, function should skip rounding branch
    codeflash_output = _fix_metadata_field_precision([el_zero])
    results_zero = codeflash_output  # 80.6μs -> 23.0μs (251% faster)
    fixed_zero = results_zero[0]


def test_coordinates_object_with_both_none_points_triggers_assertion():
    # Edge case: A CoordinatesMetadata object may be present with both points=None and system=None.
    # In that case, el.metadata.coordinates is truthy (object exists), but .points is None and the
    # function asserts points is not None. This should raise an AssertionError.
    coords_obj = CoordinatesMetadata(points=None, system=None)  # allowed by constructor
    meta = ElementMetadata(coordinates=coords_obj)
    el = Element(metadata=meta)

    # Passing such an element should cause an AssertionError inside the function.
    with pytest.raises(AssertionError):
        _fix_metadata_field_precision([el])  # 78.4μs -> 32.5μs (142% faster)


def test_negative_and_boundary_coordinate_values_rounding_behavior():
    # Test negative values and tie-breaking behavior at .5 and .005 boundaries.
    system = CoordinateSystem(50, 50)  # non-pixel -> 2 decimal precision
    # Values chosen to exercise boundary rounding rules (Python's round uses "bankers rounding").
    coords = ((-1.2345, 2.3455), (1.255, -2.255))
    el = Element(coordinates=coords, coordinate_system=system)
    # detection probability negative (edge acceptance) and should be rounded as well.
    el.metadata.detection_class_prob = -0.1234567

    fixed = _fix_metadata_field_precision([el])[0]  # 106μs -> 42.2μs (153% faster)

    # Ensure rounding to 2 decimals for coordinates, using Python's round behavior.
    expected_points = tuple((round(x, 2), round(y, 2)) for (x, y) in coords)


def test_original_elements_unchanged_after_fix():
    # Ensure the function returns deep-copied elements and does not mutate input elements.
    system = PixelSpace(10, 10)
    coords = ((9.876, 0.123),)
    el = Element(coordinates=coords, coordinate_system=system)
    el.metadata.detection_class_prob = 0.9999999

    # Keep deep copy of original metadata for comparison after function call.
    original_metadata = deepcopy(el.metadata)

    codeflash_output = _fix_metadata_field_precision([el])
    results = codeflash_output  # 79.7μs -> 30.1μs (165% faster)
    fixed = results[0]


def test_large_scale_mixed_coordinate_systems_performance_and_correctness():
    # Large-scale test: create a sequence of many elements (below the 1000-element instruction limit)
    # and ensure the function processes them all correctly and returns the expected count.
    count = 500  # well under 1000 as requested
    elements = []
    for i in range(count):
        # Alternate between PixelSpace and CoordinateSystem to exercise both precisions.
        if i % 2 == 0:
            system = PixelSpace(100, 100)
            # Coordinate to be rounded to 1 decimal
            coords = ((i + 0.12345, i + 0.54321),)
        else:
            system = CoordinateSystem(100, 100)
            # Coordinate to be rounded to 2 decimals
            coords = ((i + 0.98765, i + 0.56789),)
        el = Element(coordinates=coords, coordinate_system=system)
        # Provide a deterministic detection_class_prob; some will be falsy (0.0), others truthy.
        el.metadata.detection_class_prob = (
            i % 5
        ) * 0.00001  # yields 0.0 sometimes, small truthy values others
        elements.append(el)

    codeflash_output = _fix_metadata_field_precision(elements)
    results = codeflash_output  # 33.3ms -> 8.13ms (310% faster)
    for original, fixed in zip(elements, results):
        # Verify coordinate rounding according to system type.
        if isinstance(original.metadata.coordinates.system, PixelSpace):
            expected = tuple(
                (round(x, 1), round(y, 1)) for (x, y) in original.metadata.coordinates.points
            )
        else:
            expected = tuple(
                (round(x, 2), round(y, 2)) for (x, y) in original.metadata.coordinates.points
            )


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from copy import deepcopy
from typing import Optional

from unstructured.documents.coordinates import CoordinateSystem, PixelSpace
from unstructured.documents.elements import CoordinatesMetadata, Element, ElementMetadata
from unstructured.staging.base import _fix_metadata_field_precision


# Concrete Element subclass for testing
class Text(Element):
    """Concrete Text element for testing."""

    def __init__(
        self,
        text: str = "",
        element_id: Optional[str] = None,
        coordinates: Optional[tuple[tuple[float, float], ...]] = None,
        coordinate_system: Optional[CoordinateSystem] = None,
        metadata: Optional[ElementMetadata] = None,
        detection_origin: Optional[str] = None,
    ):
        self.text = text
        super().__init__(
            element_id=element_id,
            coordinates=coordinates,
            coordinate_system=coordinate_system,
            metadata=metadata,
            detection_origin=detection_origin,
        )


def test_basic_empty_elements_list():
    """Test that function handles an empty list of elements correctly."""
    codeflash_output = _fix_metadata_field_precision([])
    result = codeflash_output  # 915ns -> 920ns (0.543% slower)


def test_basic_single_element_no_coordinates_no_detection():
    """Test a single element with no coordinates and no detection class prob."""
    element = Text("Hello World")
    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 73.7μs -> 18.2μs (305% faster)


def test_basic_pixel_space_coordinates_rounding():
    """Test that PixelSpace coordinates are rounded to 1 decimal place."""
    points = ((10.123, 20.456), (30.789, 40.999))
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Test", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 106μs -> 42.4μs (151% faster)
    rounded_points = result[0].metadata.coordinates.points


def test_basic_generic_coordinate_system_rounding():
    """Test that non-PixelSpace coordinates are rounded to 2 decimal places."""
    points = ((10.123, 20.456), (30.789, 40.999))
    coord_system = CoordinateSystem(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Test", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 105μs -> 41.9μs (151% faster)
    rounded_points = result[0].metadata.coordinates.points


def test_basic_detection_class_prob_rounding():
    """Test that detection_class_prob is rounded to 5 decimal places."""
    metadata = ElementMetadata(detection_class_prob=0.123456789)
    element = Text("Test", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 77.4μs -> 32.9μs (136% faster)


def test_basic_multiple_elements():
    """Test processing multiple elements with mixed metadata."""
    points1 = ((1.111, 2.222),)
    coord_system = PixelSpace(width=50, height=50)
    metadata1 = ElementMetadata(
        coordinates=CoordinatesMetadata(points=points1, system=coord_system)
    )
    element1 = Text("First", metadata=metadata1)

    points2 = ((3.333, 4.444),)
    metadata2 = ElementMetadata(
        coordinates=CoordinatesMetadata(points=points2, system=coord_system)
    )
    element2 = Text("Second", metadata=metadata2)

    codeflash_output = _fix_metadata_field_precision([element1, element2])
    result = codeflash_output  # 173μs -> 59.9μs (190% faster)


def test_basic_both_coordinates_and_detection_prob():
    """Test element with both coordinates and detection_class_prob."""
    points = ((5.5555, 6.6666),)
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(
        coordinates=CoordinatesMetadata(points=points, system=coord_system),
        detection_class_prob=0.987654321,
    )
    element = Text("Both", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 103μs -> 40.3μs (156% faster)


def test_edge_zero_coordinates():
    """Test coordinates with zero values."""
    points = ((0.0, 0.0), (0.123, 0.456))
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Zero", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 105μs -> 41.5μs (154% faster)


def test_edge_negative_coordinates():
    """Test coordinates with negative values."""
    points = ((-10.123, -20.456), (-0.999, -0.111))
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Negative", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 105μs -> 41.7μs (153% faster)


def test_edge_large_coordinate_values():
    """Test coordinates with very large values."""
    points = ((999999.999, 888888.888),)
    coord_system = PixelSpace(width=1000000, height=1000000)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Large", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 101μs -> 40.6μs (149% faster)


def test_edge_very_small_coordinate_values():
    """Test coordinates with very small decimal values."""
    points = ((0.00001, 0.00002), (0.000001, 0.000002))
    coord_system = CoordinateSystem(width=1, height=1)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Small", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 106μs -> 42.1μs (152% faster)


def test_edge_single_point():
    """Test with a single point in coordinates."""
    points = ((42.5678, 99.1234),)
    coord_system = PixelSpace(width=200, height=200)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("SinglePoint", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 101μs -> 39.7μs (156% faster)


def test_edge_many_points():
    """Test with many points in coordinates."""
    points = tuple((i + 0.123, i + 0.456) for i in range(100))
    coord_system = PixelSpace(width=1000, height=1000)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("ManyPoints", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 412μs -> 125μs (228% faster)

    rounded_points = result[0].metadata.coordinates.points


def test_edge_detection_class_prob_zero():
    """Test detection_class_prob with zero value."""
    metadata = ElementMetadata(detection_class_prob=0.0)
    element = Text("Zero Prob", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 72.5μs -> 16.5μs (338% faster)


def test_edge_detection_class_prob_one():
    """Test detection_class_prob with value of 1.0."""
    metadata = ElementMetadata(detection_class_prob=1.0)
    element = Text("Full Prob", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 77.0μs -> 32.6μs (136% faster)


def test_edge_detection_class_prob_very_small():
    """Test detection_class_prob with very small value."""
    metadata = ElementMetadata(detection_class_prob=0.0000012345)
    element = Text("Tiny Prob", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 77.3μs -> 32.6μs (137% faster)


def test_edge_detection_class_prob_rounding_up():
    """Test detection_class_prob rounding that rounds up."""
    metadata = ElementMetadata(detection_class_prob=0.999995)
    element = Text("Prob Round Up", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 78.0μs -> 33.5μs (133% faster)


def test_edge_original_element_unchanged():
    """Test that the original element is not modified (deep copy is used)."""
    points = ((10.123, 20.456),)
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Original", metadata=metadata)
    original_points = element.metadata.coordinates.points

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 100μs -> 40.4μs (149% faster)


def test_edge_coordinates_none_with_valid_detection():
    """Test element with no coordinates but valid detection_class_prob."""
    metadata = ElementMetadata(detection_class_prob=0.555555)
    element = Text("NoCoord", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 77.2μs -> 33.1μs (133% faster)


def test_edge_integer_coordinates():
    """Test coordinates that are integers (no decimals)."""
    points = ((10, 20), (30, 40))
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Integers", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 101μs -> 37.3μs (172% faster)


def test_edge_mixed_integer_and_float_coordinates():
    """Test coordinates with mixed integers and floats."""
    points = ((10, 20.456), (30.789, 40))
    coord_system = CoordinateSystem(width=100, height=100)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("Mixed", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 104μs -> 41.3μs (154% faster)


def test_edge_subclass_of_coordinate_system():
    """Test that PixelSpace (subclass of CoordinateSystem) is correctly identified."""
    points = ((1.12345, 2.98765),)
    pixel_space = PixelSpace(width=800, height=600)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=pixel_space))
    element = Text("PixelTest", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 101μs -> 40.4μs (152% faster)


def test_edge_detection_class_prob_negative():
    """Test detection_class_prob with negative value (edge case, though not typical)."""
    metadata = ElementMetadata(detection_class_prob=-0.123456)
    element = Text("Negative Prob", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 76.8μs -> 32.5μs (136% faster)


def test_large_scale_many_elements():
    """Test processing a large number of elements."""
    elements = []
    for i in range(500):
        points = ((float(i) + 0.123, float(i) + 0.456),)
        coord_system = PixelSpace(width=1000, height=1000)
        metadata = ElementMetadata(
            coordinates=CoordinatesMetadata(points=points, system=coord_system),
            detection_class_prob=0.5 + (i % 10) * 0.01234567,
        )
        elements.append(Text(f"Element {i}", metadata=metadata))

    codeflash_output = _fix_metadata_field_precision(elements)
    result = codeflash_output  # 33.7ms -> 8.30ms (306% faster)


def test_large_scale_many_points_per_element():
    """Test elements with many points in coordinates."""
    points = tuple((i + 0.5555, i + 0.4444) for i in range(200))
    coord_system = CoordinateSystem(width=5000, height=5000)
    metadata = ElementMetadata(coordinates=CoordinatesMetadata(points=points, system=coord_system))
    element = Text("ManyPointsElement", metadata=metadata)

    codeflash_output = _fix_metadata_field_precision([element])
    result = codeflash_output  # 736μs -> 216μs (240% faster)
    # All points should be rounded to 2 decimal places
    for idx, point in enumerate(result[0].metadata.coordinates.points):
        pass


def test_large_scale_mixed_coordinate_systems():
    """Test processing elements with different coordinate systems."""
    elements = []

    # Add elements with PixelSpace
    for i in range(250):
        points = ((i + 0.111, i + 0.222),)
        metadata = ElementMetadata(
            coordinates=CoordinatesMetadata(points=points, system=PixelSpace(1920, 1080))
        )
        elements.append(Text(f"Pixel {i}", metadata=metadata))

    # Add elements with generic CoordinateSystem
    for i in range(250):
        points = ((i + 0.111, i + 0.222),)
        metadata = ElementMetadata(
            coordinates=CoordinatesMetadata(points=points, system=CoordinateSystem(100, 100))
        )
        elements.append(Text(f"Generic {i}", metadata=metadata))

    codeflash_output = _fix_metadata_field_precision(elements)
    result = codeflash_output  # 32.8ms -> 8.35ms (293% faster)


def test_large_scale_elements_without_metadata():
    """Test processing many elements without metadata or coordinates."""
    elements = [Text(f"Simple Text {i}") for i in range(500)]

    codeflash_output = _fix_metadata_field_precision(elements)
    result = codeflash_output  # 22.2ms -> 2.70ms (721% faster)
    for i, elem in enumerate(result):
        pass


def test_large_scale_all_fields_populated():
    """Test processing elements with all metadata fields populated."""
    elements = []
    for i in range(100):
        points = ((i + 0.123, i + 0.456),)
        metadata = ElementMetadata(
            coordinates=CoordinatesMetadata(points=points, system=PixelSpace(800, 600)),
            detection_class_prob=0.12345 + (i * 0.0001),
            category_depth=i,
            page_number=i % 10,
            filename=f"file_{i}.pdf",
        )
        elements.append(Text(f"Full Element {i}", metadata=metadata))

    codeflash_output = _fix_metadata_field_precision(elements)
    result = codeflash_output  # 7.14ms -> 1.72ms (314% faster)


def test_large_scale_extreme_precision_values():
    """Test rounding with extreme precision requirements."""
    elements = []
    for i in range(100):
        points = ((0.1 + i * 0.000001, 0.2 + i * 0.000002),)
        coord_system = CoordinateSystem(width=1, height=1)
        metadata = ElementMetadata(
            coordinates=CoordinatesMetadata(points=points, system=coord_system),
            detection_class_prob=0.1000001 + i * 0.00000001,
        )
        elements.append(Text(f"Precision {i}", metadata=metadata))

    codeflash_output = _fix_metadata_field_precision(elements)
    result = codeflash_output  # 6.73ms -> 1.67ms (302% faster)
    # All coordinates should be rounded to 2 decimal places
    for elem in result:
        for point in elem.metadata.coordinates.points:
            x, y = point


def test_large_scale_consistency():
    """Test that processing the same elements multiple times yields consistent results."""
    points = ((10.123456, 20.654321),)
    coord_system = PixelSpace(width=100, height=100)
    metadata = ElementMetadata(
        coordinates=CoordinatesMetadata(points=points, system=coord_system),
        detection_class_prob=0.123456789,
    )
    element = Text("Consistency", metadata=metadata)

    # Process the same element multiple times
    codeflash_output = _fix_metadata_field_precision([deepcopy(element)])
    result1 = codeflash_output  # 80.5μs -> 29.5μs (173% faster)
    codeflash_output = _fix_metadata_field_precision([deepcopy(element)])
    result2 = codeflash_output  # 71.2μs -> 21.1μs (237% faster)
    codeflash_output = _fix_metadata_field_precision([deepcopy(element)])
    result3 = codeflash_output  # 69.0μs -> 20.0μs (245% faster)


# 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-_fix_metadata_field_precision-mkrzwqch and push.