From b78f853d9ae9c8e20c305033065dbd12e44db7a5 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 4 Nov 2024 16:51:39 -0700 Subject: [PATCH 1/2] Fix groupby tests Closes #9715 --- xarray/tests/test_groupby.py | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 95ba0a3384e..c547e83ed91 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -2848,20 +2848,25 @@ def test_multiple_groupers(use_flox) -> None: xy=UniqueGrouper(labels=["a", "b", "c"]), eagerly_compute_group=eagerly_compute_group, ) + expected = xr.DataArray( + [[[1, 1, 1], [np.nan, 1, 2]]] * 4, + dims=("z", "x", "xy"), + coords={"xy": ("xy", ["a", "b", "c"], {"foo": "bar"})}, + ) with raise_if_dask_computes(max_computes=1): if eagerly_compute_group: with pytest.warns(DeprecationWarning): gb = b.groupby(**kwargs) # type: ignore[arg-type] + assert_identical(gb.count(), expected) else: gb = b.groupby(**kwargs) # type: ignore[arg-type] assert is_chunked_array(gb.encoded.codes.data) assert not gb.encoded.group_indices - expected = xr.DataArray( - [[[1, 1, 1], [np.nan, 1, 2]]] * 4, - dims=("z", "x", "xy"), - coords={"xy": ("xy", ["a", "b", "c"], {"foo": "bar"})}, - ) - assert_identical(gb.count(), expected) + if has_flox: + assert_identical(gb.count(), expected) + else: + with pytest.raises(ValueError, match="when lazily grouping"): + gb.count() @pytest.mark.parametrize("use_flox", [True, False]) @@ -3003,11 +3008,6 @@ def test_lazy_grouping(grouper, expect_index): pd.testing.assert_index_equal(encoded.full_index, expect_index) np.testing.assert_array_equal(encoded.unique_coord.values, np.array(expect_index)) - lazy = ( - xr.Dataset({"foo": data}, coords={"zoo": data}) - .groupby(zoo=grouper, eagerly_compute_group=False) - .count() - ) eager = ( xr.Dataset({"foo": data}, coords={"zoo": data.compute()}) .groupby(zoo=grouper) @@ -3017,9 +3017,16 @@ def test_lazy_grouping(grouper, expect_index): {"foo": (encoded.codes.name, np.ones(encoded.full_index.size))}, coords={encoded.codes.name: expect_index}, ) - assert_identical(eager, lazy) assert_identical(eager, expected) + if has_flox: + lazy = ( + xr.Dataset({"foo": data}, coords={"zoo": data}) + .groupby(zoo=grouper, eagerly_compute_group=False) + .count() + ) + assert_identical(eager, lazy) + @requires_dask def test_lazy_grouping_errors(): From cb33ee03f37d00c104fbebfadec70de162e070b2 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 4 Nov 2024 20:37:53 -0700 Subject: [PATCH 2/2] fix --- xarray/tests/test_groupby.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index c547e83ed91..c0eeace71af 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -2853,20 +2853,22 @@ def test_multiple_groupers(use_flox) -> None: dims=("z", "x", "xy"), coords={"xy": ("xy", ["a", "b", "c"], {"foo": "bar"})}, ) - with raise_if_dask_computes(max_computes=1): - if eagerly_compute_group: + if eagerly_compute_group: + with raise_if_dask_computes(max_computes=1): with pytest.warns(DeprecationWarning): gb = b.groupby(**kwargs) # type: ignore[arg-type] assert_identical(gb.count(), expected) - else: + else: + with raise_if_dask_computes(max_computes=0): gb = b.groupby(**kwargs) # type: ignore[arg-type] - assert is_chunked_array(gb.encoded.codes.data) - assert not gb.encoded.group_indices - if has_flox: + assert is_chunked_array(gb.encoded.codes.data) + assert not gb.encoded.group_indices + if has_flox: + with raise_if_dask_computes(max_computes=1): assert_identical(gb.count(), expected) - else: - with pytest.raises(ValueError, match="when lazily grouping"): - gb.count() + else: + with pytest.raises(ValueError, match="when lazily grouping"): + gb.count() @pytest.mark.parametrize("use_flox", [True, False])