⚡️ Speed up function _calculate_percentile by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for _calculate_percentile in mlflow/tracing/utils/__init__.py

⏱️ Runtime : 79.7 microseconds → 72.2 microseconds (best of 101 runs)

📝 Explanation and details

The optimization achieves a 10% speedup through three key changes that reduce computational overhead:

What was optimized:

1. Eliminated redundant not sorted_data check: Moved n = len(sorted_data) before the empty check and changed if not sorted_data: to if n == 0:. This avoids calling len() twice in the common non-empty case.

2. Reduced list access operations: Stored sorted_data[lower] and sorted_data[upper] in local variables (lower_value, upper_value) instead of accessing them multiple times during interpolation.

3. Simplified interpolation math: Changed from sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight to the mathematically equivalent but computationally simpler lower_value + weight * (upper_value - lower_value). This reduces from 3 arithmetic operations to 2.

Why this leads to speedup:

• Fewer function calls: Eliminates one len() call in the typical execution path
• Reduced memory access: Local variables avoid repeated list indexing operations
• Simpler arithmetic: The refactored interpolation formula requires fewer multiplications

Impact on workloads:
The function is called in a hot path within add_size_stats_to_trace_metadata() to compute P25, P50, and P75 percentiles for trace span sizes. Since this runs for every trace processed, the 10% improvement compounds significantly in high-throughput tracing scenarios.

Test case performance:
The optimization performs best on typical interpolation cases (10-20% faster) but shows slight regression on edge cases like empty lists or percentiles that hit exact boundaries. However, the common case of percentile calculation with interpolation - which represents the primary use case in production - sees consistent improvements.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from mlflow.tracing.utils.__init__ import _calculate_percentile

# unit tests

# ------------------------
# Basic Test Cases
# ------------------------

def test_percentile_basic_middle():
    # Test 50th percentile (median) in odd-length list
    data = [1, 2, 3, 4, 5]
    # Median is 3
    codeflash_output = _calculate_percentile(data, 0.5) # 1.80μs -> 1.53μs (17.2% faster)

def test_percentile_basic_quartile():
    # Test 25th percentile in sorted list
    data = [10, 20, 30, 40]
    # 0.25 * (4-1) = 0.75; lower=0, upper=1, weight=0.75
    # value = 10*0.25 + 20*0.75 = 2.5 + 15 = 17.5
    expected = 10 * (1 - 0.75) + 20 * 0.75
    codeflash_output = _calculate_percentile(data, 0.25) # 1.82μs -> 1.52μs (19.2% faster)

def test_percentile_basic_last():
    # Test 100th percentile (max)
    data = [1, 2, 3, 4, 5]
    codeflash_output = _calculate_percentile(data, 1.0) # 1.15μs -> 1.17μs (1.29% slower)

def test_percentile_basic_first():
    # Test 0th percentile (min)
    data = [1, 2, 3, 4, 5]
    codeflash_output = _calculate_percentile(data, 0.0) # 1.81μs -> 1.54μs (17.9% faster)

def test_percentile_basic_single_element():
    # Test with single element list
    data = [42]
    # Should always return the only value
    codeflash_output = _calculate_percentile(data, 0.0) # 1.24μs -> 1.21μs (2.40% faster)
    codeflash_output = _calculate_percentile(data, 0.5) # 360ns -> 392ns (8.16% slower)
    codeflash_output = _calculate_percentile(data, 1.0) # 250ns -> 236ns (5.93% faster)

def test_percentile_basic_two_elements():
    # Test with two elements, check interpolation
    data = [10, 20]
    # 50th percentile: index = 0.5 * (2-1) = 0.5
    # lower=0, upper=1, weight=0.5
    expected = 10*0.5 + 20*0.5
    codeflash_output = _calculate_percentile(data, 0.5) # 1.73μs -> 1.55μs (11.6% faster)

# ------------------------
# Edge Test Cases
# ------------------------

def test_percentile_edge_empty_list():
    # Should return 0.0 for empty list
    codeflash_output = _calculate_percentile([], 0.5) # 445ns -> 596ns (25.3% slower)

def test_percentile_edge_percentile_below_zero():
    # Negative percentile, should clamp to min value
    data = [5, 10, 15]
    # index = -0.5 * (3-1) = -1.0, lower=-1, upper=0, weight=0.0
    # But function does not clamp, so lower=-1, upper=0
    # But upper >= n is False, so interpolation between data[-1] and data[0]
    # Actually, with negative index, lower=-1, upper=0, weight=0.0
    # sorted_data[-1] = 15, sorted_data[0] = 5
    # value = 15*1 + 5*0 = 15
    # This is a bug, but as per implementation, this is what happens
    codeflash_output = _calculate_percentile(data, -0.5) # 1.99μs -> 1.78μs (11.6% faster)

