⚡️ Speed up function apply_dict_of_variables_vfunc by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for apply_dict_of_variables_vfunc in xarray/core/computation.py

⏱️ Runtime : 4.97 milliseconds → 4.50 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 10% speedup through several key micro-optimizations that reduce redundant computations and attribute lookups:

Core optimizations:

1. collect_dict_values: The most impactful change - caches is_dict_like results to avoid repeated calls per object-key pair. When all objects are dict-like (common case), uses a fast path with direct .get() calls. This eliminates O(n*k) is_dict_like calls where n=objects, k=keys.

2. _as_variables_or_variable: Replaces nested try/except with getattr(..., None) checks, which are faster in Python since they avoid exception handling overhead.

3. _check_core_dims: Hoists signature.input_core_dims lookup outside the loop and uses getattr instead of hasattr + attribute access, reducing repeated attribute lookups.

4. join_dict_keys: Eliminates intermediate list creation by building all_keys incrementally rather than using a list comprehension.

Performance impact by test case:

• Large-scale tests show the biggest gains (9-14% faster) where the optimizations compound
• Basic operations with few keys see moderate improvements (5-7% faster)
• Edge cases with missing dimensions benefit from faster attribute access (2-10% faster)

Hot path relevance:
Based on the function reference showing apply_dict_of_variables_vfunc is called from apply_dataset_vfunc, this optimization benefits xarray's core dataset operations. The 10% improvement becomes significant when processing large datasets with many variables, as these functions are in the critical path for most xarray computations.

The optimizations particularly excel when processing datasets with many variables (500+ keys showed 13-14% improvements) or when objects are consistently dict-like, making this a valuable optimization for typical xarray workloads.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections import OrderedDict

# imports
import pytest
from xarray.core.computation import apply_dict_of_variables_vfunc


# --- Function and helpers to test ---
# Minimal stubs for Variable and _UFuncSignature, since we cannot import xarray here.
class Variable:
    def __init__(self, dims, data):
        self.dims = tuple(dims)
        self.data = data

    def __eq__(self, other):
        return (
            isinstance(other, Variable)
            and self.dims == other.dims
            and self.data == other.data
        )

    def __repr__(self):
        return f"Variable(dims={self.dims}, data={self.data})"


class _UFuncSignature:
    def __init__(self, input_core_dims, num_outputs=1):
        self.input_core_dims = input_core_dims
        self.num_outputs = num_outputs


from xarray.core.computation import apply_dict_of_variables_vfunc

# --- Unit Tests ---

# Basic Test Cases


