⚡️ Speed up function pivot_table by 3,237%

[codeflash-ai]

📄 3,237% (32.37x) speedup for pivot_table in src/data_processing/transformations.py

⏱️ Runtime : 206 milliseconds → 6.18 milliseconds (best of 93 runs)

📝 Explanation and details

The optimized code achieves a 32x speedup by eliminating the primary bottleneck: repeated df.iloc[i] calls within the loop. In the original implementation, each df.iloc[i] triggers pandas overhead to extract a single row as a Series, which is extremely expensive when repeated thousands of times (accounting for ~70% of runtime in the line profiler).

Key optimizations:

1. Vectorized data extraction: The optimization pre-extracts entire columns as NumPy arrays using df[column].values before the loop. This converts pandas Series to raw NumPy arrays, which have minimal access overhead.

2. Direct array iteration with zip(): Instead of for i in range(len(df)) followed by df.iloc[i], the code uses zip(index_data, column_data, value_data) to iterate directly over array values. This eliminates per-row pandas indexing overhead entirely.

3. Simplified dictionary operations with setdefault(): The nested dictionary initialization is streamlined using setdefault(), which combines the existence check and default assignment into a single operation, reducing redundant dictionary lookups.

Performance characteristics:

• Small DataFrames (1-5 rows): The optimization shows marginal improvement or slight regression (~20-50μs vs ~40-100μs) because the upfront cost of extracting NumPy arrays dominates when there are few rows to process.

• Large DataFrames (1000+ rows): The optimization excels dramatically, showing 50-80x speedups (e.g., 14.5ms → 200μs). The fixed overhead of array extraction (~38ms total across three columns based on line profiler) is amortized over many rows, while eliminating the quadratic-like cost of repeated .iloc[] calls.

• All aggregation functions (mean, sum, count) benefit equally since the bottleneck was in the grouping phase, not the aggregation phase.

Impact considerations:

The function processes DataFrames to create pivot table-like aggregations. If this function is called in data processing pipelines or repeated analytics workflows with moderately-sized DataFrames (hundreds to thousands of rows), the optimization will significantly reduce processing time. The speedup scales linearly with DataFrame size, making it particularly valuable for batch processing or real-time analytics on non-trivial datasets.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from typing import Any

# function to test
# src/data_processing/transformations.py
import pandas as pd  # used to create DataFrames for our tests

# imports
import pytest  # used for our unit tests
from src.data_processing.transformations import pivot_table

# unit tests

# 1. Basic Test Cases


def test_basic_mean_aggregation():
    # Test mean aggregation on 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  # 97.5μs -> 51.7μs (88.6% faster)


def test_basic_sum_aggregation():
    # Test sum aggregation
    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  # 94.9μs -> 49.3μs (92.5% faster)


