Function test_shuffle_simple

xarray/tests/test_groupby.py:3363–3376

()

@requires_dask_ge_2024_08_1
def test_shuffle_simple() -> None:
    import dask

    da = xr.DataArray(
        dims="x",
        data=dask.array.from_array([1, 2, 3, 4, 5, 6], chunks=2),
        coords={"label": ("x", ["a", "b", "c", "a", "b", "c"])},
    )
    actual = da.groupby(label=UniqueGrouper()).shuffle_to_chunks()
    expected = da.isel(x=[0, 3, 1, 4, 2, 5])
    assert_identical(actual, expected)

    with pytest.raises(ValueError):
        da.chunk(x=2, eagerly_load_group=False).groupby("label").shuffle_to_chunks()


@requires_dask_ge_2024_08_1