def test_basic_two_dicts_inner_join():
    # Two dicts, inner join, both have 'a'
    d1 = {"a": Variable(["x"], [1, 2]), "b": Variable(["y"], [3, 4])}
    d2 = {"a": Variable(["x"], [10, 20]), "c": Variable(["z"], [30, 40])}

    def func(v1, v2):
        return Variable(v1.dims, [v1.data[0] + v2.data[0], v1.data[1] + v2.data[1]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(func, d1, d2, signature=sig)
    result = codeflash_output  # 17.9μs -> 16.7μs (6.85% faster)


def test_basic_outer_join_with_fill_value():
    # Outer join, fill_value supplied
    d1 = {"a": Variable(["x"], [1, 2])}
    d2 = {"b": Variable(["y"], [3, 4])}

    def func(v1, v2):
        # If v1 or v2 is fill_value, return Variable with data=[-1]
        if v1 == "fill" or v2 == "fill":
            return Variable([], [-1])
        return Variable(v1.dims, [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["y"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, join="outer", fill_value="fill"
    )
    result = codeflash_output  # 17.4μs -> 16.4μs (6.13% faster)


def test_basic_multiple_outputs():
    # Function returns tuple, num_outputs=2
    d1 = {"a": Variable(["x"], [1, 2])}
    d2 = {"a": Variable(["x"], [10, 20])}

    def func(v1, v2):
        return (
            Variable(v1.dims, [v1.data[0] + v2.data[0], v1.data[1] + v2.data[1]]),
            Variable(v1.dims, [v1.data[0] * v2.data[0], v1.data[1] * v2.data[1]]),
        )

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=2)
    out1, out2 = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig
    )  # 18.0μs -> 18.1μs (0.276% slower)


def test_basic_exact_join_fail():
    # exact join, keys mismatch should raise
    d1 = {"a": Variable(["x"], [1]), "b": Variable(["y"], [2])}
    d2 = {"a": Variable(["x"], [3]), "c": Variable(["z"], [4])}

    def func(v1, v2):
        return Variable(v1.dims, [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    with pytest.raises(ValueError):
        apply_dict_of_variables_vfunc(
            func, d1, d2, signature=sig, join="exact"
        )  # 13.0μs -> 10.9μs (19.4% faster)


# Edge Test Cases


def test_edge_empty_dicts():
    # Both dicts empty
    d1 = {}
    d2 = {}

    def func(v1, v2):
        return Variable([], [0])

    sig = _UFuncSignature([[], []], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(func, d1, d2, signature=sig)
    result = codeflash_output  # 9.72μs -> 9.74μs (0.277% slower)


def test_edge_one_empty_one_nonempty():
    # One dict empty, one non-empty, inner join
    d1 = {}
    d2 = {"a": Variable(["x"], [1])}

    def func(v1, v2):
        return Variable(["x"], [v2.data[0]])

    sig = _UFuncSignature([[], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(func, d1, d2, signature=sig)
    result = codeflash_output  # 9.92μs -> 9.59μs (3.45% faster)


def test_edge_fill_value_with_missing_keys():
    # Outer join, fill_value used, function must handle fill_value
    d1 = {"a": Variable(["x"], [1])}
    d2 = {}

    def func(v1, v2):
        if v2 == "missing":
            return Variable(["x"], [-1])
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, join="outer", fill_value="missing"
    )
    result = codeflash_output  # 14.0μs -> 12.9μs (8.64% faster)


def test_edge_missing_core_dim_raise():
    # Variable missing required core dim, on_missing_core_dim='raise'
    d1 = {"a": Variable(["x"], [1])}
    d2 = {"a": Variable(["y"], [2])}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    with pytest.raises(ValueError) as e:
        apply_dict_of_variables_vfunc(
            func, d1, d2, signature=sig, on_missing_core_dim="raise"
        )  # 19.4μs -> 19.0μs (1.92% faster)


def test_edge_missing_core_dim_copy():
    # Variable missing required core dim, on_missing_core_dim='copy'
    d1 = {"a": Variable(["x"], [1])}
    d2 = {"a": Variable(["y"], [2])}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, on_missing_core_dim="copy"
    )
    result = codeflash_output  # 18.6μs -> 18.7μs (0.438% slower)


def test_edge_missing_core_dim_drop():
    # Variable missing required core dim, on_missing_core_dim='drop'
    d1 = {"a": Variable(["x"], [1])}
    d2 = {"a": Variable(["y"], [2])}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, on_missing_core_dim="drop"
    )
    result = codeflash_output  # 19.5μs -> 17.7μs (10.5% faster)


def test_edge_invalid_on_missing_core_dim_value():
    # Invalid value for on_missing_core_dim should raise
    d1 = {"a": Variable(["x"], [1])}
    d2 = {"a": Variable(["y"], [2])}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    with pytest.raises(ValueError):
        apply_dict_of_variables_vfunc(
            func, d1, d2, signature=sig, on_missing_core_dim="invalid"
        )  # 19.7μs -> 18.9μs (4.04% faster)


def test_large_scale_many_keys():
    # Many keys, check performance and correctness
    N = 500
    d1 = {f"k{i}": Variable(["x"], [i]) for i in range(N)}
    d2 = {f"k{i}": Variable(["x"], [2 * i]) for i in range(N)}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(func, d1, d2, signature=sig)
    result = codeflash_output  # 705μs -> 646μs (9.06% faster)
    for i in range(N):
        key = f"k{i}"


def test_large_scale_missing_core_dim_drop():
    # Many keys, some missing core dims, drop
    N = 100
    d1 = {f"k{i}": Variable(["x"], [i]) for i in range(N)}
    d2 = {f"k{i}": Variable(["y"], [2 * i]) for i in range(N)}

    def func(v1, v2):
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, on_missing_core_dim="drop"
    )
    result = codeflash_output  # 242μs -> 227μs (6.69% faster)


def test_large_scale_outer_join_with_fill_value():
    # Many keys, outer join, fill_value
    N = 200
    d1 = {f"k{i}": Variable(["x"], [i]) for i in range(N)}
    d2 = {f"k{i}": Variable(["x"], [2 * i]) for i in range(N // 2, N + N // 2)}

    def func(v1, v2):
        if v1 == "fill" or v2 == "fill":
            return Variable([], [-1])
        return Variable(["x"], [v1.data[0] + v2.data[0]])

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=1)
    codeflash_output = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig, join="outer", fill_value="fill"
    )
    result = codeflash_output  # 411μs -> 360μs (13.9% faster)
    # All keys from both dicts present
    expected_keys = set(
        [f"k{i}" for i in range(N)] + [f"k{i}" for i in range(N // 2, N + N // 2)]
    )
    # Keys only in one dict should have Variable([], [-1])
    for key in expected_keys:
        if key not in d1 or key not in d2:
            pass


def test_large_scale_multiple_outputs():
    # Many keys, multiple outputs
    N = 50
    d1 = {f"k{i}": Variable(["x"], [i]) for i in range(N)}
    d2 = {f"k{i}": Variable(["x"], [2 * i]) for i in range(N)}

    def func(v1, v2):
        return (
            Variable(["x"], [v1.data[0] + v2.data[0]]),
            Variable(["x"], [v1.data[0] * v2.data[0]]),
        )

    sig = _UFuncSignature([["x"], ["x"]], num_outputs=2)
    out1, out2 = apply_dict_of_variables_vfunc(
        func, d1, d2, signature=sig
    )  # 110μs -> 106μs (3.90% faster)
    for i in range(N):
        key = f"k{i}"

import operator

# imports
import pytest
from xarray.core.computation import apply_dict_of_variables_vfunc


# Minimal stubs for Variable and _UFuncSignature, since we cannot import xarray
class Variable:
    def __init__(self, dims, data):
        self.dims = tuple(dims)
        self.data = data

    def __eq__(self, other):
        # Compare dims and data for equality
        return (
            isinstance(other, Variable)
            and self.dims == other.dims
            and self.data == other.data
        )

    def __repr__(self):
        return f"Variable(dims={self.dims}, data={self.data})"


class _UFuncSignature:
    def __init__(self, input_core_dims, num_outputs=1):
        self.input_core_dims = input_core_dims
        self.num_outputs = num_outputs


from xarray.core.computation import apply_dict_of_variables_vfunc

# ------------------- UNIT TESTS -------------------

# Basic Test Cases


def test_basic_single_dict_addition():
    # Test with two dicts with same keys and simple addition
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = {"x": Variable(("dim",), 10), "y": Variable(("dim",), 20)}

    def add_func(v1, v2):
        # Add data, keep dims from v1
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature
    )
    result = codeflash_output  # 17.8μs -> 16.6μs (7.42% faster)


def test_basic_fill_value():
    # Test with fill_value for missing keys
    a = {"x": Variable(("dim",), 1)}
    b = {"x": Variable(("dim",), 10), "y": Variable(("dim",), 20)}

    def add_func(v1, v2):
        # If v1 is fill_value, treat as 0
        v1_data = v1.data if isinstance(v1, Variable) else 0
        v2_data = v2.data if isinstance(v2, Variable) else 0
        return Variable(("dim",), v1_data + v2_data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func,
        a,
        b,
        signature=signature,
        join="outer",
        fill_value=Variable(("dim",), 0),
    )
    result = codeflash_output  # 16.6μs -> 15.7μs (5.59% faster)


def test_basic_single_dict_with_ndarray():
    # Test with dict and a single Variable (not dict)
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = Variable(("dim",), 10)

    def add_func(v1, v2):
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature
    )
    result = codeflash_output  # 15.6μs -> 16.6μs (6.10% slower)


def test_basic_exact_join():
    # Test with exact join
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = {"x": Variable(("dim",), 10), "y": Variable(("dim",), 20)}

    def add_func(v1, v2):
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, join="exact"
    )
    result = codeflash_output  # 15.0μs -> 14.1μs (6.16% faster)


# Edge Test Cases


def test_missing_core_dim_raises():
    # Test missing core dim with 'raise'
    a = {"x": Variable(("dim",), 1)}
    b = {"x": Variable(("other_dim",), 10)}

    def add_func(v1, v2):
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    with pytest.raises(ValueError) as excinfo:
        apply_dict_of_variables_vfunc(
            add_func, a, b, signature=signature
        )  # 19.5μs -> 19.0μs (2.63% faster)


def test_missing_core_dim_copy():
    # Test missing core dim with 'copy'
    a = {"x": Variable(("dim",), 1)}
    b = {"x": Variable(("other_dim",), 10)}

    def add_func(v1, v2):
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, on_missing_core_dim="copy"
    )
    result = codeflash_output  # 19.0μs -> 18.6μs (2.57% faster)


def test_missing_core_dim_drop():
    # Test missing core dim with 'drop'
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = {"x": Variable(("other_dim",), 10), "y": Variable(("dim",), 20)}

    def add_func(v1, v2):
        return Variable(v1.dims, v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, on_missing_core_dim="drop"
    )
    result = codeflash_output  # 21.5μs -> 20.9μs (2.65% faster)


def test_empty_dicts():
    # Test with empty dicts
    a = {}
    b = {}

    def add_func(v1, v2):
        return Variable(("dim",), v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature
    )
    result = codeflash_output  # 10.2μs -> 9.96μs (2.56% faster)


def test_left_join():
    # Test left join
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = {"x": Variable(("dim",), 10)}

    def add_func(v1, v2):
        v2_data = v2.data if isinstance(v2, Variable) else 0
        return Variable(v1.dims, v1.data + v2_data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func,
        a,
        b,
        signature=signature,
        join="left",
        fill_value=Variable(("dim",), 0),
    )
    result = codeflash_output  # 15.4μs -> 14.5μs (5.93% faster)


def test_right_join():
    # Test right join
    a = {"x": Variable(("dim",), 1)}
    b = {"x": Variable(("dim",), 10), "y": Variable(("dim",), 20)}

    def add_func(v1, v2):
        v1_data = v1.data if isinstance(v1, Variable) else 0
        v2_data = v2.data if isinstance(v2, Variable) else 0
        return Variable(("dim",), v1_data + v2_data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func,
        a,
        b,
        signature=signature,
        join="right",
        fill_value=Variable(("dim",), 0),
    )
    result = codeflash_output  # 15.5μs -> 13.7μs (13.5% faster)


def test_multiple_outputs():
    # Test with multiple outputs
    a = {"x": Variable(("dim",), 1)}
    b = {"x": Variable(("dim",), 10)}

    def func(v1, v2):
        return (
            Variable(v1.dims, v1.data + v2.data),
            Variable(v1.dims, v1.data - v2.data),
        )

    signature = _UFuncSignature([["dim"], ["dim"]], num_outputs=2)
    out1, out2 = apply_dict_of_variables_vfunc(
        func, a, b, signature=signature
    )  # 18.4μs -> 17.1μs (7.56% faster)


def test_multiple_outputs_with_drop():
    # Multiple outputs with drop
    a = {"x": Variable(("dim",), 1), "y": Variable(("dim",), 2)}
    b = {"x": Variable(("other_dim",), 10), "y": Variable(("dim",), 20)}

    def func(v1, v2):
        return (
            Variable(v1.dims, v1.data + v2.data),
            Variable(v1.dims, v1.data - v2.data),
        )

    signature = _UFuncSignature([["dim"], ["dim"]], num_outputs=2)
    out1, out2 = apply_dict_of_variables_vfunc(
        func, a, b, signature=signature, on_missing_core_dim="drop"
    )  # 25.1μs -> 24.5μs (2.41% faster)


# Large Scale Test Cases


def test_large_dicts_outer_join():
    # Test with large dicts and outer join
    N = 500
    a = {f"k{i}": Variable(("dim",), i) for i in range(N)}
    b = {f"k{i}": Variable(("dim",), 2 * i) for i in range(N // 2, N + N // 2)}

    def add_func(v1, v2):
        v1_data = v1.data if isinstance(v1, Variable) else 0
        v2_data = v2.data if isinstance(v2, Variable) else 0
        return Variable(("dim",), v1_data + v2_data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func,
        a,
        b,
        signature=signature,
        join="outer",
        fill_value=Variable(("dim",), 0),
    )
    result = codeflash_output  # 1.02ms -> 895μs (13.4% faster)
    # All keys from both dicts should be present
    all_keys = set(a.keys()).union(b.keys())
    # Spot check: key only in b
    k_last = f"k{N+N//2-1}"


def test_large_dicts_inner_join():
    # Test with large dicts and inner join
    N = 500
    a = {f"k{i}": Variable(("dim",), i) for i in range(N)}
    b = {f"k{i}": Variable(("dim",), 2 * i) for i in range(N)}

    def add_func(v1, v2):
        return Variable(("dim",), v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, join="inner"
    )
    result = codeflash_output  # 669μs -> 589μs (13.6% faster)


def test_large_multiple_outputs():
    # Test with large dicts and multiple outputs
    N = 300
    a = {f"k{i}": Variable(("dim",), i) for i in range(N)}
    b = {f"k{i}": Variable(("dim",), 2 * i) for i in range(N)}

    def func(v1, v2):
        return (
            Variable(v1.dims, v1.data + v2.data),
            Variable(v1.dims, v1.data - v2.data),
        )

    signature = _UFuncSignature([["dim"], ["dim"]], num_outputs=2)
    out1, out2 = apply_dict_of_variables_vfunc(
        func, a, b, signature=signature
    )  # 534μs -> 473μs (13.0% faster)


def test_large_missing_core_dim_drop():
    # Large dicts, drop missing core dims
    N = 200
    a = {f"k{i}": Variable(("dim",), i) for i in range(N)}
    b = {f"k{i}": Variable(("other_dim",), 2 * i) for i in range(N)}

    def add_func(v1, v2):
        return Variable(("dim",), v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, on_missing_core_dim="drop"
    )
    result = codeflash_output  # 438μs -> 411μs (6.43% faster)


def test_large_missing_core_dim_copy():
    # Large dicts, copy missing core dims
    N = 200
    a = {f"k{i}": Variable(("dim",), i) for i in range(N)}
    b = {f"k{i}": Variable(("other_dim",), 2 * i) for i in range(N)}

    def add_func(v1, v2):
        return Variable(("dim",), v1.data + v2.data)

    signature = _UFuncSignature([["dim"], ["dim"]])
    codeflash_output = apply_dict_of_variables_vfunc(
        add_func, a, b, signature=signature, on_missing_core_dim="copy"
    )
    result = codeflash_output  # 451μs -> 421μs (7.13% faster)
    for k in a:
        pass

To edit these changes git checkout codeflash/optimize-apply_dict_of_variables_vfunc-miyqb9hn and push.