⚡️ Speed up function apply_function by 1,957%

[codeflash-ai]

📄 1,957% (19.57x) speedup for apply_function in src/numpy_pandas/dataframe_operations.py

⏱️ Runtime : 48.4 milliseconds → 2.35 milliseconds (best of 857 runs)

📝 Explanation and details

The optimization replaces a manual row-by-row iteration with pandas' vectorized map() operation, resulting in a dramatic ~20x speedup.

Key changes:

• Eliminated df.iloc[i][column] access pattern: The original code uses df.iloc[i][column] inside a loop, which is extremely inefficient. Each iloc call triggers pandas' positional indexing machinery, creating significant overhead for every row access.
• Leveraged vectorized operations: The optimized version uses df[column].map(func), which operates directly on the pandas Series using optimized C code paths rather than Python iteration.
• Removed explicit loop and list building: Instead of manually appending to a result list, the operation is performed in a single vectorized call and converted to a list at the end.

Why this is faster:
The original implementation has O(n) calls to iloc, each with substantial overhead for index resolution and type checking. The line profiler shows that df.iloc[i][column] consumes 98% of the execution time (396ms out of 405ms total). In contrast, Series.map() leverages pandas' internal optimizations and vectorized operations, processing the entire column at once with minimal per-element overhead.

Performance characteristics by test case:

• Small DataFrames (3-4 rows): Modest improvements or slight regressions due to vectorization overhead being comparable to the small dataset size
• Large DataFrames (1000+ rows): Massive speedups (2600-4200% faster) where vectorization truly shines, as the fixed overhead is amortized across many elements
• Edge cases: Consistent behavior with slightly better performance for exception handling scenarios due to faster failure detection

The optimization is most effective for medium to large datasets where the vectorization benefits outweigh the setup costs.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Callable, List

import pandas as pd
# imports
import pytest  # used for our unit tests
from src.numpy_pandas.dataframe_operations import apply_function

# unit tests

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

def test_apply_function_basic_integers_square():
    # Test squaring integers in a column
    df = pd.DataFrame({'a': [1, 2, 3, 4]})
    codeflash_output = apply_function(df, 'a', lambda x: x * x); result = codeflash_output # 32.3μs -> 29.6μs (9.15% faster)

def test_apply_function_basic_strings_upper():
    # Test converting strings to uppercase
    df = pd.DataFrame({'name': ['alice', 'bob', 'carol']})
    codeflash_output = apply_function(df, 'name', lambda x: x.upper()); result = codeflash_output # 26.6μs -> 30.8μs (13.6% slower)

def test_apply_function_basic_mixed_types():
    # Test with mixed types and a string conversion function
    df = pd.DataFrame({'val': [1, 'a', 3.5, None]})
    codeflash_output = apply_function(df, 'val', str); result = codeflash_output # 31.8μs -> 31.4μs (1.20% faster)

def test_apply_function_basic_identity():
    # Test with identity function
    df = pd.DataFrame({'x': [10, 20, 30]})
    codeflash_output = apply_function(df, 'x', lambda x: x); result = codeflash_output # 26.7μs -> 29.3μs (8.96% slower)

def test_apply_function_basic_constant():
    # Test with a function that returns a constant
    df = pd.DataFrame({'foo': [1, 2, 3]})
    codeflash_output = apply_function(df, 'foo', lambda x: 42); result = codeflash_output # 26.5μs -> 29.2μs (9.27% slower)

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

def test_apply_function_empty_dataframe():
    # Test with an empty DataFrame
    df = pd.DataFrame({'a': []})
    codeflash_output = apply_function(df, 'a', lambda x: x); result = codeflash_output # 500ns -> 26.7μs (98.1% slower)

def test_apply_function_column_not_found():
    # Test with a non-existent column
    df = pd.DataFrame({'a': [1, 2, 3]})
    with pytest.raises(KeyError):
        apply_function(df, 'b', lambda x: x) # 15.2μs -> 9.79μs (54.9% faster)

