Skip to content

Commit 0384363

Browse files
authored
Fix groupby tests (#9716)
Closes #9715
1 parent af7c6d4 commit 0384363

File tree

1 file changed

+24
-15
lines changed

1 file changed

+24
-15
lines changed

xarray/tests/test_groupby.py

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2848,20 +2848,27 @@ def test_multiple_groupers(use_flox) -> None:
28482848
xy=UniqueGrouper(labels=["a", "b", "c"]),
28492849
eagerly_compute_group=eagerly_compute_group,
28502850
)
2851-
with raise_if_dask_computes(max_computes=1):
2852-
if eagerly_compute_group:
2853-
with pytest.warns(DeprecationWarning):
2854-
gb = b.groupby(**kwargs) # type: ignore[arg-type]
2855-
else:
2856-
gb = b.groupby(**kwargs) # type: ignore[arg-type]
2857-
assert is_chunked_array(gb.encoded.codes.data)
2858-
assert not gb.encoded.group_indices
28592851
expected = xr.DataArray(
28602852
[[[1, 1, 1], [np.nan, 1, 2]]] * 4,
28612853
dims=("z", "x", "xy"),
28622854
coords={"xy": ("xy", ["a", "b", "c"], {"foo": "bar"})},
28632855
)
2864-
assert_identical(gb.count(), expected)
2856+
if eagerly_compute_group:
2857+
with raise_if_dask_computes(max_computes=1):
2858+
with pytest.warns(DeprecationWarning):
2859+
gb = b.groupby(**kwargs) # type: ignore[arg-type]
2860+
assert_identical(gb.count(), expected)
2861+
else:
2862+
with raise_if_dask_computes(max_computes=0):
2863+
gb = b.groupby(**kwargs) # type: ignore[arg-type]
2864+
assert is_chunked_array(gb.encoded.codes.data)
2865+
assert not gb.encoded.group_indices
2866+
if has_flox:
2867+
with raise_if_dask_computes(max_computes=1):
2868+
assert_identical(gb.count(), expected)
2869+
else:
2870+
with pytest.raises(ValueError, match="when lazily grouping"):
2871+
gb.count()
28652872

28662873

28672874
@pytest.mark.parametrize("use_flox", [True, False])
@@ -3003,11 +3010,6 @@ def test_lazy_grouping(grouper, expect_index):
30033010
pd.testing.assert_index_equal(encoded.full_index, expect_index)
30043011
np.testing.assert_array_equal(encoded.unique_coord.values, np.array(expect_index))
30053012

3006-
lazy = (
3007-
xr.Dataset({"foo": data}, coords={"zoo": data})
3008-
.groupby(zoo=grouper, eagerly_compute_group=False)
3009-
.count()
3010-
)
30113013
eager = (
30123014
xr.Dataset({"foo": data}, coords={"zoo": data.compute()})
30133015
.groupby(zoo=grouper)
@@ -3017,9 +3019,16 @@ def test_lazy_grouping(grouper, expect_index):
30173019
{"foo": (encoded.codes.name, np.ones(encoded.full_index.size))},
30183020
coords={encoded.codes.name: expect_index},
30193021
)
3020-
assert_identical(eager, lazy)
30213022
assert_identical(eager, expected)
30223023

3024+
if has_flox:
3025+
lazy = (
3026+
xr.Dataset({"foo": data}, coords={"zoo": data})
3027+
.groupby(zoo=grouper, eagerly_compute_group=False)
3028+
.count()
3029+
)
3030+
assert_identical(eager, lazy)
3031+
30233032

30243033
@requires_dask
30253034
def test_lazy_grouping_errors():

0 commit comments

Comments
 (0)