Function test_vectorized_indexing_dask_array

xarray/tests/test_indexing.py:1136–1163

()

@requires_dask
def test_vectorized_indexing_dask_array() -> None:
    # https://github.com/pydata/xarray/issues/2511#issuecomment-563330352
    darr = DataArray(data=[0.2, 0.4, 0.6], coords={"z": range(3)}, dims=("z",))
    indexer = DataArray(
        data=np.random.randint(0, 3, 8).reshape(4, 2).astype(int),
        coords={"y": range(4), "x": range(2)},
        dims=("y", "x"),
    )
    expected = darr[indexer]

    # fails because we can't index pd.Index lazily (yet).
    # We could make this succeed by auto-chunking the values
    # and constructing a lazy index variable, and not automatically
    # create an index for it.
    with pytest.raises(ValueError, match="Cannot index with"):
        with raise_if_dask_computes():
            darr.chunk()[indexer.chunk({"y": 2})]
    with pytest.raises(ValueError, match="Cannot index with"):
        with raise_if_dask_computes():
            actual = darr[indexer.chunk({"y": 2})]

    with raise_if_dask_computes():
        actual = darr.drop_vars("z").chunk()[indexer.chunk({"y": 2})]
    assert_identical(actual, expected.drop_vars("z"))

    with raise_if_dask_computes():
        actual_variable = darr.variable.chunk()[indexer.variable.chunk({"y": 2})]
    assert_identical(actual_variable, expected.variable)


@requires_dask