⚡️ Speed up function _merge_gen_kwargs by 59%

[codeflash-ai]

📄 59% (0.59x) speedup for _merge_gen_kwargs in src/datasets/utils/sharding.py

⏱️ Runtime : 278 microseconds → 174 microseconds (best of 120 runs)

📝 Explanation and details

The optimized code achieves a 59% runtime improvement (278μs → 174μs) by replacing nested list comprehensions with explicit loops and leveraging Python's built-in list.extend() method.

Key Optimizations

1. Eliminated Nested List Comprehension Overhead

The original code uses a nested list comprehension inside a dict comprehension:

[value for gen_kwargs in gen_kwargs_list for value in gen_kwargs[key]]

This creates significant overhead because:

• Python must repeatedly evaluate the comprehension expression for each key
• The nested structure performs repeated dictionary lookups (gen_kwargs[key])
• The line profiler shows 98.3% of time spent in the comprehension (1.71ms of 1.74ms total)

The optimized version uses explicit loops, allowing the Python interpreter to optimize the iteration more effectively.

2. Leveraged list.extend() for Efficient Merging

The optimized code uses merged.extend(gen_kwargs[key]), which is implemented in C and significantly faster than the comprehension-based concatenation. The line profiler shows this operation takes only 40.4% of total time (581μs of 1.44ms), with the overhead distributed across clearer, more predictable operations.

3. Early Type Detection

By checking isinstance(value, list) once per key using the first dictionary's value, the optimized code avoids redundant type checks in the inner loop. This is more efficient than the original's conditional expression evaluated during comprehension construction.

Test Case Performance

The optimization excels across all test scenarios:

• Single dict cases: 52-83% faster (e.g., single key scalar: 66.5% faster)
• Multiple dict merges: 54-74% faster (e.g., mixed types: 74.1% faster)
• Large-scale operations: 111-250% faster (e.g., 10 dicts with 100 strings each: 250% faster)

The performance gains increase with dataset size, making this especially valuable for the function's use case.

Impact on Workloads

Based on function_references, this function is called in data sharding hot paths:

• ExamplesIterable.shard_data_sources() - used when distributing dataset examples across workers
• ArrowExamplesIterable.shard_data_sources() - used for Arrow-backed dataset sharding

Both call _merge_gen_kwargs() with filtered generator kwargs during dataset loading and multi-process data loading. Since dataset sharding occurs frequently during distributed training and parallel data loading, this 59% speedup directly reduces data pipeline initialization overhead, benefiting any workflow that uses sharded iterable datasets.

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

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
from src.datasets.utils.sharding import _merge_gen_kwargs

def test_basic_concatenation_and_preserve_first_non_list():
    # Basic scenario: two dicts with a list-valued key 'a' should be concatenated
    # and a non-list key 'b' should be taken from the first dict only.
    input_list = [
        {"a": [1, 2], "b": "first"},
        {"a": [3, 4], "b": "second"},
    ]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 2.93μs -> 1.89μs (54.8% faster)

def test_single_element_list_returns_same_structure():
    # With a single-element input list, list-typed keys should return that same list's contents,
    # and non-list keys should return the first value as-is.
    input_list = [{"x": [42, 43], "y": 7}]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 2.76μs -> 1.66μs (65.9% faster)

def test_empty_input_raises_index_error():
    # If an empty list is provided, accessing gen_kwargs_list[0] should raise IndexError.
    with pytest.raises(IndexError):
        _merge_gen_kwargs([]) # 1.56μs -> 1.35μs (15.7% faster)

def test_empty_lists_concatenate_to_empty_and_preserve_first_non_list():
    # When list-typed keys are empty across inputs, merged result should be empty list.
    input_list = [
        {"a": [], "b": "keep"},
        {"a": [], "b": "ignore"},
    ]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 2.81μs -> 1.77μs (58.6% faster)

