⚡️ Speed up function mean by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for mean in pandas/core/array_algos/masked_reductions.py

⏱️ Runtime : 1.92 milliseconds → 1.57 milliseconds (best of 83 runs)

📝 Explanation and details

The optimized code achieves a 22% speedup by adding a crucial fast-path optimization that eliminates redundant work for arrays with no missing values.

Key optimization: Early exit for mask.any() == False

The primary improvement is adding a fast-path check if not mask.any(): at the beginning. When there are no masked values (a very common case), the optimized code:

• Immediately calls func(values, axis=axis, **kwargs) after checking min_count
• Avoids the expensive where=~mask parameter in numpy operations
• Skips the object dtype check and other conditional logic

Why this matters:

• The original code always computed where=~mask even when no values were masked, forcing numpy to process the mask unnecessarily
• Line profiler shows the return func(values, where=~mask, axis=axis, **kwargs) line took 84% of execution time in the original vs only 47.3% in the optimized version
• The fast-path optimization shows dramatic improvements in test cases with no missing values (67-94% faster)

Performance characteristics:

• Best case: Arrays with no missing values see 67-94% speedup (most test cases)
• Neutral: Arrays with all missing values see minimal change (2-6% improvement)
• Slight regression: Arrays with partial missing values are 2-11% slower due to additional mask.any() check, but this is offset by the common case gains

The optimization is particularly effective because mask.any() is a highly optimized numpy operation that can short-circuit, making the additional check very cheap compared to the avoided where parameter overhead.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from pandas._libs import missing as libmissing
from pandas.core.array_algos.masked_reductions import mean

# unit tests

# Basic Test Cases

