Function test_apply_dask_parallelized_one_arg

xarray/tests/test_computation.py:1213–1229

()

@requires_dask
def test_apply_dask_parallelized_one_arg() -> None:
    import dask.array as da

    array = da.ones((2, 2), chunks=(1, 1))
    data_array = xr.DataArray(array, dims=("x", "y"))

    def parallel_identity(x):
        return apply_ufunc(identity, x, dask="parallelized", output_dtypes=[x.dtype])

    actual = parallel_identity(data_array)
    assert isinstance(actual.data, da.Array)
    assert actual.data.chunks == array.chunks
    assert_identical(data_array, actual)

    computed = data_array.compute()
    actual = parallel_identity(computed)
    assert_identical(computed, actual)


@requires_dask