def test_basic_count_aggregation():
    # Test count aggregation
    df = pd.DataFrame(
        {
            "A": ["foo", "foo", "bar", "bar", "foo"],
            "B": ["one", "one", "two", "two", "two"],
            "C": [1, 2, 3, 4, 5],
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="count"
    )
    result = codeflash_output  # 110μs -> 48.5μs (127% faster)


def test_basic_multiple_values_per_cell():
    # Test aggregation when multiple values per cell
    df = pd.DataFrame(
        {
            "A": ["foo", "foo", "foo", "bar", "bar"],
            "B": ["one", "one", "two", "two", "two"],
            "C": [1, 2, 3, 4, 5],
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="mean"
    )
    result = codeflash_output  # 109μs -> 49.1μs (122% faster)


# 2. Edge Test Cases


def test_empty_dataframe():
    # Test with 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  # 1.79μs -> 49.0μs (96.3% slower)


def test_single_row_dataframe():
    # 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  # 39.8μs -> 47.9μs (16.8% slower)


def test_missing_column():
    # Test with missing column name
    df = pd.DataFrame({"A": ["foo"], "B": ["bar"], "C": [42]})
    with pytest.raises(KeyError):
        pivot_table(
            df, index="X", columns="B", values="C", aggfunc="mean"
        )  # 41.0μs -> 25.1μs (63.3% faster)


def test_non_numeric_values_for_mean():
    # Test with non-numeric values for mean aggregation
    df = pd.DataFrame({"A": ["foo", "foo"], "B": ["bar", "baz"], "C": ["a", "b"]})
    with pytest.raises(TypeError):
        pivot_table(
            df, index="A", columns="B", values="C", aggfunc="mean"
        )  # 52.0μs -> 49.8μs (4.43% faster)


def test_unsupported_aggfunc():
    # Test with unsupported aggregation function
    df = pd.DataFrame({"A": ["foo"], "B": ["bar"], "C": [42]})
    with pytest.raises(ValueError):
        pivot_table(
            df, index="A", columns="B", values="C", aggfunc="max"
        )  # 1.33μs -> 1.17μs (14.2% faster)


def test_nan_values():
    # Test with NaN values in the values column
    df = pd.DataFrame(
        {"A": ["foo", "foo"], "B": ["bar", "baz"], "C": [1, float("nan")]}
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="mean"
    )
    result = codeflash_output  # 69.0μs -> 53.9μs (28.2% faster)


def test_column_with_all_nan():
    # Test with all NaN values in a group
    df = pd.DataFrame(
        {"A": ["foo", "foo"], "B": ["bar", "bar"], "C": [float("nan"), float("nan")]}
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="mean"
    )
    result = codeflash_output  # 66.2μs -> 52.3μs (26.6% faster)


# 3. Large Scale Test Cases


def test_large_scale_sum():
    # Test with a large DataFrame for sum aggregation
    N = 1000
    df = pd.DataFrame(
        {
            "A": ["foo"] * (N // 2) + ["bar"] * (N // 2),
            "B": ["one"] * (N // 4)
            + ["two"] * (N // 4)
            + ["one"] * (N // 4)
            + ["two"] * (N // 4),
            "C": list(range(N)),
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="sum"
    )
    result = codeflash_output  # 14.6ms -> 212μs (6783% faster)
    # Check that the sums are correct
    foo_one_sum = sum(range(0, N // 4))
    foo_two_sum = sum(range(N // 4, N // 2))
    bar_one_sum = sum(range(N // 2, N * 3 // 4))
    bar_two_sum = sum(range(N * 3 // 4, N))


def test_large_scale_count():
    # Test with a large DataFrame for count aggregation
    N = 1000
    df = pd.DataFrame(
        {
            "A": ["foo"] * (N // 2) + ["bar"] * (N // 2),
            "B": ["one"] * (N // 4)
            + ["two"] * (N // 4)
            + ["one"] * (N // 4)
            + ["two"] * (N // 4),
            "C": list(range(N)),
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="count"
    )
    result = codeflash_output  # 14.5ms -> 180μs (7931% faster)


def test_large_scale_mean():
    # Test with a large DataFrame for mean aggregation
    N = 1000
    df = pd.DataFrame(
        {
            "A": ["foo"] * (N // 2) + ["bar"] * (N // 2),
            "B": ["one"] * (N // 4)
            + ["two"] * (N // 4)
            + ["one"] * (N // 4)
            + ["two"] * (N // 4),
            "C": list(range(N)),
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="mean"
    )
    result = codeflash_output  # 14.5ms -> 207μs (6885% faster)
    # Check that the means are correct
    foo_one_vals = list(range(0, N // 4))
    foo_two_vals = list(range(N // 4, N // 2))
    bar_one_vals = list(range(N // 2, N * 3 // 4))
    bar_two_vals = list(range(N * 3 // 4, N))


def test_large_scale_distinct_index_and_column():
    # Test with many distinct index and column values
    N = 1000
    df = pd.DataFrame(
        {
            "A": [f"idx_{i}" for i in range(N)],
            "B": [f"col_{i}" for i in range(N)],
            "C": list(range(N)),
        }
    )
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="sum"
    )
    result = codeflash_output  # 15.1ms -> 579μs (2500% faster)
    # Each cell should contain its own value
    for i in range(N):
        pass


def test_large_scale_all_same_group():
    # Test with all rows in the same group
    N = 1000
    df = pd.DataFrame({"A": ["foo"] * N, "B": ["bar"] * N, "C": list(range(N))})
    codeflash_output = pivot_table(
        df, index="A", columns="B", values="C", aggfunc="sum"
    )
    result = codeflash_output  # 14.5ms -> 207μs (6878% faster)


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

import math
from typing import Any

import pandas as pd

# imports
import pytest  # used for our unit tests
from src.data_processing.transformations import pivot_table

# unit tests

# ============================================================================
# BASIC TEST CASES - Fundamental functionality under normal conditions
# ============================================================================


def test_basic_mean_aggregation():
    """Test basic pivot table with mean aggregation (default)."""
    # Create a simple dataframe with sales data
    df = pd.DataFrame(
        {
            "region": ["North", "North", "South", "South"],
            "product": ["A", "B", "A", "B"],
            "sales": [100, 200, 150, 250],
        }
    )

    # Pivot with region as index, product as columns, sales as values
    codeflash_output = pivot_table(
        df, index="region", columns="product", values="sales"
    )
    result = codeflash_output  # 96.7μs -> 50.0μs (93.6% faster)


def test_basic_sum_aggregation():
    """Test pivot table with sum aggregation."""
    # Create dataframe with multiple entries per group
    df = pd.DataFrame(
        {
            "category": ["X", "X", "Y", "Y"],
            "type": ["T1", "T1", "T2", "T2"],
            "amount": [10, 20, 30, 40],
        }
    )

    # Pivot with sum aggregation
    codeflash_output = pivot_table(
        df, index="category", columns="type", values="amount", aggfunc="sum"
    )
    result = codeflash_output  # 94.1μs -> 49.3μs (90.9% faster)


def test_basic_count_aggregation():
    """Test pivot table with count aggregation."""
    # Create dataframe with varying counts per group
    df = pd.DataFrame(
        {
            "group": ["A", "A", "A", "B", "B"],
            "subgroup": ["S1", "S1", "S2", "S1", "S2"],
            "value": [1, 2, 3, 4, 5],
        }
    )

    # Pivot with count aggregation
    codeflash_output = pivot_table(
        df, index="group", columns="subgroup", values="value", aggfunc="count"
    )
    result = codeflash_output  # 110μs -> 49.6μs (122% faster)


def test_single_row_dataframe():
    """Test pivot table with a single row."""
    # Create minimal dataframe with one row
    df = pd.DataFrame({"idx": ["I1"], "col": ["C1"], "val": [42]})

    # Pivot the single row
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 39.6μs -> 47.8μs (17.2% slower)


def test_multiple_values_same_group():
    """Test that multiple values in the same group are aggregated correctly."""
    # Create dataframe with multiple values for same index-column combination
    df = pd.DataFrame(
        {
            "index_col": ["A", "A", "A"],
            "column_col": ["X", "X", "X"],
            "value_col": [10, 20, 30],
        }
    )

    # Pivot with mean (should average all three values)
    codeflash_output = pivot_table(
        df, index="index_col", columns="column_col", values="value_col"
    )
    result = codeflash_output  # 77.6μs -> 48.5μs (59.9% faster)


# ============================================================================
# EDGE TEST CASES - Extreme or unusual conditions
# ============================================================================


def test_empty_dataframe():
    """Test pivot table with an empty dataframe."""
    # Create empty dataframe with correct columns
    df = pd.DataFrame({"idx": [], "col": [], "val": []})

    # Pivot should return empty dict
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 1.75μs -> 48.7μs (96.4% slower)


def test_unsupported_aggregation_function():
    """Test that unsupported aggregation functions raise ValueError."""
    # Create simple dataframe
    df = pd.DataFrame({"idx": ["A"], "col": ["B"], "val": [10]})

    # Attempt to use unsupported aggregation function
    with pytest.raises(ValueError) as exc_info:
        pivot_table(
            df, index="idx", columns="col", values="val", aggfunc="median"
        )  # 1.21μs -> 1.08μs (11.5% faster)


def test_negative_values():
    """Test pivot table with negative numeric values."""
    # Create dataframe with negative numbers
    df = pd.DataFrame(
        {"idx": ["A", "A", "B"], "col": ["X", "Y", "X"], "val": [-10, -20, -30]}
    )

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 81.4μs -> 50.6μs (60.7% faster)


def test_zero_values():
    """Test pivot table with zero values."""
    # Create dataframe with zeros
    df = pd.DataFrame(
        {"idx": ["A", "A", "B"], "col": ["X", "X", "Y"], "val": [0, 0, 0]}
    )

    # Pivot with sum
    codeflash_output = pivot_table(
        df, index="idx", columns="col", values="val", aggfunc="sum"
    )
    result = codeflash_output  # 78.0μs -> 48.8μs (59.8% faster)


def test_floating_point_values():
    """Test pivot table with floating point values."""
    # Create dataframe with float values
    df = pd.DataFrame({"idx": ["A", "A"], "col": ["X", "X"], "val": [1.5, 2.5]})

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 62.5μs -> 48.5μs (29.1% faster)


def test_mixed_positive_negative_values():
    """Test pivot table with mixed positive and negative values."""
    # Create dataframe with mixed signs
    df = pd.DataFrame(
        {"idx": ["A", "A", "A"], "col": ["X", "X", "X"], "val": [-10, 10, 5]}
    )

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 77.0μs -> 48.5μs (58.8% faster)


def test_string_index_and_columns():
    """Test pivot table with string values for index and columns."""
    # Create dataframe with string identifiers
    df = pd.DataFrame(
        {
            "category": ["Electronics", "Electronics", "Furniture"],
            "store": ["Store_A", "Store_B", "Store_A"],
            "revenue": [1000, 1500, 800],
        }
    )

    # Pivot with strings
    codeflash_output = pivot_table(
        df, index="category", columns="store", values="revenue"
    )
    result = codeflash_output  # 77.6μs -> 49.0μs (58.3% faster)


def test_numeric_index_and_columns():
    """Test pivot table with numeric values for index and columns."""
    # Create dataframe with numeric identifiers
    df = pd.DataFrame(
        {"year": [2020, 2020, 2021], "quarter": [1, 2, 1], "sales": [100, 200, 150]}
    )

    # Pivot with numeric keys
    codeflash_output = pivot_table(df, index="year", columns="quarter", values="sales")
    result = codeflash_output  # 65.7μs -> 51.4μs (27.9% faster)


def test_sparse_data():
    """Test pivot table with sparse data (not all combinations present)."""
    # Create dataframe where not all index-column combinations exist
    df = pd.DataFrame(
        {"idx": ["A", "A", "B"], "col": ["X", "Y", "X"], "val": [10, 20, 30]}
    )

    # Pivot the sparse data
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 87.9μs -> 52.7μs (66.9% faster)


def test_single_unique_index():
    """Test pivot table where all rows have the same index value."""
    # Create dataframe with single index value
    df = pd.DataFrame(
        {"idx": ["A", "A", "A"], "col": ["X", "Y", "Z"], "val": [10, 20, 30]}
    )

    # Pivot with single index
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 79.0μs -> 49.2μs (60.5% faster)


def test_single_unique_column():
    """Test pivot table where all rows have the same column value."""
    # Create dataframe with single column value
    df = pd.DataFrame(
        {"idx": ["A", "B", "C"], "col": ["X", "X", "X"], "val": [10, 20, 30]}
    )

    # Pivot with single column
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 78.5μs -> 49.2μs (59.6% faster)


def test_very_large_values():
    """Test pivot table with very large numeric values."""
    # Create dataframe with large numbers
    df = pd.DataFrame({"idx": ["A", "A"], "col": ["X", "X"], "val": [1e15, 2e15]})

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 61.1μs -> 48.7μs (25.4% faster)


def test_very_small_values():
    """Test pivot table with very small numeric values."""
    # Create dataframe with small numbers
    df = pd.DataFrame({"idx": ["A", "A"], "col": ["X", "X"], "val": [1e-15, 2e-15]})

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 60.4μs -> 48.3μs (25.0% faster)


def test_duplicate_rows():
    """Test pivot table with completely duplicate rows."""
    # Create dataframe with duplicate rows
    df = pd.DataFrame(
        {"idx": ["A", "A", "A"], "col": ["X", "X", "X"], "val": [10, 10, 10]}
    )

    # Pivot with mean
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 76.7μs -> 48.2μs (59.3% faster)


def test_boolean_index_columns():
    """Test pivot table with boolean values as index/columns."""
    # Create dataframe with boolean identifiers
    df = pd.DataFrame(
        {
            "flag1": [True, True, False],
            "flag2": [True, False, True],
            "val": [10, 20, 30],
        }
    )

    # Pivot with boolean keys
    codeflash_output = pivot_table(df, index="flag1", columns="flag2", values="val")
    result = codeflash_output  # 80.2μs -> 51.2μs (56.8% faster)


def test_special_characters_in_strings():
    """Test pivot table with special characters in string values."""
    # Create dataframe with special characters
    df = pd.DataFrame(
        {
            "idx": ["A@B", "A@B", "C#D"],
            "col": ["X$Y", "Z%W", "X$Y"],
            "val": [10, 20, 30],
        }
    )

    # Pivot with special characters
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 79.5μs -> 49.5μs (60.5% faster)


def test_whitespace_in_strings():
    """Test pivot table with whitespace in string values."""
    # Create dataframe with whitespace
    df = pd.DataFrame(
        {"idx": ["A B", " C", "D "], "col": ["X Y", " Z", "W "], "val": [10, 20, 30]}
    )

    # Pivot with whitespace
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 78.2μs -> 48.8μs (60.2% faster)


def test_unicode_characters():
    """Test pivot table with unicode characters."""
    # Create dataframe with unicode
    df = pd.DataFrame(
        {
            "idx": ["北京", "上海", "北京"],
            "col": ["产品A", "产品B", "产品A"],
            "val": [100, 200, 150],
        }
    )

    # Pivot with unicode
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 78.2μs -> 49.2μs (58.9% faster)


def test_mixed_type_values():
    """Test pivot table with mixed numeric types (int and float)."""
    # Create dataframe with mixed int and float
    df = pd.DataFrame({"idx": ["A", "A"], "col": ["X", "X"], "val": [10, 20.5]})

    # Pivot with mixed types
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 60.7μs -> 48.6μs (25.0% faster)


# ============================================================================
# LARGE SCALE TEST CASES - Performance and scalability
# ============================================================================


def test_large_number_of_rows():
    """Test pivot table with a large number of rows."""
    # Create dataframe with 1000 rows
    num_rows = 1000
    df = pd.DataFrame(
        {"idx": ["A"] * num_rows, "col": ["X"] * num_rows, "val": list(range(num_rows))}
    )

    # Pivot large dataset
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 14.7ms -> 209μs (6891% faster)

    # Verify correct aggregation (mean of 0 to 999)
    expected_mean = sum(range(num_rows)) / num_rows


def test_large_number_of_unique_indices():
    """Test pivot table with many unique index values."""
    # Create dataframe with 500 unique indices
    num_indices = 500
    df = pd.DataFrame(
        {
            "idx": [f"idx_{i}" for i in range(num_indices)],
            "col": ["X"] * num_indices,
            "val": list(range(num_indices)),
        }
    )

    # Pivot with many indices
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 7.63ms -> 336μs (2166% faster)


def test_large_number_of_unique_columns():
    """Test pivot table with many unique column values."""
    # Create dataframe with 500 unique columns
    num_columns = 500
    df = pd.DataFrame(
        {
            "idx": ["A"] * num_columns,
            "col": [f"col_{i}" for i in range(num_columns)],
            "val": list(range(num_columns)),
        }
    )

    # Pivot with many columns
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 7.50ms -> 257μs (2815% faster)


def test_large_sparse_matrix():
    """Test pivot table creating a large sparse result matrix."""
    # Create dataframe that results in sparse matrix (100x100 with only 200 entries)
    num_entries = 200
    df = pd.DataFrame(
        {
            "idx": [f"idx_{i % 100}" for i in range(num_entries)],
            "col": [f"col_{i // 2}" for i in range(num_entries)],
            "val": list(range(num_entries)),
        }
    )

    # Pivot to create sparse matrix
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 3.05ms -> 151μs (1911% faster)
    # Count total entries in result
    total_entries = sum(len(cols) for cols in result.values())


def test_large_aggregation_groups():
    """Test pivot table with large groups requiring aggregation."""
    # Create dataframe where each group has 100 values to aggregate
    group_size = 100
    num_groups = 10
    data = []
    for i in range(num_groups):
        for j in range(group_size):
            data.append({"idx": f"group_{i}", "col": "X", "val": j})
    df = pd.DataFrame(data)

    # Pivot with large groups
    codeflash_output = pivot_table(
        df, index="idx", columns="col", values="val", aggfunc="sum"
    )
    result = codeflash_output  # 14.7ms -> 238μs (6055% faster)

    # Verify sum for each group (sum of 0 to 99 is 4950)
    expected_sum = sum(range(group_size))
    for i in range(num_groups):
        pass


def test_many_unique_combinations():
    """Test pivot table with many unique index-column combinations."""
    # Create dataframe with 900 unique combinations (30x30)
    size = 30
    data = []
    for i in range(size):
        for j in range(size):
            data.append({"idx": f"idx_{i}", "col": f"col_{j}", "val": i * size + j})
    df = pd.DataFrame(data)

    # Pivot with many combinations
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 13.5ms -> 453μs (2875% faster)


def test_large_sum_aggregation():
    """Test pivot table with sum aggregation on large values."""
    # Create dataframe with 500 rows to sum
    num_rows = 500
    df = pd.DataFrame(
        {"idx": ["A"] * num_rows, "col": ["X"] * num_rows, "val": [1000] * num_rows}
    )

    # Pivot with sum
    codeflash_output = pivot_table(
        df, index="idx", columns="col", values="val", aggfunc="sum"
    )
    result = codeflash_output  # 7.35ms -> 129μs (5576% faster)


def test_large_count_aggregation():
    """Test pivot table with count aggregation on large dataset."""
    # Create dataframe with 800 rows to count
    num_rows = 800
    df = pd.DataFrame(
        {"idx": ["A"] * num_rows, "col": ["X"] * num_rows, "val": list(range(num_rows))}
    )

    # Pivot with count
    codeflash_output = pivot_table(
        df, index="idx", columns="col", values="val", aggfunc="count"
    )
    result = codeflash_output  # 11.8ms -> 152μs (7611% faster)


def test_multiple_large_groups():
    """Test pivot table with multiple large groups."""
    # Create dataframe with 5 groups, each with 150 entries
    num_groups = 5
    group_size = 150
    data = []
    for i in range(num_groups):
        for j in range(group_size):
            data.append({"idx": f"group_{i}", "col": f"col_{i}", "val": j})
    df = pd.DataFrame(data)

    # Pivot with multiple large groups
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 11.1ms -> 202μs (5365% faster)

    # Verify each group's mean (mean of 0 to 149)
    expected_mean = sum(range(group_size)) / group_size
    for i in range(num_groups):
        pass


def test_wide_dataframe():
    """Test pivot table resulting in a wide output (many columns)."""
    # Create dataframe that results in 200 columns for single index
    num_cols = 200
    df = pd.DataFrame(
        {
            "idx": ["A"] * num_cols,
            "col": [f"col_{i}" for i in range(num_cols)],
            "val": list(range(num_cols)),
        }
    )

    # Pivot to create wide result
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 3.03ms -> 132μs (2187% faster)


def test_tall_dataframe():
    """Test pivot table resulting in a tall output (many rows)."""
    # Create dataframe that results in 200 rows for single column
    num_rows = 200
    df = pd.DataFrame(
        {
            "idx": [f"idx_{i}" for i in range(num_rows)],
            "col": ["X"] * num_rows,
            "val": list(range(num_rows)),
        }
    )

    # Pivot to create tall result
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 3.08ms -> 159μs (1826% faster)


def test_balanced_large_pivot():
    """Test pivot table with balanced dimensions (square-ish result)."""
    # Create dataframe resulting in roughly 25x25 result
    size = 25
    data = []
    for i in range(size):
        for j in range(size):
            # Add 2 entries per combination for aggregation
            data.append({"idx": f"idx_{i}", "col": f"col_{j}", "val": i + j})
            data.append({"idx": f"idx_{i}", "col": f"col_{j}", "val": i + j + 1})
    df = pd.DataFrame(data)

    # Pivot balanced dataset
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 18.6ms -> 471μs (3842% faster)


def test_large_dataset_with_repeated_patterns():
    """Test pivot table with large dataset containing repeated patterns."""
    # Create dataframe with repeating pattern (10 indices, 10 columns, 10 repetitions)
    num_indices = 10
    num_cols = 10
    num_reps = 10
    data = []
    for rep in range(num_reps):
        for i in range(num_indices):
            for j in range(num_cols):
                data.append(
                    {
                        "idx": f"idx_{i}",
                        "col": f"col_{j}",
                        "val": rep * 100 + i * 10 + j,
                    }
                )
    df = pd.DataFrame(data)

    # Pivot with repeated patterns
    codeflash_output = pivot_table(df, index="idx", columns="col", values="val")
    result = codeflash_output  # 14.7ms -> 280μs (5143% faster)
    # Each cell should have mean of 10 values
    # For idx_0, col_0: values are 0, 100, 200, ..., 900
    expected_mean = sum(rep * 100 for rep in range(num_reps)) / num_reps


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