From c9e2a62cc43308f397126e549afc8992049984d3 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 2 Nov 2022 15:29:05 -0600 Subject: [PATCH] WIP --- xarray/core/groupby.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 54fe767b717..d4833d72c65 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -596,8 +596,19 @@ def __repr__(self) -> str: def _iter_grouped(self) -> Iterator[T_Xarray]: """Iterate over each element in this group""" - for indices in self._group_indices: - yield self._obj.isel({self._group_dim: indices}) + print(self._group_indices) + if isinstance(self._group_indices[0], slice): + for indices in self._group_indices: + yield self._obj.isel({self._group_dim: indices}) + else: + for code in np.sort(pd.unique(self._codes.data.reshape(-1))): + if code == -1: + continue + subset = self._original_obj.where(self._codes == code, drop=True) + if self._squeeze: + squeeze_dims = tuple(dim for dim in self._codes.dims if subset.sizes[dim] == 1) + subset = subset.squeeze(squeeze_dims) + yield subset def _infer_concat_args(self, applied_example): if self._group_dim in applied_example.dims: