Cleanup

Author: dcherian

xarray/core/duck_array_ops.py

@@ -835,13 +835,9 @@ def chunked_nanlast(darray, axis):
 
 def shuffle_array(array, indices: list[list[int]], axis: int):
     # TODO: do chunk manager dance here.
-    if is_duck_dask_array(array):
-        if not module_available("dask", minversion="2024.08.0"):
-            raise ValueError(
-                "This method is very inefficient on dask<2024.08.0. Please upgrade."
-            )
-        # TODO: handle dimensions
-        return array.shuffle(indexer=indices, axis=axis)
+    if is_chunked_array(array):
+        chunkmanager = get_chunked_array_type(array)
+        return chunkmanager.shuffle(array, indexer=indices, axis=axis)
     else:
         indexer = np.concatenate(indices)
         # TODO: Do the array API thing here.

xarray/core/groupby.py

@@ -538,6 +538,8 @@ def shuffle(self) -> None:
         if all(isinstance(idx, slice) for idx in self._group_indices):
             return
 
+        indices: tuple[list[int]] = self._group_indices  # type: ignore[assignment]
+
         was_array = isinstance(self._obj, DataArray)
         as_dataset = self._obj._to_temp_dataset() if was_array else self._obj
 
@@ -547,20 +549,24 @@ def shuffle(self) -> None:
                 shuffled[name] = var
                 continue
             shuffled_data = shuffle_array(
-                var._data, list(self._group_indices), axis=var.get_axis_num(dim)
+                var._data, list(indices), axis=var.get_axis_num(dim)
             )
             shuffled[name] = var._replace(data=shuffled_data)
 
         # Replace self._group_indices with slices
         slices = []
         start = 0
         for idxr in self._group_indices:
+            if TYPE_CHECKING:
+                assert not isinstance(idxr, slice)
             slices.append(slice(start, start + len(idxr)))
             start += len(idxr)
         # TODO: we have now broken the invariant
         # self._group_indices ≠ self.groupers[0].group_indices
         self._group_indices = tuple(slices)
         if was_array:
+            if TYPE_CHECKING:
+                assert isinstance(self._obj, DataArray)
             self._obj = self._obj._from_temp_dataset(shuffled)
         else:
             self._obj = shuffled

xarray/core/types.py

@@ -297,7 +297,7 @@ def copy(
 ZarrWriteModes = Literal["w", "w-", "a", "a-", "r+", "r"]
 
 GroupKey = Any
-GroupIndex = Union[int, slice, list[int]]
+GroupIndex = Union[slice, list[int]]
 GroupIndices = tuple[GroupIndex, ...]
 Bins = Union[
     int, Sequence[int], Sequence[float], Sequence[pd.Timestamp], np.ndarray, pd.Index

xarray/namedarray/daskmanager.py

@@ -251,3 +251,12 @@ def store(
             targets=targets,
             **kwargs,
         )
+
+    def shuffle(self, x: DaskArray, indexer: list[list[int]], axis: int) -> DaskArray:
+        import dask.array
+
+        if not module_available("dask", minversion="2024.08.0"):
+            raise ValueError(
+                "This method is very inefficient on dask<2024.08.0. Please upgrade."
+            )
+        return dask.array.shuffle(x, indexer, axis)

xarray/namedarray/parallelcompat.py

@@ -364,6 +364,11 @@ def compute(
         """
         raise NotImplementedError()
 
+    def shuffle(
+        self, x: T_ChunkedArray, indexer: list[list[int]], axis: int
+    ) -> T_ChunkedArray:
+        raise NotImplementedError()
+
     @property
     def array_api(self) -> Any:
         """