⚡️ Speed up function apply_array_ufunc by 18%

[codeflash-ai]

📄 18% (0.18x) speedup for apply_array_ufunc in xarray/core/computation.py

⏱️ Runtime : 1.70 milliseconds → 1.44 milliseconds (best of 5 runs)

📝 Explanation and details

The optimization restructures the is_chunked_array function to improve performance through strategic reordering of checks and early short-circuiting.

Key optimizations:

1. Early Dask detection: The function now checks is_duck_dask_array(x) first since Dask arrays are a specific subset that always returns True. This creates an immediate return path for the most common chunked array type.

2. Attribute check before duck array validation: Instead of calling is_duck_array(x) and hasattr(x, "chunks"), the optimized version checks hasattr(x, "chunks") first as a lightweight filter. If an object lacks the chunks attribute, it can immediately return False without the expensive duck array validation.

3. Reduced redundant checks: The original logic performed is_duck_array(x) twice for objects with chunks (once in is_duck_dask_array and once in the second condition). The optimized version eliminates this redundancy.

Performance impact: The 17% speedup is most pronounced in scenarios with large numbers of arguments, as seen in the test results where test_apply_array_ufunc_large_list improved by 53.8% and test_apply_array_ufunc_large_duck_arrays by 19.5%. The optimization is particularly effective when apply_array_ufunc processes many non-chunked arrays, as the faster is_chunked_array check reduces overhead in the any() loop.

Hot path relevance: Given that apply_array_ufunc is called from the high-level apply_ufunc function (a core xarray API used extensively for array operations), this optimization benefits any xarray computation involving multiple arrays, making it especially valuable for data science workflows processing large datasets.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from xarray.core.computation import apply_array_ufunc


# Mocks for duckarray and dask array detection
class DummyDuckArray:
    """A minimal duck array that mimics a numpy array."""

    def __init__(self, shape=(1,), dtype=float):
        self.ndim = len(shape)
        self.shape = shape
        self.dtype = dtype
        self.__array_function__ = True
        self.__array_ufunc__ = True


class DummyDuckDaskArray(DummyDuckArray):
    """A minimal duck dask array that mimics a dask array."""

    def __init__(self, shape=(1,), dtype=float, chunks=(1,)):
        super().__init__(shape, dtype)
        self.chunks = chunks
        self._is_dask_collection = True  # marker for dask


from xarray.core.computation import apply_array_ufunc

# unit tests

# --- BASIC TEST CASES ---


def test_apply_array_ufunc_basic_scalar():
    # Test with simple scalar arguments
    def add(x, y):
        return x + y

    codeflash_output = apply_array_ufunc(add, 2, 3)  # 5.03μs -> 4.91μs (2.46% faster)


def test_apply_array_ufunc_basic_duck_array():
    # Test with duck arrays
    def add(x, y):
        return x.shape[0] + y.shape[0]

    arr1 = DummyDuckArray(shape=(2,))
    arr2 = DummyDuckArray(shape=(3,))
    codeflash_output = apply_array_ufunc(
        add, arr1, arr2
    )  # 12.5μs -> 12.3μs (2.06% faster)


def test_apply_array_ufunc_basic_mixed_types():
    # Test with mixed types (int and duck array)
    def multiply(x, y):
        return x * y.shape[0]

    arr = DummyDuckArray(shape=(4,))
    codeflash_output = apply_array_ufunc(
        multiply, 2, arr
    )  # 7.37μs -> 7.60μs (3.00% slower)


def test_apply_array_ufunc_basic_kwargs():
    # Test with kwargs passed to function
    def add(x, y, z=0):
        return x + y + z

    codeflash_output = apply_array_ufunc(
        lambda x, y: add(x, y, z=5), 1, 2
    )  # 3.95μs -> 3.76μs (5.16% faster)


# --- EDGE TEST CASES ---


def test_apply_array_ufunc_empty_args():
    # Test with no arguments (should raise TypeError)
    def nothing():
        return "done"

    codeflash_output = apply_array_ufunc(nothing)  # 1.50μs -> 1.36μs (10.6% faster)


def test_apply_array_ufunc_single_arg():
    # Test with a single argument
    def square(x):
        return x * x

    codeflash_output = apply_array_ufunc(square, 4)  # 3.50μs -> 3.27μs (7.14% faster)


