File tree 1 file changed +13
-2
lines changed
1 file changed +13
-2
lines changed Original file line number Diff line number Diff line change @@ -596,8 +596,19 @@ def __repr__(self) -> str:
596
596
597
597
def _iter_grouped (self ) -> Iterator [T_Xarray ]:
598
598
"""Iterate over each element in this group"""
599
- for indices in self ._group_indices :
600
- yield self ._obj .isel ({self ._group_dim : indices })
599
+ print (self ._group_indices )
600
+ if isinstance (self ._group_indices [0 ], slice ):
601
+ for indices in self ._group_indices :
602
+ yield self ._obj .isel ({self ._group_dim : indices })
603
+ else :
604
+ for code in np .sort (pd .unique (self ._codes .data .reshape (- 1 ))):
605
+ if code == - 1 :
606
+ continue
607
+ subset = self ._original_obj .where (self ._codes == code , drop = True )
608
+ if self ._squeeze :
609
+ squeeze_dims = tuple (dim for dim in self ._codes .dims if subset .sizes [dim ] == 1 )
610
+ subset = subset .squeeze (squeeze_dims )
611
+ yield subset
601
612
602
613
def _infer_concat_args (self , applied_example ):
603
614
if self ._group_dim in applied_example .dims :
You can’t perform that action at this time.
0 commit comments