GroupBy(multiple groupers)

pydata · Aug 21, 2024 · da7e67b · da7e67b
1 parent 56124ae
commit da7e67b
Show file tree

Hide file tree

Showing 4 changed files with 332 additions and 88 deletions.
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -6801,27 +6801,22 @@ def groupby(
             groupers = either_dict_or_kwargs(group, groupers, "groupby")  # type: ignore
             group = None
 
-        grouper: Grouper
+        rgroupers: tuple[ResolvedGrouper, ...]
         if group is not None:
             if groupers:
                 raise ValueError(
                     "Providing a combination of `group` and **groupers is not supported."
                 )
-            grouper = UniqueGrouper()
+            rgroupers = (ResolvedGrouper(UniqueGrouper(), group, self),)
         else:
-            if len(groupers) > 1:
-                raise ValueError("grouping by multiple variables is not supported yet.")
             if not groupers:
                 raise ValueError("Either `group` or `**groupers` must be provided.")
-            group, grouper = next(iter(groupers.items()))
-
-        rgrouper = ResolvedGrouper(grouper, group, self)
+            rgroupers = tuple(
+                ResolvedGrouper(grouper, group, self)
+                for group, grouper in groupers.items()
+            )
 
-        return DataArrayGroupBy(
-            self,
-            (rgrouper,),
-            restore_coord_dims=restore_coord_dims,
-        )
+        return DataArrayGroupBy(self, rgroupers, restore_coord_dims=restore_coord_dims)
 
     @_deprecate_positional_args("v2024.07.0")
     def groupby_bins(

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -10383,25 +10383,22 @@ def groupby(
             groupers = either_dict_or_kwargs(group, groupers, "groupby")  # type: ignore
             group = None
 
+        rgroupers: tuple[ResolvedGrouper, ...]
         if group is not None:
             if groupers:
                 raise ValueError(
                     "Providing a combination of `group` and **groupers is not supported."
                 )
-            rgrouper = ResolvedGrouper(UniqueGrouper(), group, self)
+            rgroupers = (ResolvedGrouper(UniqueGrouper(), group, self),)
         else:
-            if len(groupers) > 1:
-                raise ValueError("Grouping by multiple variables is not supported yet.")
-            elif not groupers:
+            if not groupers:
                 raise ValueError("Either `group` or `**groupers` must be provided.")
-            for group, grouper in groupers.items():
-                rgrouper = ResolvedGrouper(grouper, group, self)
+            rgroupers = tuple(
+                ResolvedGrouper(grouper, group, self)
+                for group, grouper in groupers.items()
+            )
 
-        return DatasetGroupBy(
-            self,
-            (rgrouper,),
-            restore_coord_dims=restore_coord_dims,
-        )
+        return DatasetGroupBy(self, rgroupers, restore_coord_dims=restore_coord_dims)
 
     @_deprecate_positional_args("v2024.07.0")
     def groupby_bins(