Skip to content

Commit

Permalink
Fix groupby tests
Browse files Browse the repository at this point in the history
Closes #9715
  • Loading branch information
dcherian committed Nov 4, 2024
1 parent af7c6d4 commit b78f853
Showing 1 changed file with 19 additions and 12 deletions.
31 changes: 19 additions & 12 deletions xarray/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / macos-latest py3.10

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / macos-latest py3.10

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.10 min-all-deps

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.10 min-all-deps

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / macos-latest py3.12

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / macos-latest py3.12

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.12 all-but-numba

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.12 all-but-numba

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.10

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.10

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.12

test_multiple_groupers[True] RuntimeError: Too many computes. Total: 2 > max: 1.

Check failure on line 2866 in xarray/tests/test_groupby.py

View workflow job for this annotation

GitHub Actions / ubuntu-latest py3.12

test_multiple_groupers[False] RuntimeError: Too many computes. Total: 2 > max: 1.
else:
with pytest.raises(ValueError, match="when lazily grouping"):
gb.count()


@pytest.mark.parametrize("use_flox", [True, False])
Expand Down Expand Up @@ -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)
Expand All @@ -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():
Expand Down

0 comments on commit b78f853

Please sign in to comment.