From 321c5608a3be3cd4b6a4de3b658d1e2d164c0409 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Mon, 21 Mar 2022 16:26:20 +0100 Subject: [PATCH] fix DataArray groupby returning a Dataset (#6394) --- xarray/core/groupby.py | 2 +- xarray/tests/test_groupby.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 9212549d4ad..df78b7789f7 100644 --- 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 index 327c8d34b18..b1f14d6be2d 100644 --- 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)