def test_percentile_edge_percentile_above_one():
    # Percentile > 1, should clamp to max value
    data = [5, 10, 15]
    # index = 1.5 * (3-1) = 3.0, lower=3, upper=4 (upper >= n)
    # Should return sorted_data[-1] = 15
    codeflash_output = _calculate_percentile(data, 1.5) # 1.21μs -> 1.18μs (2.37% faster)

def test_percentile_edge_non_integer_percentile():
    # Test with non-integer percentile
    data = [0, 10, 20, 30, 40]
    # 0.33 * 4 = 1.32, lower=1, upper=2, weight=0.32
    expected = 10 * (1 - 0.32) + 20 * 0.32
    codeflash_output = _calculate_percentile(data, 0.33) # 1.90μs -> 1.63μs (16.9% faster)

def test_percentile_edge_unsorted_input():
    # Function expects sorted data; test with unsorted input
    data = [30, 10, 20]
    # Should not sort internally, so result is based on input order
    # 50th percentile: index = 0.5 * (3-1) = 1.0, lower=1, upper=2, weight=0.0
    # value = data[1] = 10
    codeflash_output = _calculate_percentile(data, 0.5) # 1.86μs -> 1.57μs (18.6% faster)

def test_percentile_edge_all_equal():
    # All elements equal
    data = [7, 7, 7, 7, 7]
    # Any percentile should return 7
    for p in [0.0, 0.25, 0.5, 0.75, 1.0]:
        codeflash_output = _calculate_percentile(data, p) # 3.59μs -> 3.32μs (8.23% faster)

def test_percentile_edge_float_values():
    # Test with float values
    data = [1.1, 2.2, 3.3, 4.4]
    # 75th percentile: index = 0.75 * 3 = 2.25, lower=2, upper=3, weight=0.25
    expected = 3.3 * (1 - 0.25) + 4.4 * 0.25
    codeflash_output = _calculate_percentile(data, 0.75) # 1.71μs -> 1.52μs (12.5% faster)

def test_percentile_edge_percentile_exact_index():
    # Percentile that lands exactly on an index
    data = [1, 2, 3, 4]
    # 2nd index: percentile = 2 / (4-1) = 0.666...
    percentile = 2 / 3
    codeflash_output = _calculate_percentile(data, percentile) # 1.89μs -> 1.59μs (19.1% faster)

# ------------------------
# Large Scale Test Cases
# ------------------------

def test_percentile_large_scale_uniform():
    # Large uniform list, percentile should interpolate between min and max
    data = list(range(1000))  # 0 to 999
    # 50th percentile: index = 0.5 * (999) = 499.5
    # lower=499, upper=500, weight=0.5
    expected = 499 * 0.5 + 500 * 0.5
    codeflash_output = _calculate_percentile(data, 0.5) # 2.01μs -> 1.75μs (14.8% faster)

def test_percentile_large_scale_skewed():
    # Large list with repeated values, test upper percentile
    data = [1]*500 + [100]*500
    # 90th percentile: index = 0.9 * (999) = 899.1, lower=899, upper=900, weight=0.1
    # data[899]=100, data[900]=100
    expected = 100 * (1 - 0.1) + 100 * 0.1
    codeflash_output = _calculate_percentile(data, 0.9) # 2.03μs -> 1.73μs (17.5% faster)

def test_percentile_large_scale_min_max():
    # Test min and max percentiles on large data
    data = [i for i in range(1000)]
    codeflash_output = _calculate_percentile(data, 0.0) # 2.01μs -> 1.75μs (15.0% faster)
    codeflash_output = _calculate_percentile(data, 1.0) # 680ns -> 691ns (1.59% slower)

def test_percentile_large_scale_edge():
    # Test with percentile very close to 0 and 1
    data = list(range(1000))
    # percentile just above 0
    p = 1e-6
    index = p * (1000 - 1)
    lower = int(index)
    upper = lower + 1
    weight = index - lower
    expected = data[lower] * (1 - weight) + data[upper] * weight
    codeflash_output = _calculate_percentile(data, p) # 1.28μs -> 1.40μs (8.55% slower)

    # percentile just below 1
    p = 1 - 1e-6
    index = p * (1000 - 1)
    lower = int(index)
    upper = lower + 1
    weight = index - lower
    expected = data[lower] * (1 - weight) + data[upper] * weight
    codeflash_output = _calculate_percentile(data, p) # 657ns -> 764ns (14.0% slower)

def test_percentile_large_scale_reverse():
    # Test with large reversed list
    data = list(reversed(range(1000)))
    # 25th percentile
    p = 0.25
    index = p * (1000 - 1)
    lower = int(index)
    upper = lower + 1
    weight = index - lower
    expected = data[lower] * (1 - weight) + data[upper] * weight
    codeflash_output = _calculate_percentile(data, p) # 1.26μs -> 1.38μs (8.36% slower)

