⚡️ Speed up function pivot_table by 2,181%

[codeflash-ai]

📄 2,181% (21.81x) speedup for pivot_table in src/numpy_pandas/dataframe_operations.py

⏱️ Runtime : 35.9 milliseconds → 1.57 milliseconds (best of 436 runs)

📝 Explanation and details

The optimization achieves a 2180% speedup by eliminating the most expensive operation in the original code: repeatedly calling df.iloc[i] to access DataFrame rows.

Key Optimization: Vectorized Column Extraction

The critical change replaces the inefficient row-by-row DataFrame access:

# Original: Expensive row access (71.1% of total time)
for i in range(len(df)):
    row = df.iloc[i]  # This line alone took 244ms out of 344ms total
    index_val = row[index]
    column_val = row[columns] 
    value = row[values]

With direct NumPy array extraction and zip iteration:

# Optimized: Extract entire columns as arrays once
index_arr = df[index].values      # 2.4ms
columns_arr = df[columns].values  # 1.3ms  
values_arr = df[values].values    # 1.3ms

# Then iterate over arrays directly
for index_val, column_val, value in zip(index_arr, columns_arr, values_arr):

Why This Works

1. DataFrame.iloc[i] is extremely slow - it creates a new Series object for each row access and involves significant pandas overhead for indexing operations
2. Array access is fast - NumPy arrays provide direct memory access with minimal overhead
3. Bulk extraction is efficient - Getting entire columns at once leverages pandas' optimized column operations

Performance Impact by Test Case

The optimization excels across all test scenarios:

• Large-scale tests see massive gains: 3543-6406% speedup for datasets with 1000+ rows
• Medium datasets (100-900 rows): 1560-5350% speedup
• Small datasets: 57-129% speedup
• Edge cases: Generally 19-92% faster, though very small datasets (single row, empty) show minimal or slightly negative impact due to the overhead of array extraction

The optimization is particularly effective for scenarios with many rows since it eliminates the O(n) DataFrame row access overhead, making the algorithm scale much better with dataset size.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any

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

# unit tests

# ---------------------
# 1. BASIC TEST CASES
# ---------------------

def test_basic_mean_aggregation():
    # Test mean aggregation with simple data
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar", "bar"],
        "B": ["one", "two", "one", "two"],
        "C": [1, 2, 3, 4]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 52.5μs -> 27.2μs (92.7% faster)

