Function test_eval_dask_complex_expression

xarray/tests/test_eval.py:491–510

Test that complex expressions preserve dask backing.

()

@requires_dask
def test_eval_dask_complex_expression() -> None:
    """Test that complex expressions preserve dask backing."""
    from xarray.core.utils import is_duck_dask_array

    rng = np.random.default_rng(42)
    ds = Dataset(
        {
            "x": (["time", "lat", "lon"], rng.random((3, 4, 5))),
            "y": (["time", "lat", "lon"], rng.random((3, 4, 5))),
        }
    ).chunk({"time": 1, "lat": 2, "lon": 5})

    with raise_if_dask_computes():
        result = ds.eval("x * 2 + y ** 2")

    assert is_duck_dask_array(result.data)

    # Verify correctness when computed
    expected = ds["x"] * 2 + ds["y"] ** 2
    assert_equal(result, expected)


@requires_dask