fix DataArray groupby returning a Dataset (#6394)

pydata · Mar 21, 2022 · 321c560 · 321c560
1 parent c604ee1
commit 321c560
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
@@ -866,7 +866,7 @@ def _combine(self, applied, shortcut=False):
         if coord is not None and dim not in applied_example.dims:
             index, index_vars = create_default_index_implicit(coord)
             indexes = {k: index for k in index_vars}
-            combined = combined._overwrite_indexes(indexes, coords=index_vars)
+            combined = combined._overwrite_indexes(indexes, index_vars)
         combined = self._maybe_restore_empty_groups(combined)
         combined = self._maybe_unstack(combined)
         return combined

diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
@@ -936,9 +936,17 @@ def test_groupby_dataset_assign():
 
 def test_groupby_dataset_map_dataarray_func():
     # regression GH6379
-    ds = xr.Dataset({"foo": ("x", [1, 2, 3, 4])}, coords={"x": [0, 0, 1, 1]})
+    ds = Dataset({"foo": ("x", [1, 2, 3, 4])}, coords={"x": [0, 0, 1, 1]})
     actual = ds.groupby("x").map(lambda grp: grp.foo.mean())
-    expected = xr.DataArray([1.5, 3.5], coords={"x": [0, 1]}, dims="x", name="foo")
+    expected = DataArray([1.5, 3.5], coords={"x": [0, 1]}, dims="x", name="foo")
+    assert_identical(actual, expected)
+
+
+def test_groupby_dataarray_map_dataset_func():
+    # regression GH6379
+    da = DataArray([1, 2, 3, 4], coords={"x": [0, 0, 1, 1]}, dims="x", name="foo")
+    actual = da.groupby("x").map(lambda grp: grp.mean().to_dataset())
+    expected = xr.Dataset({"foo": ("x", [1.5, 3.5])}, coords={"x": [0, 1]})
     assert_identical(actual, expected)