def test_apply_function_nan_values():
    # Test with NaN values in the column
    import math
    df = pd.DataFrame({'val': [1, float('nan'), 3]})
    codeflash_output = apply_function(df, 'val', lambda x: 0 if isinstance(x, float) and math.isnan(x) else x); result = codeflash_output # 27.1μs -> 29.8μs (9.08% slower)

def test_apply_function_none_values():
    # Test with None values in the column
    df = pd.DataFrame({'val': [None, 2, None]})
    codeflash_output = apply_function(df, 'val', lambda x: 99 if x is None else x); result = codeflash_output # 26.3μs -> 29.2μs (9.97% slower)

def test_apply_function_column_with_all_same_values():
    # Test with all values being the same
    df = pd.DataFrame({'col': [7, 7, 7, 7]})
    codeflash_output = apply_function(df, 'col', lambda x: x + 1); result = codeflash_output # 31.8μs -> 29.5μs (7.78% faster)

def test_apply_function_column_with_all_nan():
    # Test with all values as NaN
    import math
    df = pd.DataFrame({'col': [float('nan')] * 5})
    codeflash_output = apply_function(df, 'col', lambda x: 0 if isinstance(x, float) and math.isnan(x) else x); result = codeflash_output # 36.6μs -> 30.0μs (21.9% faster)

def test_apply_function_column_with_lists():
    # Test with lists as elements
    df = pd.DataFrame({'col': [[1, 2], [3, 4], []]})
    codeflash_output = apply_function(df, 'col', len); result = codeflash_output # 26.6μs -> 29.2μs (8.86% slower)

def test_apply_function_column_with_dicts():
    # Test with dicts as elements
    df = pd.DataFrame({'col': [{'a': 1}, {}, {'b': 2, 'c': 3}]})
    codeflash_output = apply_function(df, 'col', lambda d: sorted(d.keys())); result = codeflash_output # 27.1μs -> 31.0μs (12.6% slower)

def test_apply_function_function_raises():
    # Test when the function raises an exception
    df = pd.DataFrame({'col': [1, 2, 3]})
    def bad_func(x):
        if x == 2:
            raise ValueError("bad value")
        return x
    with pytest.raises(ValueError):
        apply_function(df, 'col', bad_func) # 21.7μs -> 16.2μs (33.3% faster)

def test_apply_function_column_with_boolean():
    # Test with boolean values
    df = pd.DataFrame({'flag': [True, False, True]})
    codeflash_output = apply_function(df, 'flag', lambda x: not x); result = codeflash_output # 26.4μs -> 29.6μs (10.8% slower)

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

def test_apply_function_large_dataframe_sum():
    # Test with a large DataFrame and sum function
    N = 1000
    df = pd.DataFrame({'a': list(range(N))})
    codeflash_output = apply_function(df, 'a', lambda x: x + 1); result = codeflash_output # 4.77ms -> 167μs (2742% faster)

def test_apply_function_large_dataframe_strings():
    # Test with a large DataFrame of strings
    N = 1000
    df = pd.DataFrame({'s': [str(i) for i in range(N)]})
    codeflash_output = apply_function(df, 's', lambda x: x.zfill(4)); result = codeflash_output # 4.75ms -> 113μs (4072% faster)

def test_apply_function_large_dataframe_mixed_types():
    # Test with a large DataFrame with alternating types
    N = 1000
    data = [i if i % 2 == 0 else str(i) for i in range(N)]
    df = pd.DataFrame({'mix': data})
    codeflash_output = apply_function(df, 'mix', lambda x: str(x) + 'X'); result = codeflash_output # 4.76ms -> 117μs (3963% faster)

def test_apply_function_large_dataframe_all_none():
    # Test with a large DataFrame with all None values
    N = 1000
    df = pd.DataFrame({'col': [None] * N})
    codeflash_output = apply_function(df, 'col', lambda x: 0 if x is None else x); result = codeflash_output # 4.70ms -> 152μs (2984% faster)