def test_apply_array_ufunc_duck_dask_array_forbidden():
    # Test with a duck dask array and dask='forbidden' (should raise ValueError)
    arr = DummyDuckDaskArray(shape=(5,), chunks=(2, 3))

    def identity(x):
        return x

    with pytest.raises(ValueError, match="apply_ufunc encountered a dask array"):
        apply_array_ufunc(
            identity, arr, dask="forbidden"
        )  # 8.01μs -> 8.42μs (4.84% slower)


def test_apply_array_ufunc_duck_dask_array_allowed():
    # Test with a duck dask array and dask='allowed' (should not raise)
    arr = DummyDuckDaskArray(shape=(5,), chunks=(2, 3))

    def identity(x):
        return x

    codeflash_output = apply_array_ufunc(
        identity, arr, dask="allowed"
    )  # 7.35μs -> 7.65μs (3.89% slower)


def test_apply_array_ufunc_duck_dask_array_unknown_dask():
    # Test with a duck dask array and unknown dask argument (should raise ValueError)
    arr = DummyDuckDaskArray(shape=(5,), chunks=(2, 3))

    def identity(x):
        return x

    with pytest.raises(ValueError, match="unknown setting for dask array handling"):
        apply_array_ufunc(
            identity, arr, dask="unsupported"
        )  # 7.72μs -> 7.81μs (1.22% slower)


def test_apply_array_ufunc_duck_dask_array_parallelized():
    # Test with a duck dask array and dask='parallelized' (should raise ValueError)
    arr = DummyDuckDaskArray(shape=(5,), chunks=(2, 3))

    def identity(x):
        return x

    with pytest.raises(ValueError, match="cannot use dask='parallelized'"):
        apply_array_ufunc(
            identity, arr, dask="parallelized"
        )  # 7.99μs -> 7.70μs (3.71% faster)


def test_apply_array_ufunc_duck_array_with_chunks():
    # Test with a duck array that has chunks attribute but is not dask
    arr = DummyDuckArray(shape=(5,))
    arr.chunks = (5,)

    def identity(x):
        return x

    # Should raise ValueError since chunks attribute triggers is_chunked_array
    with pytest.raises(ValueError, match="apply_ufunc encountered a dask array"):
        apply_array_ufunc(
            identity, arr, dask="forbidden"
        )  # 7.49μs -> 8.07μs (7.25% slower)


def test_apply_array_ufunc_multiple_args_one_chunked():
    # Test with multiple arguments, one chunked array
    arr1 = DummyDuckArray(shape=(2,))
    arr2 = DummyDuckDaskArray(shape=(3,), chunks=(3,))

    def add(x, y):
        return x.shape[0] + y.shape[0]

    with pytest.raises(ValueError, match="apply_ufunc encountered a dask array"):
        apply_array_ufunc(
            add, arr1, arr2, dask="forbidden"
        )  # 9.39μs -> 9.42μs (0.329% slower)


def test_apply_array_ufunc_multiple_args_all_chunked_allowed():
    # Test with multiple chunked arrays, dask='allowed'
    arr1 = DummyDuckDaskArray(shape=(2,), chunks=(2,))
    arr2 = DummyDuckDaskArray(shape=(3,), chunks=(3,))

    def add(x, y):
        return x.shape[0] + y.shape[0]

    codeflash_output = apply_array_ufunc(
        add, arr1, arr2, dask="allowed"
    )  # 7.71μs -> 7.65μs (0.876% faster)


def test_apply_array_ufunc_non_duck_array_with_chunks():
    # Test with a non-duck array with a chunks attribute
    class WeirdArray:
        def __init__(self):
            self.chunks = (1,)

    arr = WeirdArray()

    def identity(x):
        return x

    # Should not raise, since it's not a duck array
    codeflash_output = apply_array_ufunc(
        identity, arr
    )  # 3.48μs -> 3.56μs (2.19% slower)


def test_apply_array_ufunc_func_raises():
    # Test when the function itself raises an exception
    def explode(x):
        raise RuntimeError("boom")

    with pytest.raises(RuntimeError, match="boom"):
        apply_array_ufunc(explode, 1)  # 3.73μs -> 3.82μs (2.38% slower)


# --- LARGE SCALE TEST CASES ---


