fix

Author: dcherian

xarray/core/groupby.py

@@ -704,14 +704,14 @@ def distributed_shuffle(self, chunks: T_Chunks = None) -> T_Xarray:
 
         Examples
         --------
-        >>> import dask
+        >>> import dask.array
         >>> da = xr.DataArray(
         ...     dims="x",
         ...     data=dask.array.arange(10, chunks=3),
         ...     coords={"x": [1, 2, 3, 1, 2, 3, 1, 2, 3, 0]},
         ...     name="a",
         ... )
-        >>> shuffled = da.groupby("x")
+        >>> shuffled = da.groupby("x").distributed_shuffle()
         >>> shuffled.groupby("x").quantile(q=0.5).compute()
         <xarray.DataArray 'a' (x: 4)> Size: 32B
         array([9., 3., 4., 5.])
@@ -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,

xarray/tests/test_groupby.py

@@ -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)
@@ -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"))