# ------------------------
# Mutation-sensitive test cases
# ------------------------

@pytest.mark.parametrize("data,percentile,expected", [
    ([1, 2, 3, 4], 0.25, 1 * 0.75 + 2 * 0.25),
    ([1, 2, 3, 4], 0.75, 3 * 0.25 + 4 * 0.75),
    ([10, 20, 30, 40], 0.5, 20),
    ([10, 20, 30, 40], 0.99, 40),
])
def test_percentile_mutation_sensitive(data, percentile, expected):
    # These should fail if implementation is changed
    codeflash_output = _calculate_percentile(data, percentile) # 7.56μs -> 6.41μs (17.9% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from mlflow.tracing.utils.__init__ import _calculate_percentile

# unit tests

class TestCalculatePercentile:
    # --- Basic Test Cases ---

    def test_empty_list_returns_zero(self):
        """Test that an empty list returns 0.0 for any percentile"""
        codeflash_output = _calculate_percentile([], 0.5) # 429ns -> 609ns (29.6% slower)
        codeflash_output = _calculate_percentile([], 0.0) # 190ns -> 203ns (6.40% slower)
        codeflash_output = _calculate_percentile([], 1.0) # 118ns -> 146ns (19.2% slower)

    def test_single_element_list(self):
        """Test that a single-element list always returns that element"""
        codeflash_output = _calculate_percentile([42], 0.0) # 1.26μs -> 1.20μs (5.77% faster)
        codeflash_output = _calculate_percentile([42], 0.5) # 383ns -> 378ns (1.32% faster)
        codeflash_output = _calculate_percentile([42], 1.0) # 236ns -> 252ns (6.35% slower)

    def test_percentile_at_bounds(self):
        """Test percentile at 0.0 and 1.0 returns first and last element"""
        data = [1, 2, 3, 4, 5]
        codeflash_output = _calculate_percentile(data, 0.0) # 1.93μs -> 1.65μs (16.5% faster)
        codeflash_output = _calculate_percentile(data, 1.0) # 538ns -> 573ns (6.11% slower)

    def test_percentile_middle(self):
        """Test that 0.5 percentile returns the median for odd and even length"""
        odd_data = [1, 2, 3, 4, 5]
        even_data = [1, 2, 3, 4]
        # For odd: median is 3
        codeflash_output = _calculate_percentile(odd_data, 0.5) # 1.79μs -> 1.62μs (10.0% faster)
        # For even: median is (2 + 3)/2 = 2.5
        codeflash_output = _calculate_percentile(even_data, 0.5) # 600ns -> 555ns (8.11% faster)

    def test_quartiles(self):
        """Test 25th and 75th percentiles on sample data"""
        data = [10, 20, 30, 40, 50]
        # 25th percentile: index = 0.25 * 4 = 1.0, so data[1] = 20
        codeflash_output = _calculate_percentile(data, 0.25) # 1.90μs -> 1.58μs (19.9% faster)
        # 75th percentile: index = 0.75 * 4 = 3.0, so data[3] = 40
        codeflash_output = _calculate_percentile(data, 0.75) # 599ns -> 545ns (9.91% faster)

    # --- Edge Test Cases ---

    def test_percentile_exact_index(self):
        """Test percentiles that hit exact indices (no interpolation)"""
        data = [0, 10, 20, 30, 40]
        # 40th percentile: index = 0.4*4 = 1.6, so interpolate between data[1]=10 and data[2]=20
        expected = 10 * (1 - 0.6) + 20 * 0.6  # 10*0.4 + 20*0.6 = 4 + 12 = 16
        codeflash_output = _calculate_percentile(data, 0.4) # 1.91μs -> 1.58μs (20.3% faster)

    def test_percentile_requires_interpolation(self):
        """Test percentiles that require interpolation between two values"""
        data = [100, 200, 300, 400]
        # 25th percentile: index = 0.25*3 = 0.75, interpolate between 100 and 200
        expected = 100 * (1 - 0.75) + 200 * 0.75  # 25 + 150 = 175
        codeflash_output = _calculate_percentile(data, 0.25) # 1.85μs -> 1.63μs (13.5% faster)

    def test_percentile_near_zero(self):
        """Test percentile very close to 0"""
        data = [1, 2, 3, 4, 5]
        # percentile=0.0001, index=0.0004, interpolate between 1 and 2
        expected = 1 * (1 - 0.0004) + 2 * 0.0004
        codeflash_output = _calculate_percentile(data, 0.0001) # 1.95μs -> 1.67μs (17.1% faster)

    def test_percentile_near_one(self):
        """Test percentile very close to 1"""
        data = [1, 2, 3, 4, 5]
        # percentile=0.9999, index=3.9996, interpolate between 4 and 5
        expected = 4 * (1 - 0.9996) + 5 * 0.9996
        codeflash_output = _calculate_percentile(data, 0.9999) # 1.85μs -> 1.59μs (16.2% faster)

    def test_percentile_out_of_range(self):
        """Test percentiles less than 0 and greater than 1 (should clamp to bounds)"""
        data = [1, 2, 3, 4, 5]
        # Negative percentile returns first element
        codeflash_output = _calculate_percentile(data, -0.5) # 1.96μs -> 1.68μs (16.8% faster)
        # Percentile > 1 returns last element
        codeflash_output = _calculate_percentile(data, 1.5) # 558ns -> 556ns (0.360% faster)

    def test_unsorted_input(self):
        """Test that function expects sorted input (behavior not defined for unsorted)"""
        data = [10, 1, 5, 3]
        # The function does not sort input, so percentile is based on input order
        codeflash_output = _calculate_percentile(data, 0.5); result = codeflash_output # 1.81μs -> 1.58μs (14.2% faster)

    def test_identical_elements(self):
        """Test that all identical elements always return that value"""
        data = [7, 7, 7, 7, 7]
        for p in [0.0, 0.25, 0.5, 0.75, 1.0]:
            codeflash_output = _calculate_percentile(data, p) # 3.56μs -> 3.29μs (8.33% faster)

    def test_negative_and_zero_values(self):
        """Test data with negative and zero values"""
        data = [-10, 0, 10, 20]
        # 50th percentile: index=1.5, interpolate between 0 and 10
        expected = 0.5 * 0 + 0.5 * 10
        codeflash_output = _calculate_percentile(data, 0.5) # 1.85μs -> 1.64μs (12.7% faster)

    # --- Large Scale Test Cases ---

    def test_large_data_set(self):
        """Test with a large sorted data set (1000 elements)"""
        data = list(range(1000))  # [0, 1, ..., 999]
        # 90th percentile: index = 0.9*999 = 899.1, interpolate between 899 and 900
        expected = 899 * (1 - 0.1) + 900 * 0.1  # 899*0.9 + 900*0.1 = 809.1 + 90 = 899.1
        codeflash_output = _calculate_percentile(data, 0.9) # 2.04μs -> 1.85μs (10.1% faster)

    def test_large_data_set_first_and_last(self):
        """Test 0th and 100th percentile on large data set"""
        data = [i * 2 for i in range(1000)]  # [0, 2, 4, ..., 1998]
        codeflash_output = _calculate_percentile(data, 0.0) # 2.06μs -> 1.73μs (19.0% faster)
        codeflash_output = _calculate_percentile(data, 1.0) # 693ns -> 714ns (2.94% slower)

    def test_large_data_set_interpolation(self):
        """Test interpolation on large data set with floats"""
        data = [i * 0.5 for i in range(1000)]  # [0.0, 0.5, ..., 499.5]
        # 33rd percentile: index = 0.33*999 = 329.67, interpolate between 329 and 330
        lower = 329 * 0.5
        upper = 330 * 0.5
        weight = 0.67
        expected = lower * (1 - weight) + upper * weight
        codeflash_output = _calculate_percentile(data, 0.33) # 1.57μs -> 1.49μs (5.65% faster)

    def test_large_data_set_all_identical(self):
        """Test large data set with all identical values"""
        data = [123.456] * 1000
        codeflash_output = _calculate_percentile(data, 0.5) # 1.76μs -> 1.55μs (13.4% faster)
        codeflash_output = _calculate_percentile(data, 0.99) # 554ns -> 526ns (5.32% faster)

    def test_large_data_set_negative_values(self):
        """Test large data set with all negative values"""
        data = [-i for i in range(1000)]  # [0, -1, -2, ..., -999]
        data.sort()  # Ensure sorted order: [-999, ..., -1, 0]
        # 75th percentile: index = 0.75*999 = 749.25, interpolate between 749th and 750th
        lower = data[749]
        upper = data[750]
        weight = 0.25
        expected = lower * (1 - weight) + upper * weight
        codeflash_output = _calculate_percentile(data, 0.75) # 1.53μs -> 1.60μs (4.68% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from mlflow.tracing.utils.__init__ import _calculate_percentile

def test__calculate_percentile():
    _calculate_percentile([0.0, 0.0], 0.0)

def test__calculate_percentile_2():
    _calculate_percentile([0.0], 0.0)

def test__calculate_percentile_3():
    _calculate_percentile([], 0.0)

To edit these changes git checkout codeflash/optimize-_calculate_percentile-mhwzg0wi and push.