Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Nov 19, 2024
1 parent c77d7c5 commit 6951e6a
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 7 deletions.
8 changes: 5 additions & 3 deletions xarray/core/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -740,9 +740,11 @@ def _shuffle_obj(self, chunks: T_Chunks) -> T_Xarray:
shuffled = as_dataset._shuffle(
dim=self._group_dim, indices=self.encoded.group_indices, chunks=chunks
)
shuffled = self._maybe_unstack(shuffled)
new_obj = self._obj._from_temp_dataset(shuffled) if was_array else shuffled
return new_obj
unstacked: Dataset = self._maybe_unstack(shuffled)
if was_array:
return self._obj._from_temp_dataset(unstacked)
else:
return unstacked # type: ignore[return-value]

def map(
self,
Expand Down
9 changes: 5 additions & 4 deletions xarray/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -1891,10 +1891,10 @@ def resample_as_pandas(array, *args, **kwargs):

rs = array.resample(time="24h", closed="right")
actual = rs.mean()
shuffled = rs.distributed_shuffle().resample(time="24h", closed="right").mean()
shuffled = rs.distributed_shuffle().resample(time="24h", closed="right")
expected = resample_as_pandas(array, "24h", closed="right")
assert_identical(expected, actual)
assert_identical(expected, shuffled)
assert_identical(expected, shuffled.mean())

with pytest.raises(ValueError, match=r"Index must be monotonic"):
array[[2, 0, 1]].resample(time=resample_freq)
Expand Down Expand Up @@ -2883,9 +2883,10 @@ def test_multiple_groupers(use_flox: bool, shuffle: bool) -> None:
coords={"xy": (("x", "y"), [["a", "b", "c"], ["b", "c", "c"]], {"foo": "bar"})},
dims=["x", "y", "z"],
)
gb = b.groupby(x=UniqueGrouper(), y=UniqueGrouper())
groupers = dict(x=UniqueGrouper(), y=UniqueGrouper())
gb = b.groupby(groupers)
if shuffle:
gb = gb.distributed_shuffle()
gb = gb.distributed_shuffle().groupby(groupers)
repr(gb)
with xr.set_options(use_flox=use_flox):
assert_identical(gb.mean("z"), b.mean("z"))
Expand Down

0 comments on commit 6951e6a

Please sign in to comment.