⚡️ Speed up function dataframe_merge by 1,305%

[codeflash-ai]

📄 1,305% (13.05x) speedup for dataframe_merge in src/numpy_pandas/dataframe_operations.py

⏱️ Runtime : 119 milliseconds → 8.48 milliseconds (best of 267 runs)

📝 Explanation and details

The optimized code achieves a 13x speedup by eliminating pandas' slowest operations and leveraging NumPy arrays for data access.

Key optimizations:

1. Eliminated .iloc[] calls: The original code used left.iloc[i] and right.iloc[right_idx] for every row access, which are extremely expensive operations. The optimized version extracts the underlying NumPy arrays once using .values and accesses rows directly via array indexing.

2. Pre-cached column indices: Instead of repeatedly looking up column names during the merge loop, the optimized code pre-computes column indices using get_loc() and stores them in dictionaries for O(1) lookup.

3. Vectorized right-side dictionary building: Uses enumerate(right_values[:, right_on_idx]) to build the key mapping in one pass, avoiding individual .iloc[] calls for each right DataFrame row.

Performance impact by test case:

• Large-scale tests show the most dramatic improvements (1300-3800% faster) because the .iloc[] overhead scales poorly with DataFrame size
• Basic operations with small DataFrames still see 40-100% speedups
• Edge cases with empty DataFrames show minimal improvement since there's less data to process

The line profiler confirms this: the original code spent 33.4% of time in right.iloc[i][right_on] and 27.6% in left.iloc[i] calls, while the optimized version eliminates these bottlenecks entirely. The optimization is particularly effective for datasets with hundreds or thousands of rows where pandas overhead becomes the dominant cost.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import random
import string

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

# unit tests

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

