⚡️ Speed up function dataframe_merge by 2,074%

[codeflash-ai]

📄 2,074% (20.74x) speedup for dataframe_merge in src/numpy_pandas/dataframe_operations.py

⏱️ Runtime : 251 milliseconds → 11.5 milliseconds (best of 245 runs)

📝 Explanation and details

The optimized code achieves a 20x speedup by replacing slow pandas .iloc[] operations with much faster tuple-based row iteration.

Key optimizations:

1. Replaced .iloc[] with itertuples(): The original code used right.iloc[i] and left.iloc[i] for row access, which creates new Series objects and involves expensive indexing operations. The optimized version uses itertuples(index=False, name=None) which returns lightweight tuples directly.

2. Pre-computed column index mappings: Instead of repeatedly accessing columns by name on Series objects, the optimized code creates left_col_idx and right_col_idx dictionaries that map column names to tuple positions, enabling direct integer indexing like row[left_col_idx[col]].

3. Stored tuples instead of indices: The right_dict now stores the actual row tuples rather than row indices, eliminating the need for additional .iloc[] lookups during the merge phase.

From the line profiler results, the most expensive operations in the original code were:

• right.iloc[right_idx] (47.3% of total time)
• left.iloc[i] (14.9% of total time)
• right.iloc[i][right_on] (12.4% of total time)

These are completely eliminated in the optimized version.

Performance characteristics: The optimization is most effective for larger datasets (as seen in the test results where large-scale tests show 2000-4000% speedups) but maintains good performance even for small datasets. Edge cases with empty DataFrames show slightly slower performance due to setup overhead, but normal use cases benefit significantly.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

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

# unit tests

# 1. Basic Test Cases