def test_inconsistent_types_second_not_iterable_raises_type_error():
    # If the first dict has a list for a key but a subsequent dict has a non-iterable
    # for the same key (e.g., int), the list-comprehension will attempt to iterate the int
    # and a TypeError should be raised.
    input_list = [
        {"k": [1, 2]},
        {"k": 5},  # not iterable -> should cause TypeError during merge
    ]
    with pytest.raises(TypeError):
        _merge_gen_kwargs(input_list) # 3.41μs -> 2.52μs (35.3% faster)

def test_non_list_first_ignores_others_and_returns_first_reference():
    # If the first dict's value for a key is non-list, the function must return that first value
    # unchanged, ignoring subsequent dict values entirely.
    input_list = [
        {"k": "first_value"},
        {"k": ["this", "is", "ignored"]},
    ]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 1.97μs -> 1.19μs (65.4% faster)

def test_merged_list_is_independent_of_original_list_append():
    # The merged result for list-typed keys should be a newly created list (contents copied).
    # Appending to the original per-dict list AFTER calling the function should not mutate the merged result.
    mutable = [100]
    input_list = [
        {"vals": mutable},
        {"vals": [200]},
    ]
    codeflash_output = _merge_gen_kwargs(input_list); merged = codeflash_output # 2.55μs -> 1.51μs (68.9% faster)
    # mutate the original list by appending a new element
    mutable.append(999)

def test_non_list_reference_is_same_object_and_reflects_mutation():
    # For non-list keys, the function returns the original object from the first dict.
    # If that object is mutable (e.g., a dict), modifying it afterwards should be visible via result.
    first_obj = {"inner": 1}
    input_list = [
        {"cfg": first_obj},
        {"cfg": {"inner": 2}},
    ]
    codeflash_output = _merge_gen_kwargs(input_list); merged = codeflash_output # 2.02μs -> 1.31μs (54.0% faster)
    # Mutate the referenced object and verify the change is visible through merged result
    first_obj["inner"] = 999

def test_large_scale_merge_concatenates_many_small_lists():
    # Large-ish input under the 1000-iteration guideline:
    # create 500 dicts each with a single-element list for key 'a' and ensure merged list has 500 elements.
    n = 500  # number of dicts, safe under 1000
    input_list = [{"a": [i], "meta": "first"} for i in range(n)]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 32.8μs -> 28.9μs (13.6% faster)

def test_large_scale_multiple_keys_and_small_sublists():
    # Another larger test: multiple keys where each list contains 2 elements but fewer dicts,
    # ensuring total iterations remain reasonable (here: 300 dicts * 2 elements = 600 iterations).
    dicts_count = 300
    input_list = [
        {"k1": [i, i + 1000], "k2": ["const"], "flag": False}
        for i in range(dicts_count)
    ]
    codeflash_output = _merge_gen_kwargs(input_list); result = codeflash_output # 45.4μs -> 35.2μs (29.0% faster)
    # k1 must be concatenation of all sublists in order
    expected_k1 = []
    for i in range(dicts_count):
        expected_k1.extend([i, i + 1000])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from src.datasets.utils.sharding import _merge_gen_kwargs

