⚡️ Speed up function dataframe_merge by 1,247%

[codeflash-ai]

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

⏱️ Runtime : 404 milliseconds → 30.0 milliseconds (best of 161 runs)

📝 Explanation and details

Here is an optimized version of your program, keeping the logic, function signature, and all behaviors identical. I replaced the slow, repeated use of .iloc[] with direct NumPy array access, batched per-column lookups, and rewrote the merge loop with list comprehensions and index-based lookups. This way, the function avoids thousands of slow Pandas Series creation steps, and directly accesses the data under the hood.

All comments are kept verbatim (none existed before). Only internal algorithm and data structure are changed.

Key optimizations.

• Vectorized access: Operating directly on NumPy arrays via .values, which is much faster than pandas .iloc[].
• Avoid repeated Series/dict creation: Build the base 'left row' dict once, copy for each matching join.
• Precompute column indices: Avoid repeated get_loc lookups in the merge loops.

This will typically result in 10-40x speedup for medium-to-large dataframes. All previous functionality is preserved.

✅ 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 Details

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_single_match():
    # Merge on a single matching row
    left = pd.DataFrame({'id': [1], 'val': ['a']})
    right = pd.DataFrame({'key': [1], 'data': ['x']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1], 'val': ['a'], 'data': ['x']})

def test_basic_multiple_matches():
    # Merge where multiple rows match on the key
    left = pd.DataFrame({'id': [1, 2], 'val': ['a', 'b']})
    right = pd.DataFrame({'key': [2, 1], 'data': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 2], 'val': ['a', 'b'], 'data': ['y', 'x']})

def test_basic_duplicate_keys_in_right():
    # Right dataframe has duplicate keys, should produce multiple rows per match
    left = pd.DataFrame({'id': [1], 'val': ['a']})
    right = pd.DataFrame({'key': [1, 1], 'data': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 1], 'val': ['a', 'a'], 'data': ['x', 'y']})

def test_basic_duplicate_keys_in_left():
    # Left dataframe has duplicate keys, should produce multiple rows per match
    left = pd.DataFrame({'id': [1, 1], 'val': ['a', 'b']})
    right = pd.DataFrame({'key': [1], 'data': ['x']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 1], 'val': ['a', 'b'], 'data': ['x', 'x']})