def test_apply_array_ufunc_large_list():
    # Test with a large list of integers
    def sum_all(*args):
        return sum(args)

    args = list(range(1000))
    codeflash_output = apply_array_ufunc(
        sum_all, *args
    )  # 267μs -> 173μs (53.8% faster)


def test_apply_array_ufunc_large_duck_arrays():
    # Test with many duck arrays
    def total_size(*arrays):
        return sum(arr.shape[0] for arr in arrays)

    arrays = [DummyDuckArray(shape=(i + 1,)) for i in range(1000)]
    codeflash_output = apply_array_ufunc(
        total_size, *arrays
    )  # 1.02ms -> 855μs (19.5% faster)


def test_apply_array_ufunc_large_duck_dask_arrays_allowed():
    # Test with many duck dask arrays, dask='allowed'
    def total_size(*arrays):
        return sum(arr.shape[0] for arr in arrays)

    arrays = [DummyDuckDaskArray(shape=(i + 1,), chunks=(i + 1,)) for i in range(1000)]
    codeflash_output = apply_array_ufunc(
        total_size, *arrays, dask="allowed"
    )  # 45.3μs -> 46.1μs (1.71% slower)


def test_apply_array_ufunc_large_duck_dask_arrays_forbidden():
    # Test with many duck dask arrays, dask='forbidden' (should raise)
    def total_size(*arrays):
        return sum(arr.shape[0] for arr in arrays)

    arrays = [DummyDuckDaskArray(shape=(i + 1,), chunks=(i + 1,)) for i in range(1000)]
    with pytest.raises(ValueError, match="apply_ufunc encountered a dask array"):
        apply_array_ufunc(
            total_size, *arrays, dask="forbidden"
        )  # 12.0μs -> 11.9μs (0.986% faster)


def test_apply_array_ufunc_large_mixed_arrays():
    # Test with a mix of duck arrays and duck dask arrays, dask='allowed'
    def total_size(*arrays):
        return sum(arr.shape[0] for arr in arrays)

    arrays = [
        (
            DummyDuckArray(shape=(i + 1,))
            if i % 2 == 0
            else DummyDuckDaskArray(shape=(i + 1,), chunks=(i + 1,))
        )
        for i in range(1000)
    ]
    codeflash_output = apply_array_ufunc(
        total_size, *arrays, dask="allowed"
    )  # 50.4μs -> 50.6μs (0.259% slower)


def test_apply_array_ufunc_large_mixed_arrays_forbidden():
    # Test with a mix of duck arrays and duck dask arrays, dask='forbidden' (should raise)
    def total_size(*arrays):
        return sum(arr.shape[0] for arr in arrays)

    arrays = [
        (
            DummyDuckArray(shape=(i + 1,))
            if i % 2 == 0
            else DummyDuckDaskArray(shape=(i + 1,), chunks=(i + 1,))
        )
        for i in range(1000)
    ]
    with pytest.raises(ValueError, match="apply_ufunc encountered a dask array"):
        apply_array_ufunc(
            total_size, *arrays, dask="forbidden"
        )  # 12.7μs -> 12.9μs (1.87% slower)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from xarray.core.computation import apply_array_ufunc

# --- Mocks and minimal stubs for duckarray and dask detection ---
# These are minimal stubs to allow the code to run without external dependencies.
# They are *not* mocks of the function under test, but allow us to simulate duck arrays and dask arrays.


class DummyDuckArray:
    def __init__(self, shape=(2, 2), dtype="float64"):
        self.ndim = len(shape)
        self.shape = shape
        self.dtype = dtype
        self.__array_function__ = True
        self.__array_ufunc__ = True


class DummyDaskArray(DummyDuckArray):
    def __init__(self, shape=(2, 2), dtype="float64"):
        super().__init__(shape, dtype)
        self.chunks = ((1, 1), (1, 1))  # Simulate chunked array


from xarray.core.computation import apply_array_ufunc

# --- Unit tests ---

# Basic Test Cases


def test_apply_array_ufunc_basic_add():
    # Test with two regular Python lists, should behave like normal function call
    codeflash_output = apply_array_ufunc(
        lambda x, y: [a + b for a, b in zip(x, y)], [1, 2], [3, 4]
    )
    result = codeflash_output  # 4.91μs -> 5.19μs (5.56% slower)