def test_basic_inner_join_one_match():
    # Test simple join with one matching row
    left = pd.DataFrame({'id': [1], 'val_left': ['a']})
    right = pd.DataFrame({'key': [1], 'val_right': ['b']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 93.8μs -> 68.7μs (36.6% faster)

def test_basic_inner_join_multiple_matches():
    # Test join with multiple matching rows
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [2, 1], 'val_right': ['c', 'd']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 117μs -> 67.5μs (73.9% faster)
    # Check that all expected combinations exist
    expected = [
        {'id': 1, 'val_left': 'a', 'val_right': 'd'},
        {'id': 2, 'val_left': 'b', 'val_right': 'c'}
    ]
    for row in expected:
        pass

def test_basic_no_match():
    # Test join where there are no matching keys
    left = pd.DataFrame({'id': [1], 'val_left': ['a']})
    right = pd.DataFrame({'key': [2], 'val_right': ['b']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 96.6μs -> 85.5μs (13.0% faster)

def test_basic_duplicate_keys():
    # Test join where right has duplicate keys (should produce cartesian product)
    left = pd.DataFrame({'id': [1], 'val_left': ['a']})
    right = pd.DataFrame({'key': [1, 1], 'val_right': ['b', 'c']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 109μs -> 67.8μs (61.1% faster)
    vals = set(result['val_right'])

def test_basic_left_duplicate_keys():
    # Test join where left has duplicate keys
    left = pd.DataFrame({'id': [1, 1], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1], 'val_right': ['c']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 108μs -> 67.1μs (61.7% faster)
    vals = set(result['val_left'])

def test_basic_both_duplicate_keys():
    # Test join where both left and right have duplicate keys
    left = pd.DataFrame({'id': [1, 1], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 1], 'val_right': ['c', 'd']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 134μs -> 68.0μs (98.0% faster)
    # All combinations should exist
    expected = set((l, r) for l in ['a', 'b'] for r in ['c', 'd'])
    actual = set(zip(result['val_left'], result['val_right']))

def test_basic_column_overlap():
    # Test join where left and right have overlapping column names (other than join key)
    left = pd.DataFrame({'id': [1], 'val': ['a']})
    right = pd.DataFrame({'key': [1], 'val': ['b']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 84.0μs -> 60.0μs (40.1% faster)

def test_basic_non_integer_keys():
    # Test join with string keys
    left = pd.DataFrame({'id': ['x', 'y'], 'val_left': [1, 2]})
    right = pd.DataFrame({'key': ['y', 'z'], 'val_right': [3, 4]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 105μs -> 64.5μs (62.9% faster)

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

def test_edge_empty_left():
    # Test join with empty left DataFrame
    left = pd.DataFrame({'id': [], 'val_left': []})
    right = pd.DataFrame({'key': [1], 'val_right': ['a']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 82.3μs -> 80.8μs (1.96% faster)

def test_edge_empty_right():
    # Test join with empty right DataFrame
    left = pd.DataFrame({'id': [1], 'val_left': ['a']})
    right = pd.DataFrame({'key': [], 'val_right': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 80.8μs -> 79.2μs (1.95% faster)

def test_edge_both_empty():
    # Test join with both DataFrames empty
    left = pd.DataFrame({'id': [], 'val_left': []})
    right = pd.DataFrame({'key': [], 'val_right': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 62.6μs -> 72.0μs (13.1% slower)

def test_edge_missing_left_on_column():
    # Test error when left_on column is missing
    left = pd.DataFrame({'idx': [1], 'val_left': ['a']})
    right = pd.DataFrame({'key': [1], 'val_right': ['b']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 35.8μs -> 6.88μs (421% faster)

def test_edge_missing_right_on_column():
    # Test error when right_on column is missing
    left = pd.DataFrame({'id': [1], 'val_left': ['a']})
    right = pd.DataFrame({'k': [1], 'val_right': ['b']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 20.8μs -> 8.88μs (134% faster)

def test_edge_null_keys():
    # Test join with NaN/null keys
    left = pd.DataFrame({'id': [1, None], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, None], 'val_right': ['c', 'd']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 110μs -> 68.9μs (61.0% faster)

def test_edge_mixed_types_keys():
    # Test join with mixed types in join keys
    left = pd.DataFrame({'id': [1, '1'], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, '1'], 'val_right': ['c', 'd']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 94.9μs -> 55.4μs (71.3% faster)
    pairs = set(zip(result['val_left'], result['val_right']))

def test_edge_non_unique_columns():
    # Test join with duplicate column names in right (should not happen in pandas, but test for robustness)
    left = pd.DataFrame({'id': [1], 'val': ['a']})
    right = pd.DataFrame({'key': [1], 'val': ['b'], 'val2': ['c']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 89.8μs -> 67.1μs (33.8% faster)

def test_edge_column_order_preserved():
    # Test that column order is left columns, then right (excluding join key)
    left = pd.DataFrame({'id': [1], 'a': [2]})
    right = pd.DataFrame({'key': [1], 'b': [3], 'c': [4]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 72.8μs -> 49.6μs (46.6% faster)

def test_edge_right_on_is_also_left_col():
    # Test if right_on column is also present in left, it should not be duplicated
    left = pd.DataFrame({'id': [1], 'key': [5]})
    right = pd.DataFrame({'key': [1], 'val': [10]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 67.6μs -> 45.6μs (48.1% faster)

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

def test_large_scale_all_match():
    # Test join with 1000 rows, all keys match
    n = 1000
    left = pd.DataFrame({'id': list(range(n)), 'left_val': list(range(n, 2*n))})
    right = pd.DataFrame({'key': list(range(n)), 'right_val': list(range(2*n, 3*n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 16.5ms -> 1.14ms (1339% faster)
    # Check a few random rows
    for i in [0, n//2, n-1]:
        row = result[result['id'] == i].iloc[0]

def test_large_scale_no_match():
    # Test join with 1000 rows, no keys match
    n = 1000
    left = pd.DataFrame({'id': list(range(n)), 'left_val': list(range(n))})
    right = pd.DataFrame({'key': list(range(n, 2*n)), 'right_val': list(range(n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 9.65ms -> 244μs (3852% faster)

def test_large_scale_some_matches():
    # Test join with 1000 rows, half keys match
    n = 1000
    left = pd.DataFrame({'id': list(range(n)), 'left_val': list(range(n))})
    right = pd.DataFrame({'key': list(range(n//2, n + n//2)), 'right_val': list(range(n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 13.0ms -> 671μs (1835% faster)
    # Check a few matches
    for i in [n//2, n//2+10, n-1]:
        if i < n:
            idx = i - n//2
            if 0 <= idx < n:
                row = result[result['id'] == i]
                if not row.empty:
                    pass

def test_large_scale_duplicates():
    # Test join where right has duplicate keys, left has unique
    n = 200
    left = pd.DataFrame({'id': list(range(n)), 'left_val': list(range(n))})
    # Each key in right appears twice
    right = pd.DataFrame({'key': [i//2 for i in range(2*n)], 'right_val': list(range(2*n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 5.58ms -> 426μs (1210% faster)
    for i in range(n):
        matches = result[result['id'] == i]
        vals = set(matches['right_val'])
        expected = {2*i, 2*i+1}

def test_large_scale_string_keys():
    # Test join with large number of string keys
    n = 500
    keys = [''.join(random.choices(string.ascii_letters, k=5)) for _ in range(n)]
    left = pd.DataFrame({'id': keys, 'left_val': list(range(n))})
    right = pd.DataFrame({'key': keys, 'right_val': list(range(n, 2*n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 11.9ms -> 410μs (2808% faster)
    # Check a few random rows
    for i in [0, n//2, n-1]:
        row = result[result['id'] == keys[i]].iloc[0]

def test_large_scale_performance():
    # Test join with 1000 rows, random keys, some matches, ensure completes in reasonable time
    n = 1000
    keys_left = [random.randint(0, 1499) for _ in range(n)]
    keys_right = [random.randint(0, 1499) for _ in range(n)]
    left = pd.DataFrame({'id': keys_left, 'left_val': list(range(n))})
    right = pd.DataFrame({'key': keys_right, 'right_val': list(range(n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 14.0ms -> 832μs (1582% faster)
    # All rows in result must have id == key from left/right
    for idx, row in result.iterrows():
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import random
import string
import sys

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

# unit tests

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

def test_basic_inner_join_single_match():
    # Simple inner join with one matching row
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [2, 3], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 111μs -> 69.8μs (59.8% faster)

def test_basic_inner_join_multiple_matches():
    # Multiple matches, including duplicate keys in right
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [2, 2], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 117μs -> 68.3μs (72.6% faster)
    vals = set(result['val_right'])

def test_basic_no_matches():
    # No keys match
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [3, 4], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 114μs -> 86.9μs (32.1% faster)

def test_basic_all_match():
    # All keys match
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 2], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 118μs -> 68.0μs (74.0% faster)

def test_basic_column_order_and_names():
    # Ensure columns are ordered: left columns first, then right (except join key)
    left = pd.DataFrame({'foo': [1], 'bar': [2]})
    right = pd.DataFrame({'baz': [1], 'qux': [3]})
    codeflash_output = dataframe_merge(left, right, 'foo', 'baz'); result = codeflash_output # 70.1μs -> 46.6μs (50.4% faster)

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

def test_edge_empty_left():
    # Left DataFrame is empty
    left = pd.DataFrame({'id': [], 'val_left': []})
    right = pd.DataFrame({'key': [1, 2], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 92.4μs -> 80.7μs (14.5% faster)

def test_edge_empty_right():
    # Right DataFrame is empty
    left = pd.DataFrame({'id': [1, 2], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [], 'val_right': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 90.8μs -> 79.9μs (13.7% faster)

def test_edge_both_empty():
    # Both DataFrames are empty
    left = pd.DataFrame({'id': [], 'val_left': []})
    right = pd.DataFrame({'key': [], 'val_right': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 62.4μs -> 71.8μs (13.1% slower)

def test_edge_duplicate_keys_in_left():
    # Duplicate join keys in left DataFrame
    left = pd.DataFrame({'id': [1, 1, 2], 'val_left': ['a', 'b', 'c']})
    right = pd.DataFrame({'key': [1, 2], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 135μs -> 69.2μs (95.6% faster)

def test_edge_duplicate_keys_in_both():
    # Duplicate join keys in both DataFrames
    left = pd.DataFrame({'id': [1, 1], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 1], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 136μs -> 68.4μs (100% faster)
    # All combinations present
    left_vals = set(result['val_left'])
    right_vals = set(result['val_right'])

def test_edge_non_overlapping_columns():
    # Columns in right not present in left and vice versa
    left = pd.DataFrame({'id': [1], 'foo': ['bar']})
    right = pd.DataFrame({'key': [1], 'baz': ['qux']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 91.6μs -> 67.5μs (35.7% faster)

def test_edge_left_on_not_in_left():
    # left_on does not exist in left DataFrame
    left = pd.DataFrame({'foo': [1]})
    right = pd.DataFrame({'key': [1], 'val': ['x']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 32.4μs -> 6.79μs (377% faster)

def test_edge_right_on_not_in_right():
    # right_on does not exist in right DataFrame
    left = pd.DataFrame({'id': [1]})
    right = pd.DataFrame({'foo': [1], 'val': ['x']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'key') # 21.0μs -> 8.92μs (136% faster)

def test_edge_null_values_in_join_columns():
    # Null values in join columns should not match
    left = pd.DataFrame({'id': [1, None], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, None], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 111μs -> 68.7μs (62.9% faster)

def test_edge_different_types_in_join_columns():
    # Join columns with different types should not match
    left = pd.DataFrame({'id': [1, '2'], 'val_left': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 2], 'val_right': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 104μs -> 62.1μs (68.5% faster)

def test_edge_empty_column_names():
    # Join columns with empty string as name
    left = pd.DataFrame({'': [1, 2], 'foo': ['a', 'b']})
    right = pd.DataFrame({'': [2, 3], 'bar': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, '', ''); result = codeflash_output # 108μs -> 67.3μs (61.5% faster)

def test_edge_column_name_collision():
    # Right DataFrame has a column with same name as a left column (other than join key)
    left = pd.DataFrame({'id': [1], 'val': ['a']})
    right = pd.DataFrame({'key': [1], 'val': ['b']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 84.8μs -> 60.2μs (40.9% faster)

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

def test_large_scale_all_match():
    # Large DataFrames, all keys match
    n = 500
    left = pd.DataFrame({'id': range(n), 'val_left': list(range(n))})
    right = pd.DataFrame({'key': range(n), 'val_right': list(range(n, 2*n))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 8.22ms -> 580μs (1316% faster)
    # Check a few random rows
    for i in [0, n//2, n-1]:
        row = result[result['id'] == i]

def test_large_scale_no_match():
    # Large DataFrames, no keys match
    n = 500
    left = pd.DataFrame({'id': range(n)})
    right = pd.DataFrame({'key': range(n, 2*n)})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 4.86ms -> 156μs (3013% faster)

def test_large_scale_many_to_many():
    # Many-to-many: each left key matches multiple right keys
    n = 100
    left = pd.DataFrame({'id': [i//2 for i in range(n)], 'val_left': range(n)})
    right = pd.DataFrame({'key': [i//2 for i in range(n)], 'val_right': range(n)})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 2.31ms -> 229μs (908% faster)
    # For each unique key, there are about n/2 left and n/2 right, so n/2 * n/2 matches per key
    # But since keys repeat, total rows should be sum of product of counts per key
    from collections import Counter
    left_counts = Counter(left['id'])
    right_counts = Counter(right['key'])
    expected_rows = sum(left_counts[k] * right_counts[k] for k in left_counts)

def test_large_scale_randomized_keys():
    # Randomized keys, some matches, some not
    n = 500
    keys = random.sample(range(1000), n)
    left_keys = keys[:n//2]
    right_keys = keys[n//4:n*3//4]
    left = pd.DataFrame({'id': left_keys, 'val_left': range(len(left_keys))})
    right = pd.DataFrame({'key': right_keys, 'val_right': range(len(right_keys))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 3.28ms -> 202μs (1518% faster)
    # Only keys in both left_keys and right_keys should match
    expected_keys = set(left_keys) & set(right_keys)

def test_large_scale_string_keys():
    # Large DataFrames with string keys
    n = 500
    def randstr(): return ''.join(random.choices(string.ascii_letters, k=8))
    keys = [randstr() for _ in range(n)]
    left = pd.DataFrame({'id': keys, 'val_left': range(n)})
    right = pd.DataFrame({'key': keys, 'val_right': range(n)})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 11.9ms -> 415μs (2777% faster)

def test_large_scale_performance():
    # Performance test: should not be excessively slow
    n = 900
    left = pd.DataFrame({'id': range(n), 'val_left': range(n)})
    right = pd.DataFrame({'key': range(n), 'val_right': range(n)})
    import time
    start = time.time()
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output # 14.7ms -> 1.03ms (1329% faster)
    elapsed = time.time() - start
# 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-mfejj4uk and push.