def test_basic_sum_aggregation():
    # Test sum aggregation with repeated groups
    df = pd.DataFrame({
        "A": ["foo", "foo", "foo", "bar"],
        "B": ["one", "one", "two", "two"],
        "C": [1, 2, 3, 4]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 51.6μs -> 26.9μs (91.6% faster)

def test_basic_count_aggregation():
    # Test count aggregation
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar", "bar", "bar"],
        "B": ["one", "one", "two", "one", "two"],
        "C": [10, 20, 30, 40, 50]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="count"); result = codeflash_output # 60.0μs -> 26.7μs (125% faster)

def test_basic_different_types():
    # Test with non-string index and column values
    df = pd.DataFrame({
        "A": [1, 1, 2, 2],
        "B": [True, False, True, False],
        "C": [10, 20, 30, 40]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 54.0μs -> 26.6μs (103% faster)

# ---------------------
# 2. EDGE TEST CASES
# ---------------------

def test_empty_dataframe():
    # Test with an empty DataFrame
    df = pd.DataFrame(columns=["A", "B", "C"])
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 750ns -> 26.0μs (97.1% slower)

def test_single_row():
    # Test with a single row
    df = pd.DataFrame({"A": ["foo"], "B": ["bar"], "C": [42]})
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 21.6μs -> 26.3μs (17.9% slower)

def test_missing_index_column_value():
    # Test with missing index/column values (NaN)
    df = pd.DataFrame({
        "A": ["foo", None, "bar", "bar"],
        "B": ["one", "two", None, "two"],
        "C": [1, 2, 3, 4]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 52.9μs -> 28.1μs (88.4% faster)

def test_nonexistent_column_raises():
    # Test for KeyError if index/columns/values do not exist
    df = pd.DataFrame({
        "A": ["foo", "bar"],
        "B": ["one", "two"],
        "C": [1, 2]
    })
    with pytest.raises(KeyError):
        pivot_table(df, index="X", columns="B", values="C") # 20.1μs -> 10.4μs (94.0% faster)
    with pytest.raises(KeyError):
        pivot_table(df, index="A", columns="Y", values="C") # 13.8μs -> 15.0μs (7.52% slower)
    with pytest.raises(KeyError):
        pivot_table(df, index="A", columns="B", values="Z") # 11.8μs -> 11.2μs (4.82% faster)

def test_unsupported_aggfunc():
    # Test for ValueError on unsupported aggregation function
    df = pd.DataFrame({
        "A": ["foo"],
        "B": ["bar"],
        "C": [1]
    })
    with pytest.raises(ValueError):
        pivot_table(df, index="A", columns="B", values="C", aggfunc="median") # 500ns -> 458ns (9.17% faster)

def test_duplicate_groups():
    # Test that multiple rows in the same group are aggregated
    df = pd.DataFrame({
        "A": ["foo", "foo", "foo"],
        "B": ["bar", "bar", "bar"],
        "C": [1, 2, 3]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 43.5μs -> 27.1μs (60.6% faster)

def test_column_with_all_same_value():
    # Test where all values in columns are the same
    df = pd.DataFrame({
        "A": ["foo", "foo", "foo"],
        "B": ["bar", "bar", "bar"],
        "C": [10, 20, 30]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 43.6μs -> 26.7μs (63.3% faster)

def test_column_with_all_unique_values():
    # Test where each row is its own group
    df = pd.DataFrame({
        "A": ["a", "b", "c"],
        "B": ["x", "y", "z"],
        "C": [1, 2, 3]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="count"); result = codeflash_output # 43.0μs -> 26.9μs (59.6% faster)

def test_non_numeric_values_column():
    # Test with non-numeric values for count aggregation
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar"],
        "B": ["x", "x", "y"],
        "C": ["apple", "banana", "pear"]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="count"); result = codeflash_output # 31.3μs -> 26.2μs (19.4% faster)

def test_non_numeric_values_sum_mean():
    # Test with non-numeric values for sum/mean should raise TypeError
    df = pd.DataFrame({
        "A": ["foo", "foo"],
        "B": ["x", "x"],
        "C": ["apple", "banana"]
    })
    with pytest.raises(TypeError):
        pivot_table(df, index="A", columns="B", values="C", aggfunc="sum") # 24.6μs -> 26.1μs (5.74% slower)
    with pytest.raises(TypeError):
        pivot_table(df, index="A", columns="B", values="C", aggfunc="mean") # 14.5μs -> 4.50μs (221% faster)

def test_nan_values_in_values_column():
    # Test with NaN in the values column
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar"],
        "B": ["x", "x", "y"],
        "C": [1.0, float('nan'), 3.0]
    })
    # NaN should propagate in mean/sum
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 43.4μs -> 27.1μs (60.3% faster)

# ---------------------
# 3. LARGE SCALE TEST CASES
# ---------------------


def test_large_number_of_unique_groups():
    # Test with 900 unique groups (30x30)
    n = 30
    df = pd.DataFrame({
        "I": [f"i{i}" for i in range(n) for j in range(n)],
        "J": [f"j{j}" for i in range(n) for j in range(n)],
        "V": [i * n + j for i in range(n) for j in range(n)]
    })
    codeflash_output = pivot_table(df, index="I", columns="J", values="V", aggfunc="count"); result = codeflash_output # 7.87ms -> 216μs (3543% faster)
    # Each group should have exactly one value
    for i in range(n):
        for j in range(n):
            pass

def test_large_scale_performance():
    # Test that the function runs in reasonable time for 1000 rows (not a strict perf test)
    import time
    N = 1000
    df = pd.DataFrame({
        "A": [f"g{i%20}" for i in range(N)],
        "B": [f"h{i%50}" for i in range(N)],
        "C": [i for i in range(N)]
    })
    start = time.time()
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 8.59ms -> 171μs (4898% faster)
    duration = time.time() - start
    # Spot check a group
    group = [row["C"] for _, row in df.iterrows() if row["A"] == "g0" and row["B"] == "h0"]
    if group:
        expected = sum(group) / len(group)

def test_large_sparse_groups():
    # Test with many groups but most are empty (simulate sparse data)
    n = 30
    df = pd.DataFrame({
        "I": [f"i{i}" for i in range(n)],
        "J": [f"j{j}" for j in range(n)],
        "V": [1] * n
    })
    # Only n groups are filled, but all possible i,j pairs exist
    codeflash_output = pivot_table(df, index="I", columns="J", values="V", aggfunc="sum"); result = codeflash_output # 285μs -> 36.2μs (687% faster)
    # Only diagonal should be filled
    for i in range(n):
        for j in range(n):
            if i == j:
                pass
            else:
                pass
# 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

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

# unit tests

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

def test_basic_mean_aggregation():
    # Test mean aggregation with a simple DataFrame
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar", "bar"],
        "B": ["one", "two", "one", "two"],
        "C": [1, 2, 3, 4]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 52.5μs -> 27.6μs (90.0% faster)

def test_basic_sum_aggregation():
    # Test sum aggregation with a simple DataFrame
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar", "bar"],
        "B": ["one", "one", "two", "two"],
        "C": [1, 2, 3, 4]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 52.0μs -> 27.0μs (92.7% faster)

def test_basic_count_aggregation():
    # Test count aggregation with a simple DataFrame
    df = pd.DataFrame({
        "X": ["a", "a", "b", "b", "b"],
        "Y": ["x", "y", "x", "x", "y"],
        "Z": [10, 20, 30, 40, 50]
    })
    codeflash_output = pivot_table(df, index="X", columns="Y", values="Z", aggfunc="count"); result = codeflash_output # 61.0μs -> 26.7μs (129% faster)

def test_basic_default_aggfunc():
    # Test that default aggfunc is mean
    df = pd.DataFrame({
        "I": ["a", "a", "b"],
        "J": ["x", "x", "x"],
        "K": [2, 4, 6]
    })
    codeflash_output = pivot_table(df, index="I", columns="J", values="K"); result = codeflash_output # 42.5μs -> 27.0μs (57.0% faster)

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

def test_empty_dataframe():
    # Test with an empty DataFrame
    df = pd.DataFrame(columns=["A", "B", "C"])
    codeflash_output = pivot_table(df, index="A", columns="B", values="C"); result = codeflash_output # 750ns -> 26.0μs (97.1% slower)

def test_single_row_dataframe():
    # Test with a DataFrame with a single row
    df = pd.DataFrame({"A": ["x"], "B": ["y"], "C": [42]})
    codeflash_output = pivot_table(df, index="A", columns="B", values="C"); result = codeflash_output # 21.8μs -> 26.2μs (17.0% slower)

def test_missing_combinations():
    # Test with missing combinations of index and columns
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar"],
        "B": ["one", "two", "two"],
        "C": [1, 2, 3]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="sum"); result = codeflash_output # 43.4μs -> 26.8μs (62.3% faster)

def test_non_numeric_values():
    # Test with non-numeric values for count aggregation
    df = pd.DataFrame({
        "A": ["x", "x", "y"],
        "B": ["a", "a", "b"],
        "C": ["cat", "dog", "fish"]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="count"); result = codeflash_output # 31.7μs -> 26.4μs (19.9% faster)

def test_nan_values_in_values_column():
    # Test with NaN values in the values column
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar"],
        "B": ["x", "x", "y"],
        "C": [1.0, float('nan'), 3.0]
    })
    # mean aggregation should propagate nan if present
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 43.5μs -> 26.9μs (61.5% faster)
    # 'foo'-'x': mean(1.0, nan) = nan, 'bar'-'y': 3.0
    import math

def test_unsupported_aggfunc_raises():
    # Test that an unsupported aggfunc raises ValueError
    df = pd.DataFrame({
        "A": ["a", "b"],
        "B": ["x", "y"],
        "C": [1, 2]
    })
    with pytest.raises(ValueError):
        pivot_table(df, index="A", columns="B", values="C", aggfunc="median") # 500ns -> 500ns (0.000% faster)

def test_duplicate_index_column_names():
    # Test when index and columns are the same column
    df = pd.DataFrame({
        "A": ["a", "b", "a"],
        "C": [1, 2, 3]
    })
    # Use 'A' for both index and columns
    codeflash_output = pivot_table(df, index="A", columns="A", values="C", aggfunc="sum"); result = codeflash_output # 42.7μs -> 23.0μs (85.7% faster)

def test_non_hashable_index_or_column():
    # Test with non-hashable values in index or column
    df = pd.DataFrame({
        "A": [[1,2], [1,2], [3,4]],
        "B": ["x", "y", "x"],
        "C": [10, 20, 30]
    })
    # Should work as lists are hashable in pandas, but not in dict
    with pytest.raises(TypeError):
        pivot_table(df, index="A", columns="B", values="C", aggfunc="sum") # 20.8μs -> 25.6μs (18.9% slower)

def test_values_column_with_all_nan():
    # Test when all values in the values column are NaN
    df = pd.DataFrame({
        "A": ["a", "b"],
        "B": ["x", "y"],
        "C": [float('nan'), float('nan')]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 33.9μs -> 26.8μs (26.4% faster)
    import math

def test_column_with_all_same_value():
    # Test when the column field has only one value
    df = pd.DataFrame({
        "A": ["foo", "foo", "bar"],
        "B": ["one", "one", "one"],
        "C": [5, 10, 15]
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 42.0μs -> 27.1μs (55.1% faster)

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

def test_large_scale_unique_combinations():
    # Test with 100 unique index and column combinations
    n = 10
    df = pd.DataFrame({
        "I": [f"i{i}" for i in range(n) for _ in range(n)],
        "J": [f"j{j}" for _ in range(n) for j in range(n)],
        "V": [i * j for i in range(n) for j in range(n)]
    })
    codeflash_output = pivot_table(df, index="I", columns="J", values="V", aggfunc="sum"); result = codeflash_output # 866μs -> 52.2μs (1560% faster)
    # Each (i,j) pair appears once, so value is i*j
    for i in range(n):
        for j in range(n):
            pass

def test_large_scale_duplicate_combinations():
    # Test with 1000 rows, 10 index, 10 columns, each combination appears 10 times
    n = 10
    df = pd.DataFrame({
        "idx": [f"i{i}" for i in range(n) for j in range(n) for _ in range(10)],
        "col": [f"c{j}" for i in range(n) for j in range(n) for _ in range(10)],
        "val": [i + j for i in range(n) for j in range(n) for _ in range(10)]
    })
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val", aggfunc="mean"); result = codeflash_output # 8.61ms -> 157μs (5350% faster)
    # Each (i,j) pair appears 10 times, mean is i+j
    for i in range(n):
        for j in range(n):
            pass

def test_large_scale_count():
    # Test with 1000 rows, count aggregation
    n = 100
    df = pd.DataFrame({
        "A": ["foo"] * n + ["bar"] * n,
        "B": ["x"] * (n//2) + ["y"] * (n//2) + ["x"] * (n//2) + ["y"] * (n//2),
        "C": list(range(n)) + list(range(n))
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="count"); result = codeflash_output # 1.70ms -> 39.8μs (4174% faster)

def test_large_scale_sparse_matrix():
    # Test with many missing combinations (sparse)
    n = 30
    df = pd.DataFrame({
        "row": [f"r{i}" for i in range(n) for j in range(0, n, 3)],
        "col": [f"c{j}" for i in range(n) for j in range(0, n, 3)],
        "val": [i * j for i in range(n) for j in range(0, n, 3)]
    })
    codeflash_output = pivot_table(df, index="row", columns="col", values="val", aggfunc="sum"); result = codeflash_output # 2.63ms -> 100.0μs (2531% faster)
    # Only every 3rd column exists for each row
    for i in range(n):
        for j in range(0, n, 3):
            pass

def test_large_scale_all_same_value():
    # Test with a large DataFrame where all values are the same
    n = 500
    df = pd.DataFrame({
        "A": ["a"] * n,
        "B": ["b"] * n,
        "C": [7] * n
    })
    codeflash_output = pivot_table(df, index="A", columns="B", values="C", aggfunc="mean"); result = codeflash_output # 4.25ms -> 65.2μs (6406% 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-pivot_table-mdpen2to and push.