def test_apply_function_large_dataframe_with_nan_and_none():
    # Test with a large DataFrame with alternating NaN and None
    import math
    N = 1000
    data = [float('nan') if i % 2 == 0 else None for i in range(N)]
    df = pd.DataFrame({'col': data})
    codeflash_output = apply_function(df, 'col', lambda x: 1 if (isinstance(x, float) and math.isnan(x)) or x is None else x); result = codeflash_output # 4.78ms -> 174μs (2635% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import Any, Callable, List

import pandas as pd
# imports
import pytest  # used for our unit tests
from src.numpy_pandas.dataframe_operations import apply_function

# unit tests

# -------------------- BASIC TEST CASES --------------------

def test_apply_function_basic_integers():
    # Test applying a function to a column of integers
    df = pd.DataFrame({'a': [1, 2, 3, 4]})
    codeflash_output = apply_function(df, 'a', lambda x: x * 2); result = codeflash_output # 32.0μs -> 29.6μs (8.02% faster)

def test_apply_function_basic_strings():
    # Test applying a function to a column of strings
    df = pd.DataFrame({'name': ['Alice', 'Bob', 'Charlie']})
    codeflash_output = apply_function(df, 'name', lambda x: x.upper()); result = codeflash_output # 26.7μs -> 30.8μs (13.1% slower)

def test_apply_function_basic_floats():
    # Test applying a function to a column of floats
    df = pd.DataFrame({'x': [0.5, 1.5, 2.5]})
    codeflash_output = apply_function(df, 'x', lambda x: round(x + 0.1, 1)); result = codeflash_output # 34.8μs -> 30.1μs (15.8% faster)

def test_apply_function_basic_bool():
    # Test applying a function to a column of booleans
    df = pd.DataFrame({'flag': [True, False, True]})
    codeflash_output = apply_function(df, 'flag', lambda x: not x); result = codeflash_output # 26.6μs -> 29.5μs (9.88% slower)

def test_apply_function_basic_identity():
    # Test applying identity function
    df = pd.DataFrame({'val': [10, 20, 30]})
    codeflash_output = apply_function(df, 'val', lambda x: x); result = codeflash_output # 26.5μs -> 29.4μs (9.93% slower)

# -------------------- EDGE TEST CASES --------------------

def test_apply_function_empty_dataframe():
    # Test with empty DataFrame
    df = pd.DataFrame({'a': []})
    codeflash_output = apply_function(df, 'a', lambda x: x); result = codeflash_output # 500ns -> 26.8μs (98.1% slower)

def test_apply_function_column_not_found():
    # Test with non-existent column
    df = pd.DataFrame({'a': [1, 2, 3]})
    with pytest.raises(KeyError):
        apply_function(df, 'b', lambda x: x) # 15.2μs -> 9.79μs (55.3% faster)

def test_apply_function_with_none_values():
    # Test with None values in the column
    df = pd.DataFrame({'val': [1, None, 3]})
    codeflash_output = apply_function(df, 'val', lambda x: 0 if x is None else x); result = codeflash_output # 26.7μs -> 29.5μs (9.61% slower)

def test_apply_function_with_nan_values():
    # Test with NaN values in the column
    import math
    df = pd.DataFrame({'val': [1.0, float('nan'), 3.0]})
    codeflash_output = apply_function(df, 'val', lambda x: 0 if isinstance(x, float) and math.isnan(x) else x); result = codeflash_output # 26.9μs -> 29.5μs (8.90% slower)

def test_apply_function_with_mixed_types():
    # Test with mixed types in the column
    df = pd.DataFrame({'mixed': [1, 'a', None, 3.5]})
    codeflash_output = apply_function(df, 'mixed', lambda x: str(x)); result = codeflash_output # 31.8μs -> 31.0μs (2.69% faster)