def test_apply_array_ufunc_basic_tuple_args():
    # Test with tuples
    codeflash_output = apply_array_ufunc(
        lambda x, y: tuple(a * b for a, b in zip(x, y)), (2, 3), (4, 5)
    )
    result = codeflash_output  # 5.04μs -> 5.05μs (0.139% slower)


def test_apply_array_ufunc_basic_scalar_args():
    # Test with scalar arguments
    codeflash_output = apply_array_ufunc(lambda x, y: x**y, 2, 3)
    result = codeflash_output  # 3.95μs -> 3.89μs (1.49% faster)


def test_apply_array_ufunc_basic_duckarray_args():
    # Test with duck array arguments
    arr1 = DummyDuckArray(shape=(2,), dtype="int")
    arr2 = DummyDuckArray(shape=(2,), dtype="int")
    # The function just returns the shape tuple for simplicity
    codeflash_output = apply_array_ufunc(lambda x, y: (x.shape, y.shape), arr1, arr2)
    result = codeflash_output  # 11.1μs -> 10.7μs (3.44% faster)


def test_apply_array_ufunc_basic_no_args():
    # Test with no args, just a function
    codeflash_output = apply_array_ufunc(lambda: "no args")
    result = codeflash_output  # 1.63μs -> 1.39μs (16.9% faster)


# Edge Test Cases


def test_apply_array_ufunc_empty_list():
    # Test with empty list
    codeflash_output = apply_array_ufunc(lambda x: len(x), [])
    result = codeflash_output  # 3.48μs -> 3.17μs (9.82% faster)


def test_apply_array_ufunc_empty_tuple():
    # Test with empty tuple
    codeflash_output = apply_array_ufunc(lambda x: len(x), ())
    result = codeflash_output  # 3.41μs -> 3.36μs (1.49% faster)


def test_apply_array_ufunc_none_arg():
    # Test with None as argument
    codeflash_output = apply_array_ufunc(lambda x: x is None, None)
    result = codeflash_output  # 3.19μs -> 3.31μs (3.83% slower)


def test_apply_array_ufunc_single_arg():
    # Test with single argument
    codeflash_output = apply_array_ufunc(lambda x: x + 1, 10)
    result = codeflash_output  # 3.20μs -> 3.10μs (3.23% faster)


def test_apply_array_ufunc_mixed_types():
    # Test with mixed types in args
    codeflash_output = apply_array_ufunc(lambda x, y: str(x) + str(y), 1, "a")
    result = codeflash_output  # 4.27μs -> 4.17μs (2.42% faster)


def test_apply_array_ufunc_chunked_array_forbidden():
    # Test with chunked (dask) array, dask forbidden
    arr = DummyDaskArray()
    with pytest.raises(ValueError) as excinfo:
        apply_array_ufunc(
            lambda x: x, arr, dask="forbidden"
        )  # 7.73μs -> 8.13μs (4.92% slower)


def test_apply_array_ufunc_chunked_array_parallelized():
    # Test with chunked (dask) array, dask parallelized
    arr = DummyDaskArray()
    with pytest.raises(ValueError) as excinfo:
        apply_array_ufunc(
            lambda x: x, arr, dask="parallelized"
        )  # 7.47μs -> 7.70μs (2.92% slower)


def test_apply_array_ufunc_chunked_array_allowed():
    # Test with chunked (dask) array, dask allowed
    arr = DummyDaskArray()
    # Should not raise error
    codeflash_output = apply_array_ufunc(lambda x: "allowed", arr, dask="allowed")
    result = codeflash_output  # 7.34μs -> 7.55μs (2.85% slower)


def test_apply_array_ufunc_chunked_array_unknown_dask_setting():
    # Test with chunked (dask) array, unknown dask setting
    arr = DummyDaskArray()
    with pytest.raises(ValueError) as excinfo:
        apply_array_ufunc(
            lambda x: x, arr, dask="notarealsetting"
        )  # 7.76μs -> 7.97μs (2.61% slower)


def test_apply_array_ufunc_non_chunked_array_with_chunks_attr():
    # Test with a duck array that has a 'chunks' attribute but is not dask
    class NotDaskDuckArray(DummyDuckArray):
        def __init__(self):
            super().__init__()
            self.chunks = ((1, 1),)

    arr = NotDaskDuckArray()
    # Should be treated as chunked, so forbidden should raise
    with pytest.raises(ValueError):
        apply_array_ufunc(
            lambda x: x, arr, dask="forbidden"
        )  # 7.63μs -> 8.08μs (5.54% slower)


