Introduce Grouper objects.

pydata · Mar 9, 2023 · 1147411 · 1147411
1 parent f30da34
commit 1147411
Show file tree

Hide file tree

Showing 7 changed files with 378 additions and 288 deletions.
diff --git a/xarray/core/common.py b/xarray/core/common.py
@@ -996,11 +996,16 @@ def _resample(
         if base is not None and offset is not None:
             raise ValueError("base and offset cannot be present at the same time")
 
+        index = self._indexes[dim_name].to_pandas_index()
         if base is not None:
-            index = self._indexes[dim_name].to_pandas_index()
             offset = _convert_base_to_offset(base, freq, index)
 
+        group = DataArray(
+            dim_coord, coords=dim_coord.coords, dims=dim_coord.dims, name=RESAMPLE_DIM
+        )
+
         grouper = TimeResampleGrouper(
+            group=group,
             freq=freq,
             closed=closed,
             label=label,
@@ -1009,14 +1014,10 @@ def _resample(
             loffset=loffset,
         )
 
-        group = DataArray(
-            dim_coord, coords=dim_coord.coords, dims=dim_coord.dims, name=RESAMPLE_DIM
-        )
         return resample_cls(
             self,
-            group=group,
-            dim=dim_name,
             grouper=grouper,
+            dim=dim_name,
             resample_dim=RESAMPLE_DIM,
             restore_coord_dims=restore_coord_dims,
         )

diff --git a/xarray/core/computation.py b/xarray/core/computation.py
@@ -531,8 +531,9 @@ def apply_groupby_func(func, *args):
             "grouped the same way"
         )
 
-    grouped_dim = first_groupby._group.name
-    unique_values = first_groupby._unique_coord.values
+    (grouper,) = first_groupby.groupers
+    grouped_dim = grouper.name
+    unique_values = grouper.unique_coord.values
 
     iterators = []
     for arg in args:

diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -6256,7 +6256,7 @@ def groupby(
         core.groupby.DataArrayGroupBy
         pandas.DataFrame.groupby
         """
-        from xarray.core.groupby import DataArrayGroupBy
+        from xarray.core.groupby import DataArrayGroupBy, UniqueGrouper
 
         # While we don't generally check the type of every arg, passing
         # multiple dimensions as multiple arguments is common enough, and the
@@ -6269,8 +6269,9 @@ def groupby(
                 f"`squeeze` must be True or False, but {squeeze} was supplied"
             )
 
+        grouper = UniqueGrouper(group)
         return DataArrayGroupBy(
-            self, group, squeeze=squeeze, restore_coord_dims=restore_coord_dims
+            self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
         )
 
     def groupby_bins(
@@ -6341,14 +6342,22 @@ def groupby_bins(
         ----------
         .. [1] http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
         """
-        from xarray.core.groupby import DataArrayGroupBy
+        from xarray.core.groupby import BinGrouper, DataArrayGroupBy
 
-        return DataArrayGroupBy(
-            self,
-            group,
-            squeeze=squeeze,
+        # While we don't generally check the type of every arg, passing
+        # multiple dimensions as multiple arguments is common enough, and the
+        # consequences hidden enough (strings evaluate as true) to warrant
+        # checking here.
+        # A future version could make squeeze kwarg only, but would face
+        # backward-compat issues.
+        if not isinstance(squeeze, bool):
+            raise TypeError(
+                f"`squeeze` must be True or False, but {squeeze} was supplied"
+            )
+
+        grouper = BinGrouper(
+            group=group,
             bins=bins,
-            restore_coord_dims=restore_coord_dims,
             cut_kwargs={
                 "right": right,
                 "labels": labels,
@@ -6357,6 +6366,13 @@ def groupby_bins(
             },
         )
 
+        return DataArrayGroupBy(
+            self,
+            grouper,
+            squeeze=squeeze,
+            restore_coord_dims=restore_coord_dims,
+        )
+
     def weighted(self, weights: DataArray) -> DataArrayWeighted:
         """
         Weighted DataArray operations.

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -8942,7 +8942,7 @@ def groupby(
         Dataset.resample
         DataArray.resample
         """
-        from xarray.core.groupby import DatasetGroupBy
+        from xarray.core.groupby import DatasetGroupBy, UniqueGrouper
 
         # While we don't generally check the type of every arg, passing
         # multiple dimensions as multiple arguments is common enough, and the
@@ -8955,8 +8955,10 @@ def groupby(
                 f"`squeeze` must be True or False, but {squeeze} was supplied"
             )
 
+        grouper = UniqueGrouper(group)
+
         return DatasetGroupBy(
-            self, group, squeeze=squeeze, restore_coord_dims=restore_coord_dims
+            self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
         )
 
     def groupby_bins(
@@ -9027,14 +9029,11 @@ def groupby_bins(
         ----------
         .. [1] http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
         """
-        from xarray.core.groupby import DatasetGroupBy
+        from xarray.core.groupby import BinGrouper, DatasetGroupBy
 
-        return DatasetGroupBy(
-            self,
-            group,
-            squeeze=squeeze,
+        grouper = BinGrouper(
+            group=group,
             bins=bins,
-            restore_coord_dims=restore_coord_dims,
             cut_kwargs={
                 "right": right,
                 "labels": labels,
@@ -9043,6 +9042,10 @@ def groupby_bins(
             },
         )
 
+        return DatasetGroupBy(
+            self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
+        )
+
     def weighted(self, weights: DataArray) -> DatasetWeighted:
         """
         Weighted Dataset operations.