def test_apply_function_function_raises_exception():
    # Test if function raises an exception for a specific value
    df = pd.DataFrame({'x': [1, 2, 0, 4]})
    def safe_inverse(x):
        if x == 0:
            raise ValueError("Division by zero")
        return 1 / x
    with pytest.raises(ValueError):
        apply_function(df, 'x', safe_inverse) # 27.3μs -> 16.4μs (66.5% faster)

def test_apply_function_on_non_series_column():
    # Test with a DataFrame with multiple columns, ensure correct column is used
    df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
    codeflash_output = apply_function(df, 'b', lambda x: x + 10); result = codeflash_output # 26.8μs -> 29.3μs (8.52% slower)

def test_apply_function_with_duplicate_column_names():
    # Test with DataFrame with duplicate column names (should use the first occurrence)
    df = pd.DataFrame([[1, 2], [3, 4]], columns=['a', 'a'])
    # Pandas allows duplicate column names, but iloc will still work column-wise
    codeflash_output = apply_function(df, 'a', lambda x: x * 2); result = codeflash_output # 82.4μs -> 94.6μs (12.9% slower)

def test_apply_function_column_all_nones():
    # Test where the column is all None
    df = pd.DataFrame({'a': [None, None, None]})
    codeflash_output = apply_function(df, 'a', lambda x: 1 if x is None else 0); result = codeflash_output # 26.8μs -> 29.2μs (8.42% slower)

def test_apply_function_column_all_nans():
    # Test where the column is all NaN
    import math
    df = pd.DataFrame({'a': [float('nan')] * 3})
    codeflash_output = apply_function(df, 'a', lambda x: 1 if isinstance(x, float) and math.isnan(x) else 0); result = codeflash_output # 26.6μs -> 29.7μs (10.3% slower)

def test_apply_function_empty_column_name():
    # Test with empty string as column name
    df = pd.DataFrame({'': [1, 2, 3]})
    codeflash_output = apply_function(df, '', lambda x: x + 1); result = codeflash_output # 26.5μs -> 29.2μs (9.00% slower)

def test_apply_function_with_indexed_dataframe():
    # Test with DataFrame with custom index
    df = pd.DataFrame({'a': [5, 6, 7]}, index=['x', 'y', 'z'])
    codeflash_output = apply_function(df, 'a', lambda x: x * 3); result = codeflash_output # 26.9μs -> 28.8μs (6.79% slower)

# -------------------- LARGE SCALE TEST CASES --------------------

def test_apply_function_large_dataframe():
    # Test with a large DataFrame (1000 rows)
    df = pd.DataFrame({'num': list(range(1000))})
    codeflash_output = apply_function(df, 'num', lambda x: x + 1); result = codeflash_output # 4.77ms -> 167μs (2746% faster)

def test_apply_function_large_strings():
    # Test with a large DataFrame of strings (1000 rows)
    df = pd.DataFrame({'s': ['test'] * 1000})
    codeflash_output = apply_function(df, 's', lambda x: x + '_done'); result = codeflash_output # 4.73ms -> 109μs (4227% faster)

def test_apply_function_large_mixed_types():
    # Test with a large DataFrame of mixed types
    data = [i if i % 2 == 0 else str(i) for i in range(1000)]
    df = pd.DataFrame({'mixed': data})
    codeflash_output = apply_function(df, 'mixed', lambda x: str(x) + '_x'); result = codeflash_output # 4.76ms -> 117μs (3968% faster)

def test_apply_function_large_with_nones():
    # Test with a large DataFrame with some None values
    data = [i if i % 10 != 0 else None for i in range(1000)]
    df = pd.DataFrame({'val': data})
    codeflash_output = apply_function(df, 'val', lambda x: 0 if x is None else x); result = codeflash_output # 4.77ms -> 118μs (3916% faster)

def test_apply_function_large_performance():
    # Test with a large DataFrame to check for reasonable performance (not timing, but correctness)
    df = pd.DataFrame({'a': range(1000)})
    codeflash_output = apply_function(df, 'a', lambda x: x * x); result = codeflash_output # 4.76ms -> 170μs (2690% 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-apply_function-mdperii7 and push.