def test_mean_basic_positive_integers():
    # Test mean of positive integers with no missing values
    arr = np.array([1, 2, 3, 4, 5])
    mask = np.array([False, False, False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 44.1μs -> 23.2μs (90.2% faster)

def test_mean_basic_negative_integers():
    # Test mean of negative integers with no missing values
    arr = np.array([-1, -2, -3, -4, -5])
    mask = np.array([False]*5)
    codeflash_output = mean(arr, mask); result = codeflash_output # 41.8μs -> 23.0μs (81.3% faster)

def test_mean_basic_mixed_integers():
    # Test mean of mixed positive and negative integers
    arr = np.array([-2, 0, 2])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 39.2μs -> 22.1μs (77.2% faster)

def test_mean_basic_floats():
    # Test mean of floats
    arr = np.array([1.5, 2.5, 3.5])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 38.4μs -> 21.4μs (79.7% faster)

def test_mean_basic_with_masked_values():
    # Test mean with some masked (missing) values
    arr = np.array([1, 2, 3, 4, 5])
    mask = np.array([False, True, False, True, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 37.9μs -> 45.5μs (16.7% slower)

def test_mean_basic_all_masked_with_skipna_false():
    # Test mean with all values masked and skipna=False
    arr = np.array([1, 2, 3])
    mask = np.array([True, True, True])
    codeflash_output = mean(arr, mask, skipna=False); result = codeflash_output # 6.92μs -> 6.75μs (2.50% faster)

def test_mean_basic_all_masked_with_skipna_true():
    # Test mean with all values masked and skipna=True (default)
    arr = np.array([1, 2, 3])
    mask = np.array([True, True, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.77μs -> 6.52μs (3.83% faster)

def test_mean_basic_empty_array():
    # Test mean with empty array
    arr = np.array([])
    mask = np.array([], dtype=bool)
    codeflash_output = mean(arr, mask); result = codeflash_output # 996ns -> 881ns (13.1% faster)

def test_mean_basic_single_element():
    # Test mean with single element
    arr = np.array([42])
    mask = np.array([False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 49.1μs -> 28.8μs (70.1% faster)

def test_mean_basic_single_masked_element():
    # Test mean with single masked element
    arr = np.array([42])
    mask = np.array([True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.70μs -> 6.57μs (1.98% faster)

# Edge Test Cases

def test_mean_edge_all_zeros():
    # Test mean with all zeros
    arr = np.zeros(5)
    mask = np.array([False]*5)
    codeflash_output = mean(arr, mask); result = codeflash_output # 43.4μs -> 23.9μs (81.7% faster)

def test_mean_edge_some_zeros_and_masked():
    # Test mean with zeros and masked values
    arr = np.array([0, 0, 0, 0, 0])
    mask = np.array([True, False, True, False, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 41.2μs -> 46.0μs (10.3% slower)

def test_mean_edge_object_dtype():
    # Test mean with object dtype and masked values
    arr = np.array([1, 2, 3, None], dtype=object)
    mask = np.array([False, False, False, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 27.9μs -> 29.4μs (5.12% slower)

def test_mean_edge_axis_argument_1d():
    # Test axis argument with 1D array (should behave normally)
    arr = np.array([10, 20, 30])
    mask = np.array([False, True, False])
    codeflash_output = mean(arr, mask, axis=0); result = codeflash_output # 44.3μs -> 45.5μs (2.82% slower)

def test_mean_edge_axis_argument_2d():
    # Test axis argument with 2D array
    arr = np.array([[1, 2, 3], [4, 5, 6]])
    mask = np.array([[False, True, False], [False, False, True]])
    # mean along axis 0 (columns)
    codeflash_output = mean(arr, mask, axis=0); result = codeflash_output # 49.8μs -> 49.3μs (0.967% faster)



def test_mean_edge_mask_all_false():
    # Test with mask all False (no missing values)
    arr = np.array([1, 2, 3])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 56.1μs -> 35.5μs (58.1% faster)


def test_mean_edge_large_negative_and_positive():
    # Test with large negative and positive values
    arr = np.array([-1e10, 1e10])
    mask = np.array([False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 55.5μs -> 32.9μs (68.5% faster)

def test_mean_edge_nan_in_object_dtype():
    # Test with np.nan in object dtype
    arr = np.array([1, np.nan, 3], dtype=object)
    mask = np.array([False, True, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 29.4μs -> 32.1μs (8.49% slower)

# Large Scale Test Cases

def test_mean_large_scale_all_valid():
    # Test mean with large array, all valid values
    arr = np.arange(1000)
    mask = np.zeros(1000, dtype=bool)
    codeflash_output = mean(arr, mask); result = codeflash_output # 47.7μs -> 27.4μs (74.5% faster)

def test_mean_large_scale_half_masked():
    # Test mean with large array, half masked
    arr = np.arange(1000)
    mask = np.zeros(1000, dtype=bool)
    mask[::2] = True  # mask every other value
    valid = arr[1::2]  # values not masked
    codeflash_output = mean(arr, mask); result = codeflash_output # 49.0μs -> 53.8μs (9.05% slower)

def test_mean_large_scale_all_masked():
    # Test mean with large array, all masked
    arr = np.arange(1000)
    mask = np.ones(1000, dtype=bool)
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.60μs -> 6.23μs (5.89% faster)





#------------------------------------------------
from __future__ import annotations

import numpy as np
# imports
import pytest  # used for our unit tests
from pandas.core.array_algos.masked_reductions import mean


# Minimal stub for libmissing.NA (since we can't import pandas._libs)
class _LibMissing:
    NA = None  # We'll use None to represent NA for testing

libmissing = _LibMissing()
from pandas.core.array_algos.masked_reductions import mean

# ------------------ UNIT TESTS ------------------

# Basic Test Cases

def test_mean_basic_no_missing():
    # Test with no missing values, integer array
    arr = np.array([1, 2, 3, 4, 5])
    mask = np.array([False, False, False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 51.3μs -> 30.6μs (67.8% faster)

def test_mean_basic_float():
    # Test with no missing values, float array
    arr = np.array([1.0, 2.0, 3.0])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 42.4μs -> 22.8μs (85.5% faster)

def test_mean_basic_with_missing():
    # Test with some missing values
    arr = np.array([1, 2, 3, 4, 5])
    mask = np.array([False, True, False, False, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 40.7μs -> 45.9μs (11.3% slower)

def test_mean_basic_all_missing():
    # All values are missing
    arr = np.array([1, 2, 3])
    mask = np.array([True, True, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.97μs -> 6.56μs (6.28% faster)

def test_mean_basic_empty_array():
    # Empty array
    arr = np.array([])
    mask = np.array([])
    codeflash_output = mean(arr, mask); result = codeflash_output # 950ns -> 813ns (16.9% faster)

def test_mean_basic_object_dtype():
    # Object dtype, no missing
    arr = np.array([1, 2, 3], dtype=object)
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 31.7μs -> 27.2μs (16.8% faster)

def test_mean_basic_object_dtype_with_missing():
    # Object dtype, with missing
    arr = np.array([1, 2, 3], dtype=object)
    mask = np.array([True, False, True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 27.3μs -> 29.8μs (8.16% slower)

def test_mean_basic_negative_numbers():
    # Negative numbers
    arr = np.array([-5, -10, -15])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 45.8μs -> 25.0μs (83.2% faster)

def test_mean_basic_mixed_sign():
    # Mixed positive and negative
    arr = np.array([-2, 0, 2])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 40.9μs -> 23.2μs (76.6% faster)

# Edge Test Cases

def test_mean_edge_single_element():
    # Single element, not missing
    arr = np.array([42])
    mask = np.array([False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 40.5μs -> 22.6μs (78.7% faster)

def test_mean_edge_single_element_missing():
    # Single element, missing
    arr = np.array([42])
    mask = np.array([True])
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.55μs -> 6.49μs (0.955% faster)

def test_mean_edge_all_missing_with_min_count():
    # All missing, min_count > 0
    arr = np.array([1, 2, 3])
    mask = np.array([True, True, True])
    codeflash_output = mean(arr, mask, skipna=True); result = codeflash_output # 7.00μs -> 6.81μs (2.75% faster)



def test_mean_edge_skipna_false_with_missing():
    # skipna=False, any missing value should return NA
    arr = np.array([1, 2, 3])
    mask = np.array([False, True, False])
    codeflash_output = mean(arr, mask, skipna=False); result = codeflash_output # 17.2μs -> 16.1μs (6.76% faster)

def test_mean_edge_skipna_false_no_missing():
    # skipna=False, no missing
    arr = np.array([1, 2, 3])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask, skipna=False); result = codeflash_output # 31.9μs -> 29.7μs (7.41% faster)

def test_mean_edge_nan_in_data():
    # np.nan in data, not marked missing
    arr = np.array([1.0, np.nan, 3.0])
    mask = np.array([False, False, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 45.5μs -> 23.4μs (94.4% faster)

def test_mean_edge_nan_in_data_marked_missing():
    # np.nan in data, marked missing
    arr = np.array([1.0, np.nan, 3.0])
    mask = np.array([False, True, False])
    codeflash_output = mean(arr, mask); result = codeflash_output # 40.8μs -> 45.0μs (9.41% slower)

def test_mean_edge_all_zeros():
    # All zeros
    arr = np.zeros(5)
    mask = np.array([False]*5)
    codeflash_output = mean(arr, mask); result = codeflash_output # 39.3μs -> 22.1μs (78.2% faster)

def test_mean_edge_axis_0():
    # Test axis=0 on 2D array
    arr = np.array([[1, 2], [3, 4]])
    mask = np.array([[False, False], [False, True]])
    codeflash_output = mean(arr, mask, axis=0); result = codeflash_output # 49.1μs -> 51.8μs (5.22% slower)

def test_mean_edge_axis_1():
    # Test axis=1 on 2D array
    arr = np.array([[1, 2], [3, 4]])
    mask = np.array([[False, True], [False, False]])
    codeflash_output = mean(arr, mask, axis=1); result = codeflash_output # 47.0μs -> 46.6μs (1.05% faster)


def test_mean_large_scale_no_missing():
    # Large array, no missing
    arr = np.arange(1000)
    mask = np.zeros(1000, dtype=bool)
    codeflash_output = mean(arr, mask); result = codeflash_output # 59.8μs -> 36.5μs (63.6% faster)

def test_mean_large_scale_half_missing():
    # Large array, half missing
    arr = np.arange(1000)
    mask = np.array([i%2==0 for i in range(1000)])
    codeflash_output = mean(arr, mask); result = codeflash_output # 51.8μs -> 55.0μs (5.81% slower)
    # Only odd indices are counted
    codeflash_output = np.mean(arr[~mask]); expected = codeflash_output # 10.4μs -> 9.57μs (9.17% faster)

def test_mean_large_scale_all_missing():
    # Large array, all missing
    arr = np.arange(1000)
    mask = np.ones(1000, dtype=bool)
    codeflash_output = mean(arr, mask); result = codeflash_output # 6.79μs -> 6.82μs (0.352% slower)

def test_mean_large_scale_random_missing():
    # Large array, random missing
    rng = np.random.default_rng(42)
    arr = rng.integers(0, 10000, size=1000)
    mask = rng.choice([False, True], size=1000, p=[0.8, 0.2])
    expected = np.mean(arr[~mask]) if (~mask).sum() > 0 else libmissing.NA # 14.5μs -> 14.4μs (0.347% faster)
    codeflash_output = mean(arr, mask); result = codeflash_output # 37.3μs -> 38.3μs (2.76% slower)
    if expected is libmissing.NA:
        pass
    else:
        pass

def test_mean_large_scale_axis():
    # 2D large array, axis=1
    arr = np.arange(2000).reshape(1000,2)
    mask = np.zeros((1000,2), dtype=bool)
    codeflash_output = mean(arr, mask, axis=1); result = codeflash_output # 86.9μs -> 46.5μs (86.9% faster)
    codeflash_output = np.mean(arr, axis=1); expected = codeflash_output # 26.3μs -> 22.6μs (16.6% faster)

def test_mean_large_scale_object_dtype():
    # Large array, object dtype, some missing
    arr = np.arange(1000, dtype=object)
    mask = np.array([i%3==0 for i in range(1000)])
    codeflash_output = np.mean(arr[~mask]); expected = codeflash_output # 22.3μs -> 21.4μs (3.97% faster)
    codeflash_output = mean(arr, mask); result = codeflash_output # 23.2μs -> 25.7μs (9.70% slower)
# 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-mean-mho7wi26 and push.