Fix recombination in groupby when changing size along the grouped dimension

Author: ej81

doc/whats-new.rst

@@ -35,6 +35,10 @@ New Features
 Bug fixes
 ~~~~~~~~~
 
+- Fix recombination of groups in :py:meth:`Dataset.groupby` and
+  :py:meth:`DataArray.groupby` when performing an operation that changes the
+  size of the groups along the grouped dimension. By `Eric Jansen
+  <https://github.com/ej81>`_.
 - Fix :py:meth:`Dataset.swap_dims` and :py:meth:`DataArray.swap_dims` producing
   index with name reflecting the previous dimension name instead of the new one
   (:issue:`3748`, :pull:`3752`). By `Joseph K Aicher

xarray/core/groupby.py

@@ -715,7 +715,7 @@ def assign_coords(self, coords=None, **coords_kwargs):
 def _maybe_reorder(xarray_obj, dim, positions):
     order = _inverse_permutation_indices(positions)
 
-    if order is None:
+    if order is None or len(order) != xarray_obj.sizes[dim]:
         return xarray_obj
     else:
         return xarray_obj[{dim: order}]
@@ -833,7 +833,7 @@ def _combine(self, applied, restore_coord_dims=False, shortcut=False):
         if isinstance(combined, type(self._obj)):
             # only restore dimension order for arrays
             combined = self._restore_dim_order(combined)
-        if coord is not None:
+        if coord is not None and applied_example.sizes.get(dim, 0) == 0:
             if shortcut:
                 coord_var = as_variable(coord)
                 combined._coords[coord.name] = coord_var
@@ -949,7 +949,7 @@ def _combine(self, applied):
         coord, dim, positions = self._infer_concat_args(applied_example)
         combined = concat(applied, dim)
         combined = _maybe_reorder(combined, dim, positions)
-        if coord is not None:
+        if coord is not None and applied_example.sizes.get(dim, 0) == 0:
             combined[coord.name] = coord
         combined = self._maybe_restore_empty_groups(combined)
         combined = self._maybe_unstack(combined)

xarray/tests/test_groupby.py

@@ -107,6 +107,50 @@ def test_groupby_input_mutation():
     assert_identical(array, array_copy)  # should not modify inputs
 
 
+def test_da_groupby_map_shrink_groups():
+    array = xr.DataArray([1, 2, 3, 4, 5, 6], [("x", [1, 1, 1, 2, 2, 2])])
+    expected = array.isel(x=[0, 1, 3, 4])
+    actual = array.groupby("x").map(lambda f: f.isel(x=[0, 1]))
+    assert_identical(expected, actual)
+
+
+def test_ds_groupby_map_shrink_groups():
+    dataset = xr.Dataset({"foo": ("x", [1, 2, 3, 4, 5, 6])}, {"x": [1, 1, 1, 2, 2, 2]})
+    expected = dataset.isel(x=[0, 1, 3, 4])
+    actual = dataset.groupby("x").map(lambda f: f.isel(x=[0, 1]))
+    assert_identical(expected, actual)
+
+
+def test_da_groupby_map_change_group_size():
+    array = xr.DataArray([1, 2, 3], [("x", [1, 2, 2])])
+
+    def func(group):
+        if group.sizes["x"] == 1:
+            result = group.isel(x=[0, 0])
+        else:
+            result = group.isel(x=[0])
+        return result
+
+    expected = array.isel(x=[0, 0, 1])
+    actual = array.groupby("x").map(func)
+    assert_identical(expected, actual)
+
+
+def test_ds_groupby_map_change_group_size():
+    dataset = xr.Dataset({"foo": ("x", [1, 2, 3])}, {"x": [1, 2, 2]})
+
+    def func(group):
+        if group.sizes["x"] == 1:
+            result = group.isel(x=[0, 0])
+        else:
+            result = group.isel(x=[0])
+        return result
+
+    expected = dataset.isel(x=[0, 0, 1])
+    actual = dataset.groupby("x").map(func)
+    assert_identical(expected, actual)
+
+
 def test_da_groupby_map_func_args():
     def func(arg1, arg2, arg3=0):
         return arg1 + arg2 + arg3