def test_basic_no_matches():
    # No matching keys, should return empty DataFrame
    left = pd.DataFrame({'id': [1, 2], 'val': ['a', 'b']})
    right = pd.DataFrame({'key': [3, 4], 'data': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # Should have columns from left and right (excluding right_on)
    expected_cols = ['id', 'val', 'data']

def test_basic_column_overlap():
    # Both dataframes have a column with the same name (not the key)
    left = pd.DataFrame({'id': [1], 'val': ['a'], 'shared': [10]})
    right = pd.DataFrame({'key': [1], 'shared': [20], 'data': ['x']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # The 'shared' column from right should overwrite left's if not handled, but in this implementation,
    # left's 'shared' comes first, then right's 'shared' is added (with same name), so right's will overwrite.
    # But since right_cols excludes the key, 'shared' is included.
    expected = pd.DataFrame({'id': [1], 'val': ['a'], 'shared': [20], 'data': ['x']})

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

def test_edge_empty_left():
    # Left dataframe is empty
    left = pd.DataFrame({'id': [], 'val': []})
    right = pd.DataFrame({'key': [1, 2], 'data': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected_cols = ['id', 'val', 'data']

def test_edge_empty_right():
    # Right dataframe is empty
    left = pd.DataFrame({'id': [1, 2], 'val': ['a', 'b']})
    right = pd.DataFrame({'key': [], 'data': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected_cols = ['id', 'val', 'data']

def test_edge_both_empty():
    # Both dataframes are empty
    left = pd.DataFrame({'id': [], 'val': []})
    right = pd.DataFrame({'key': [], 'data': []})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected_cols = ['id', 'val', 'data']

def test_edge_no_common_column_names():
    # No columns in common except the join keys
    left = pd.DataFrame({'id': [1, 2], 'foo': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 2], 'bar': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 2], 'foo': ['a', 'b'], 'bar': ['x', 'y']})

def test_edge_non_string_column_names():
    # Use non-string column names (e.g., integers)
    left = pd.DataFrame({0: [1, 2], 1: ['a', 'b']})
    right = pd.DataFrame({2: [1, 2], 3: ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 0, 2); result = codeflash_output
    expected = pd.DataFrame({0: [1, 2], 1: ['a', 'b'], 3: ['x', 'y']})

def test_edge_null_values_in_keys():
    # Null values in key columns should not match
    left = pd.DataFrame({'id': [1, None, 2], 'val': ['a', 'b', 'c']})
    right = pd.DataFrame({'key': [1, 2, None], 'data': ['x', 'y', 'z']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 2], 'val': ['a', 'c'], 'data': ['x', 'y']})

def test_edge_non_unique_keys():
    # Both sides have non-unique keys, should produce cartesian product for matches
    left = pd.DataFrame({'id': [1, 1], 'val': ['a', 'b']})
    right = pd.DataFrame({'key': [1, 1], 'data': ['x', 'y']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # Should produce 4 rows: (a,x), (a,y), (b,x), (b,y)
    expected = pd.DataFrame({
        'id': [1, 1, 1, 1],
        'val': ['a', 'a', 'b', 'b'],
        'data': ['x', 'y', 'x', 'y']
    })

def test_edge_key_column_missing():
    # Should raise KeyError if join key is missing in either dataframe
    left = pd.DataFrame({'foo': [1]})
    right = pd.DataFrame({'bar': [1]})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'id', 'bar')
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'foo', 'id')

def test_edge_key_column_with_nan():
    # NaN in key columns should not match
    import math
    left = pd.DataFrame({'id': [1, float('nan'), 2], 'val': ['a', 'b', 'c']})
    right = pd.DataFrame({'key': [1, 2, float('nan')], 'data': ['x', 'y', 'z']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected = pd.DataFrame({'id': [1, 2], 'val': ['a', 'c'], 'data': ['x', 'y']})

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

def test_large_scale_many_rows():
    # Merge two dataframes with 1000 rows each, matching on a single key
    left = pd.DataFrame({'id': list(range(1000)), 'val': [str(i) for i in range(1000)]})
    right = pd.DataFrame({'key': list(range(1000)), 'data': [str(i*2) for i in range(1000)]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # Spot check a few rows
    for i in [0, 499, 999]:
        row = result[result['id'] == i].iloc[0]

def test_large_scale_cartesian():
    # Both dataframes have 100 rows with the same key, should produce 100*100 = 10000 rows
    left = pd.DataFrame({'id': [1]*100, 'val': list(range(100))})
    right = pd.DataFrame({'key': [1]*100, 'data': list(range(100))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # Each left row should appear 100 times, each right row should appear 100 times
    left_counts = result['val'].value_counts()
    right_counts = result['data'].value_counts()

def test_large_scale_sparse_matches():
    # Only a few keys match out of large dataframes
    left = pd.DataFrame({'id': list(range(1000)), 'val': [str(i) for i in range(1000)]})
    right = pd.DataFrame({'key': [100, 500, 900], 'data': ['x', 'y', 'z']})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    # Only 3 rows should match
    expected = pd.DataFrame({
        'id': [100, 500, 900],
        'val': ['100', '500', '900'],
        'data': ['x', 'y', 'z']
    })

def test_large_scale_no_matches():
    # Large dataframes with no matches
    left = pd.DataFrame({'id': list(range(1000)), 'val': [str(i) for i in range(1000)]})
    right = pd.DataFrame({'key': list(range(1000, 2000)), 'data': [str(i) for i in range(1000, 2000)]})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
    expected_cols = ['id', 'val', 'data']

def test_large_scale_column_overlap():
    # Large dataframes with overlapping column names (not the key)
    left = pd.DataFrame({'id': list(range(100)), 'val': list(range(100)), 'shared': ['L']*100})
    right = pd.DataFrame({'key': list(range(100)), 'shared': ['R']*100, 'data': list(range(100, 200))})
    codeflash_output = dataframe_merge(left, right, 'id', 'key'); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

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 case: one row in each, one match
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [1], 'D': ['y']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_basic_inner_join_multiple_matches():
    # Multiple matching rows
    left = pd.DataFrame({'A': [1, 2], 'B': ['x', 'z']})
    right = pd.DataFrame({'C': [1, 2], 'D': ['y', 'w']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_basic_inner_join_duplicate_keys():
    # Duplicate keys in right, should produce cartesian product for matches
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [1, 1], 'D': ['y', 'z']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_basic_inner_join_duplicate_keys_left():
    # Duplicate keys in left, should produce cartesian product for matches
    left = pd.DataFrame({'A': [1, 1], 'B': ['x', 'w']})
    right = pd.DataFrame({'C': [1], 'D': ['y']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_basic_no_match():
    # No matching keys
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [2], 'D': ['y']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

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

def test_empty_left():
    # Left DataFrame is empty
    left = pd.DataFrame({'A': [], 'B': []})
    right = pd.DataFrame({'C': [1], 'D': ['y']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_empty_right():
    # Right DataFrame is empty
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [], 'D': []})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_both_empty():
    # Both DataFrames are empty
    left = pd.DataFrame({'A': [], 'B': []})
    right = pd.DataFrame({'C': [], 'D': []})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_column_name_collision():
    # Both DataFrames have a column (not join key) with the same name
    left = pd.DataFrame({'A': [1], 'X': ['foo']})
    right = pd.DataFrame({'C': [1], 'X': ['bar']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_non_overlapping_columns():
    # DataFrames with no columns in common except join keys
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [1], 'D': ['y']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_join_on_nonexistent_column():
    # Join on column that does not exist
    left = pd.DataFrame({'A': [1], 'B': ['x']})
    right = pd.DataFrame({'C': [1], 'D': ['y']})
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'Z', 'C')
    with pytest.raises(KeyError):
        dataframe_merge(left, right, 'A', 'Z')

def test_null_values_in_keys():
    # Null (None/NaN) values in join columns
    left = pd.DataFrame({'A': [1, None], 'B': ['x', 'z']})
    right = pd.DataFrame({'C': [1, None], 'D': ['y', 'w']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_different_data_types_in_keys():
    # Join keys with different types (should not match)
    left = pd.DataFrame({'A': [1, '2'], 'B': ['x', 'y']})
    right = pd.DataFrame({'C': [1, 2], 'D': ['z', 'w']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_multiple_duplicate_keys():
    # Multiple duplicate keys in both left and right
    left = pd.DataFrame({'A': [1, 1, 2], 'B': ['a', 'b', 'c']})
    right = pd.DataFrame({'C': [1, 1, 2], 'D': ['x', 'y', 'z']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_mixed_types_in_other_columns():
    # Columns with mixed types (object dtype)
    left = pd.DataFrame({'A': [1, 2], 'B': ['x', 42]})
    right = pd.DataFrame({'C': [1, 2], 'D': [None, 'w']})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

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

def test_large_scale_many_rows():
    # Large DataFrames with 1000 rows each, keys overlap on even numbers
    n = 1000
    left = pd.DataFrame({'A': list(range(n)), 'B': ['l']*n})
    right = pd.DataFrame({'C': list(range(0, n, 2)), 'D': ['r']*(n//2)})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_large_scale_cartesian_product():
    # Many duplicate keys: left has 100 rows of key 1, right has 10 rows of key 1
    n_left = 100
    n_right = 10
    left = pd.DataFrame({'A': [1]*n_left, 'B': range(n_left)})
    right = pd.DataFrame({'C': [1]*n_right, 'D': range(n_right)})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_large_scale_no_overlap():
    # Two large DataFrames, but no keys overlap
    n = 500
    left = pd.DataFrame({'A': list(range(n)), 'B': ['x']*n})
    right = pd.DataFrame({'C': list(range(n, 2*n)), 'D': ['y']*n})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output

def test_large_scale_column_name_collision():
    # Both DataFrames have a large number of columns, with some collisions
    n = 100
    left_cols = {f'L{i}': list(range(n)) for i in range(10)}
    left_cols['K'] = list(range(n))
    right_cols = {f'R{i}': list(range(n)) for i in range(10)}
    right_cols['K'] = list(range(n))
    left = pd.DataFrame(left_cols)
    right = pd.DataFrame(right_cols)
    # Join on 'K'
    codeflash_output = dataframe_merge(left, right, 'K', 'K'); result = codeflash_output
    # Should have all columns from left, and all right columns except 'K'
    expected_cols = set(left.columns) | (set(right.columns) - {'K'})

def test_large_scale_null_keys():
    # Many rows, some with null keys
    n = 500
    left = pd.DataFrame({'A': [None if i % 10 == 0 else i for i in range(n)], 'B': range(n)})
    right = pd.DataFrame({'C': [None if i % 10 == 0 else i for i in range(n)], 'D': range(n)})
    codeflash_output = dataframe_merge(left, right, 'A', 'C'); result = codeflash_output
    # For each i not divisible by 10, should match exactly one row
    # For i divisible by 10, None matches None, so for each None in left, matches all None in right
    expected_non_null_matches = n - n//10
    expected_nulls = n//10 * n//10
    # Check that all rows with None in 'A' also have None in 'C' (since that's the only way they match)
    null_rows = result[result['A'].isnull()]
# 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-mbhmnt7d and push.