class TestMergeGenKwargsBasic:
    """Basic test cases for _merge_gen_kwargs function."""

    def test_single_dict_with_list_values(self):
        """Test merging a single dictionary containing list values."""
        gen_kwargs_list = [{"key1": [1, 2], "key2": [3, 4]}]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.96μs -> 1.81μs (63.3% faster)

    def test_single_dict_with_scalar_values(self):
        """Test merging a single dictionary containing scalar (non-list) values."""
        gen_kwargs_list = [{"key1": "value1", "key2": 42}]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.04μs -> 1.34μs (52.5% faster)

    def test_two_dicts_with_list_values(self):
        """Test merging two dictionaries with list values."""
        gen_kwargs_list = [
            {"key1": [1, 2], "key2": [3, 4]},
            {"key1": [5, 6], "key2": [7, 8]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 3.25μs -> 2.04μs (58.8% faster)

    def test_two_dicts_with_scalar_values(self):
        """Test merging two dictionaries with scalar values preserves first dict's values."""
        gen_kwargs_list = [
            {"key1": "value1", "key2": 42},
            {"key1": "value2", "key2": 100}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.02μs -> 1.31μs (54.2% faster)

    def test_single_key_single_element_list(self):
        """Test merging dictionaries with single-element lists."""
        gen_kwargs_list = [
            {"key": [10]},
            {"key": [20]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.55μs -> 1.52μs (68.2% faster)

    def test_three_dicts_with_list_values(self):
        """Test merging three dictionaries with list values."""
        gen_kwargs_list = [
            {"data": [1], "names": ["a"]},
            {"data": [2], "names": ["b"]},
            {"data": [3], "names": ["c"]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 3.37μs -> 2.10μs (60.2% faster)

    def test_mixed_types_in_lists(self):
        """Test merging lists containing mixed types."""
        gen_kwargs_list = [
            {"mixed": [1, "string", 3.14]},
            {"mixed": [True, None, {}]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.71μs -> 1.56μs (74.1% faster)

class TestMergeGenKwargsEdgeCases:
    """Edge case test cases for _merge_gen_kwargs function."""

    def test_empty_list_values(self):
        """Test merging dictionaries with empty list values."""
        gen_kwargs_list = [
            {"key": []},
            {"key": []}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.39μs -> 1.48μs (61.7% faster)

    def test_one_dict_empty_list_one_dict_with_values(self):
        """Test merging where one dict has empty lists and another has values."""
        gen_kwargs_list = [
            {"key": []},
            {"key": [1, 2, 3]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.56μs -> 1.49μs (72.1% faster)

    def test_numeric_scalar_values(self):
        """Test with numeric scalar values from first dictionary."""
        gen_kwargs_list = [
            {"count": 10, "offset": 5},
            {"count": 20, "offset": 15}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.00μs -> 1.35μs (48.1% faster)

    def test_none_scalar_value(self):
        """Test with None as a scalar value in first dictionary."""
        gen_kwargs_list = [
            {"value": None},
            {"value": [1, 2, 3]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 1.86μs -> 1.05μs (78.0% faster)

    def test_boolean_scalar_value(self):
        """Test with boolean scalar values."""
        gen_kwargs_list = [
            {"flag": True, "enabled": False},
            {"flag": False, "enabled": True}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.45μs -> 1.81μs (35.4% faster)

    def test_string_scalar_value(self):
        """Test with string scalar values."""
        gen_kwargs_list = [
            {"mode": "train", "format": "json"},
            {"mode": "test", "format": "csv"}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.06μs -> 1.39μs (48.8% faster)

    def test_single_element_in_each_list(self):
        """Test with lists containing single elements across multiple dicts."""
        gen_kwargs_list = [
            {"ids": [100]},
            {"ids": [200]},
            {"ids": [300]},
            {"ids": [400]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.75μs -> 1.75μs (57.6% faster)

    def test_many_dicts_with_list_values(self):
        """Test merging many dictionaries with list values."""
        gen_kwargs_list = [
            {"items": [i]} for i in range(10)
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 3.22μs -> 2.16μs (49.3% faster)

    def test_dict_with_single_key_list(self):
        """Test single dictionary with one key containing a list."""
        gen_kwargs_list = [{"only_key": [1, 2, 3, 4, 5]}]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.60μs -> 1.42μs (83.4% faster)

    def test_dict_with_single_key_scalar(self):
        """Test single dictionary with one key containing a scalar."""
        gen_kwargs_list = [{"only_key": "only_value"}]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 1.81μs -> 1.08μs (66.5% faster)

    def test_nested_lists_not_flattened(self):
        """Test that nested lists are treated as single values and not flattened."""
        gen_kwargs_list = [
            {"nested": [[1, 2], [3, 4]]},
            {"nested": [[5, 6], [7, 8]]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.67μs -> 1.62μs (64.6% faster)

    def test_tuple_as_scalar_value(self):
        """Test tuple treated as a scalar (non-list) value."""
        gen_kwargs_list = [
            {"coord": (1, 2)},
            {"coord": (3, 4)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.14μs -> 1.29μs (65.9% faster)

    def test_dict_object_as_value_in_list(self):
        """Test dictionaries stored within lists are properly merged."""
        gen_kwargs_list = [
            {"configs": [{"a": 1}]},
            {"configs": [{"b": 2}]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.57μs -> 1.61μs (60.1% faster)

    def test_multiple_keys_with_different_value_types(self):
        """Test multiple keys where some have lists and some have scalars."""
        gen_kwargs_list = [
            {"list_key": [1, 2], "scalar_key": "value", "another_list": ["x"]},
            {"list_key": [3, 4], "scalar_key": "other", "another_list": ["y"]}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 3.47μs -> 2.32μs (49.5% faster)
        expected = {
            "list_key": [1, 2, 3, 4],
            "scalar_key": "value",
            "another_list": ["x", "y"]
        }

class TestMergeGenKwargsLargeScale:
    """Large scale test cases for _merge_gen_kwargs function."""

    def test_many_dicts_with_many_list_elements(self):
        """Test merging many dictionaries with lists containing many elements."""
        # Create 50 dictionaries, each with lists of 10 elements
        gen_kwargs_list = [
            {"values": list(range(i * 10, (i + 1) * 10))}
            for i in range(50)
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 14.1μs -> 6.71μs (111% faster)

    def test_large_number_of_dicts_with_scalar_values(self):
        """Test merging 100 dictionaries with scalar values."""
        gen_kwargs_list = [
            {"id": 1, "name": "first", "count": 100}
            for _ in range(100)
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.22μs -> 1.56μs (42.5% faster)

    def test_multiple_large_lists_merged(self):
        """Test merging dictionaries with multiple large lists."""
        dict_count = 30
        list_size = 20
        gen_kwargs_list = [
            {
                "list1": list(range(i * list_size, (i + 1) * list_size)),
                "list2": list(range(100 + i * list_size, 100 + (i + 1) * list_size)),
                "list3": list(range(200 + i * list_size, 200 + (i + 1) * list_size))
            }
            for i in range(dict_count)
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 38.5μs -> 13.0μs (196% faster)

    def test_large_string_lists(self):
        """Test merging dictionaries with large lists of strings."""
        gen_kwargs_list = [
            {"names": [f"name_{j}_{i}" for j in range(100)]}
            for i in range(10)
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 18.8μs -> 5.37μs (250% faster)

    def test_many_keys_in_dictionary(self):
        """Test with dictionaries containing many keys with list values."""
        # Create a dictionary with 50 keys, each with a list
        gen_kwargs_list = [
            {f"key_{i}": [i, i+1, i+2] for i in range(50)},
            {f"key_{i}": [i+10, i+11, i+12] for i in range(50)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 26.5μs -> 17.6μs (50.6% faster)
        for i in range(50):
            pass

    def test_large_mixed_types_in_lists(self):
        """Test merging large lists with mixed types."""
        values = [1, "string", 3.14, None, True, {"nested": "dict"}, [1, 2, 3]]
        gen_kwargs_list = [
            {"mixed": values.copy() for _ in range(30)},
            {"mixed": values.copy() for _ in range(30)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.86μs -> 1.62μs (76.7% faster)

    def test_dict_with_100_scalar_keys(self):
        """Test dictionary with 100 scalar keys."""
        gen_kwargs_list = [
            {f"scalar_{i}": i for i in range(100)},
            {f"scalar_{i}": i + 1000 for i in range(100)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 18.0μs -> 14.9μs (20.3% faster)
        for i in range(100):
            pass

    def test_performance_with_250_dicts(self):
        """Test that function handles 250 dictionaries efficiently."""
        # Create 250 dictionaries with lists
        gen_kwargs_list = [
            {"data": [i] for i in range(250)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.43μs -> 1.35μs (80.5% faster)

    def test_deeply_nested_structure_in_list(self):
        """Test with complex nested structures in lists."""
        gen_kwargs_list = [
            {"complex": [{"level1": {"level2": {"level3": i}}}] for i in range(50)},
            {"complex": [{"level1": {"level2": {"level3": i + 100}}}] for i in range(50)}
        ]
        codeflash_output = _merge_gen_kwargs(gen_kwargs_list); result = codeflash_output # 2.54μs -> 1.50μs (69.3% 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-_merge_gen_kwargs-mlcxjhbn and push.