def test_basic_single_match():
    # One row in each, matching key
    left = pd.DataFrame({'id': [1], 'val': ['A']})
    right = pd.DataFrame({'key': [1], 'desc': ['foo']})
    expected = pd.DataFrame({'id': [1], 'val': ['A'], 'desc': ['foo']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 171μs -> 189μs (9.26% slower)

def test_basic_multiple_matches():
    # Multiple rows, unique keys
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [1, 2], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B'], 'desc': ['foo', 'bar']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 212μs -> 190μs (11.5% faster)

def test_basic_duplicate_keys_in_right():
    # Right has duplicate keys, left has unique
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [1, 1, 2], 'desc': ['foo', 'baz', 'bar']})
    expected = pd.DataFrame({
        'id': [1, 1, 2],
        'val': ['A', 'A', 'B'],
        'desc': ['foo', 'baz', 'bar']
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 238μs -> 188μs (26.7% faster)
    # Sort for deterministic comparison
    result_sorted = result.sort_values(['id', 'desc']).reset_index(drop=True)
    expected_sorted = expected.sort_values(['id', 'desc']).reset_index(drop=True)

def test_basic_duplicate_keys_in_left():
    # Left has duplicate keys, right has unique
    left = pd.DataFrame({'id': [1, 1, 2], 'val': ['A', 'C', 'B']})
    right = pd.DataFrame({'key': [1, 2], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({
        'id': [1, 1, 2],
        'val': ['A', 'C', 'B'],
        'desc': ['foo', 'foo', 'bar']
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 239μs -> 188μs (27.0% faster)
    result_sorted = result.sort_values(['id', 'val']).reset_index(drop=True)
    expected_sorted = expected.sort_values(['id', 'val']).reset_index(drop=True)

def test_basic_multiple_matches_both_sides():
    # Both sides have duplicates
    left = pd.DataFrame({'id': [1, 1, 2], 'val': ['A', 'C', 'B']})
    right = pd.DataFrame({'key': [1, 1, 2], 'desc': ['foo', 'baz', 'bar']})
    expected = pd.DataFrame({
        'id': [1, 1, 1, 1, 2],
        'val': ['A', 'A', 'C', 'C', 'B'],
        'desc': ['foo', 'baz', 'foo', 'baz', 'bar']
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 284μs -> 190μs (49.6% faster)
    result_sorted = result.sort_values(['id', 'val', 'desc']).reset_index(drop=True)
    expected_sorted = expected.sort_values(['id', 'val', 'desc']).reset_index(drop=True)

def test_basic_different_column_names():
    # Merge on differently named columns
    left = pd.DataFrame({'left_id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'right_id': [1, 2], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'left_id': [1, 2], 'val': ['A', 'B'], 'desc': ['foo', 'bar']})
    codeflash_output = dataframe_merge(left, right, 'left_id', 'right_id'); result = codeflash_output # 211μs -> 188μs (12.4% faster)

# 2. Edge Test Cases

def test_edge_no_matches():
    # No matching keys
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [3, 4], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'id': [], 'val': [], 'desc': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 218μs -> 221μs (1.60% slower)

def test_edge_empty_left():
    # Left DataFrame is empty
    left = pd.DataFrame({'id': [], 'val': []})
    right = pd.DataFrame({'key': [1, 2], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'id': [], 'val': [], 'desc': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 173μs -> 217μs (20.2% slower)

def test_edge_empty_right():
    # Right DataFrame is empty
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [], 'desc': []})
    expected = pd.DataFrame({'id': [], 'val': [], 'desc': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 170μs -> 214μs (20.5% slower)

def test_edge_both_empty():
    # Both DataFrames are empty
    left = pd.DataFrame({'id': [], 'val': []})
    right = pd.DataFrame({'key': [], 'desc': []})
    expected = pd.DataFrame({'id': [], 'val': [], 'desc': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 117μs -> 216μs (45.8% slower)

def test_edge_missing_merge_column_left():
    # Missing merge column in left
    left = pd.DataFrame({'not_id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [1, 2], 'desc': ['foo', 'bar']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 83.0μs -> 102μs (18.9% slower)

def test_edge_missing_merge_column_right():
    # Missing merge column in right
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'not_key': [1, 2], 'desc': ['foo', 'bar']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 39.3μs -> 13.3μs (196% faster)

def test_edge_null_values_in_merge_column():
    # Null values in merge columns
    left = pd.DataFrame({'id': [1, None, 2], 'val': ['A', 'B', 'C']})
    right = pd.DataFrame({'key': [1, 2, None], 'desc': ['foo', 'bar', 'baz']})
    expected = pd.DataFrame({
        'id': [1, 2, None],
        'val': ['A', 'C', 'B'],
        'desc': ['foo', 'bar', 'baz']
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 245μs -> 191μs (27.8% faster)
    # Only rows with matching keys (including None) should be present
    result_sorted = result.sort_values(['id', 'val']).reset_index(drop=True)
    expected_sorted = expected.sort_values(['id', 'val']).reset_index(drop=True)

def test_edge_merge_column_with_different_types():
    # Merge columns with different types (should not match)
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': ['1', '2'], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'id': [], 'val': [], 'desc': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 204μs -> 218μs (6.41% slower)

def test_edge_merge_column_with_nan_and_none():
    # Merge columns with NaN and None
    left = pd.DataFrame({'id': [float('nan'), None, 2], 'val': ['A', 'B', 'C']})
    right = pd.DataFrame({'key': [float('nan'), None, 2], 'desc': ['foo', 'bar', 'baz']})
    # Only 2 should match, NaN != NaN, None == None
    expected = pd.DataFrame({
        'id': [None, 2],
        'val': ['B', 'C'],
        'desc': ['bar', 'baz']
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 220μs -> 191μs (15.3% faster)
    # Remove NaN rows for comparison
    result_non_nan = result[result['id'].notnull() | result['id'].isnull()]
    expected_non_nan = expected[expected['id'].notnull() | expected['id'].isnull()]
    result_sorted = result_non_nan.sort_values(['id', 'val']).reset_index(drop=True)
    expected_sorted = expected_non_nan.sort_values(['id', 'val']).reset_index(drop=True)

def test_edge_merge_column_is_index():
    # Merge column is index in one or both DataFrames
    left = pd.DataFrame({'val': ['A', 'B']}, index=[1, 2])
    right = pd.DataFrame({'desc': ['foo', 'bar']}, index=[1, 2])
    left = left.reset_index().rename(columns={'index': 'id'})
    right = right.reset_index().rename(columns={'index': 'key'})
    expected = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B'], 'desc': ['foo', 'bar']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 213μs -> 181μs (17.4% faster)

def test_edge_merge_with_extra_columns():
    # Extra columns in right DataFrame
    left = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B']})
    right = pd.DataFrame({'key': [1, 2], 'desc': ['foo', 'bar'], 'extra': [10, 20]})
    expected = pd.DataFrame({'id': [1, 2], 'val': ['A', 'B'], 'desc': ['foo', 'bar'], 'extra': [10, 20]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 249μs -> 237μs (5.25% faster)

def test_edge_merge_with_overlapping_column_names():
    # Overlapping column names except merge column
    left = pd.DataFrame({'id': [1, 2], 'desc': ['A', 'B']})
    right = pd.DataFrame({'id': [1, 2], 'desc': ['foo', 'bar']})
    expected = pd.DataFrame({'id': [1, 2], 'desc': ['A', 'B']})
    # Since right's 'desc' is dropped, only left's 'desc' remains
    codeflash_output = dataframe_merge(left, right, 'id', 'id'); result = codeflash_output # 197μs -> 175μs (12.4% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_rows():
    # Merge two DataFrames with 1000 rows each, all keys match
    N = 1000
    left = pd.DataFrame({'id': list(range(N)), 'val': [str(i) for i in range(N)]})
    right = pd.DataFrame({'key': list(range(N)), 'desc': ['desc'+str(i) for i in range(N)]})
    expected = pd.DataFrame({
        'id': list(range(N)),
        'val': [str(i) for i in range(N)],
        'desc': ['desc'+str(i) for i in range(N)]
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 39.5ms -> 1.12ms (3442% faster)
    # Sort for deterministic comparison
    result_sorted = result.sort_values('id').reset_index(drop=True)
    expected_sorted = expected.sort_values('id').reset_index(drop=True)

def test_large_scale_sparse_matches():
    # Merge two DataFrames with 1000 rows each, only 10 keys match
    N = 1000
    match_keys = set(range(0, N, 100))
    left = pd.DataFrame({'id': list(range(N)), 'val': [str(i) for i in range(N)]})
    right = pd.DataFrame({'key': list(match_keys), 'desc': ['desc'+str(i) for i in match_keys]})
    expected = pd.DataFrame({
        'id': sorted(match_keys),
        'val': [str(i) for i in sorted(match_keys)],
        'desc': ['desc'+str(i) for i in sorted(match_keys)]
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 12.5ms -> 363μs (3328% faster)
    result_sorted = result.sort_values('id').reset_index(drop=True)
    expected_sorted = expected.sort_values('id').reset_index(drop=True)

def test_large_scale_many_duplicates():
    # Many duplicate keys in both DataFrames
    N = 100
    left = pd.DataFrame({'id': [1]*N, 'val': [str(i) for i in range(N)]})
    right = pd.DataFrame({'key': [1]*N, 'desc': ['desc'+str(i) for i in range(N)]})
    # Each left row matches every right row
    expected_rows = []
    for i in range(N):
        for j in range(N):
            expected_rows.append({'id': 1, 'val': str(i), 'desc': 'desc'+str(j)})
    expected = pd.DataFrame(expected_rows)
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 149ms -> 5.20ms (2776% faster)
    # Sort for deterministic comparison
    result_sorted = result.sort_values(['val', 'desc']).reset_index(drop=True)
    expected_sorted = expected.sort_values(['val', 'desc']).reset_index(drop=True)

def test_large_scale_no_matches():
    # No keys match in large DataFrames
    N = 1000
    left = pd.DataFrame({'id': list(range(N)), 'val': [str(i) for i in range(N)]})
    right = pd.DataFrame({'key': [i+N for i in range(N)], 'desc': ['desc'+str(i) for i in range(N)]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 24.5ms -> 569μs (4198% faster)

def test_large_scale_extra_columns():
    # Large DataFrames with extra columns
    N = 500
    left = pd.DataFrame({'id': list(range(N)), 'val': [str(i) for i in range(N)], 'extra1': [i*2 for i in range(N)]})
    right = pd.DataFrame({'key': list(range(N)), 'desc': ['desc'+str(i) for i in range(N)], 'extra2': [i*3 for i in range(N)]})
    expected = pd.DataFrame({
        'id': list(range(N)),
        'val': [str(i) for i in range(N)],
        'extra1': [i*2 for i in range(N)],
        'desc': ['desc'+str(i) for i in range(N)],
        'extra2': [i*3 for i in range(N)]
    })
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 21.4ms -> 975μs (2097% faster)
    result_sorted = result.sort_values('id').reset_index(drop=True)
    expected_sorted = expected.sort_values('id').reset_index(drop=True)
# 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-dataframe_merge-mhb125gu and push.