def test_apply_array_ufunc_multiple_args_one_chunked():
    # Test with multiple arguments, one is chunked
    arr1 = DummyDuckArray()
    arr2 = DummyDaskArray()
    with pytest.raises(ValueError):
        apply_array_ufunc(
            lambda x, y: (x, y), arr1, arr2, dask="forbidden"
        )  # 8.91μs -> 9.12μs (2.31% slower)


def test_apply_array_ufunc_multiple_args_all_chunked_allowed():
    # Test with multiple chunked arrays, dask allowed
    arr1 = DummyDaskArray()
    arr2 = DummyDaskArray()
    codeflash_output = apply_array_ufunc(lambda x, y: "ok", arr1, arr2, dask="allowed")
    result = codeflash_output  # 7.24μs -> 7.16μs (1.05% faster)


def test_apply_array_ufunc_multiple_args_mixed_chunked():
    # Test with mixed chunked and non-chunked arrays, dask allowed
    arr1 = DummyDuckArray()
    arr2 = DummyDaskArray()
    codeflash_output = apply_array_ufunc(
        lambda x, y: "mixed", arr1, arr2, dask="allowed"
    )
    result = codeflash_output  # 8.49μs -> 8.76μs (3.07% slower)


def test_apply_array_ufunc_args_with_false_chunks_attr():
    # Test with object that has a 'chunks' attribute but is not duck array
    class FakeChunks:
        chunks = ((1, 1),)

    obj = FakeChunks()
    # Should not be treated as chunked array
    codeflash_output = apply_array_ufunc(lambda x: "not chunked", obj)
    result = codeflash_output  # 3.37μs -> 3.49μs (3.46% slower)


# Large Scale Test Cases


def test_apply_array_ufunc_large_list_sum():
    # Test with large list, sum operation
    large_list = list(range(1000))
    codeflash_output = apply_array_ufunc(lambda x: sum(x), large_list)
    result = codeflash_output  # 5.69μs -> 5.76μs (1.27% slower)


def test_apply_array_ufunc_large_tuple_product():
    # Test with large tuple, product operation
    large_tuple = tuple([2] * 500)

    # Compute product without importing math.prod (Python 3.8+), use loop
    def product(seq):
        result = 1
        for x in seq:
            result *= x
        return result

    codeflash_output = apply_array_ufunc(product, large_tuple)
    result = codeflash_output  # 18.3μs -> 17.9μs (2.37% faster)


def test_apply_array_ufunc_large_duckarray_args():
    # Test with large duck arrays
    arr1 = DummyDuckArray(shape=(1000,), dtype="float64")
    arr2 = DummyDuckArray(shape=(1000,), dtype="float64")
    # Just check shapes are preserved
    codeflash_output = apply_array_ufunc(
        lambda x, y: (x.shape[0] + y.shape[0]), arr1, arr2
    )
    result = codeflash_output  # 8.45μs -> 8.27μs (2.25% faster)


def test_apply_array_ufunc_large_chunked_array_allowed():
    # Test with large chunked (dask) array, dask allowed
    arr = DummyDaskArray(shape=(1000,), dtype="float64")
    codeflash_output = apply_array_ufunc(lambda x: x.shape[0], arr, dask="allowed")
    result = codeflash_output  # 7.17μs -> 7.80μs (8.09% slower)


def test_apply_array_ufunc_large_mixed_args():
    # Test with large mixed arguments
    arr1 = DummyDuckArray(shape=(1000,), dtype="float64")
    arr2 = DummyDaskArray(shape=(1000,), dtype="float64")
    codeflash_output = apply_array_ufunc(
        lambda x, y: x.shape[0] * y.shape[0], arr1, arr2, dask="allowed"
    )
    result = codeflash_output  # 9.38μs -> 9.46μs (0.825% slower)


def test_apply_array_ufunc_large_list_and_tuple():
    # Test with large list and tuple, zipped addition
    large_list = list(range(1000))
    large_tuple = tuple(range(1000, 2000))
    codeflash_output = apply_array_ufunc(
        lambda x, y: [a + b for a, b in zip(x, y)], large_list, large_tuple
    )
    result = codeflash_output  # 31.6μs -> 31.9μs (0.712% slower)


# 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-apply_array_ufunc-miyse1e0 and push.