flox: Properly propagate multiindex

Closes pydata#9648

Author: dcherian

xarray/core/coordinates.py

@@ -1116,3 +1116,14 @@ def create_coords_with_default_indexes(
         new_coords = Coordinates._construct_direct(coords=variables, indexes=indexes)
 
     return new_coords
+
+
+def _coordinates_from_variable(variable: Variable) -> Coordinates:
+    from xarray.core.indexes import create_default_index_implicit
+
+    (name,) = variable.dims
+    new_index, index_vars = create_default_index_implicit(variable)
+    indexes = {k: new_index for k in index_vars}
+    new_vars = new_index.create_variables()
+    new_vars[name].attrs = variable.attrs
+    return Coordinates(new_vars, indexes)

xarray/core/groupby.py

@@ -21,13 +21,13 @@
 from xarray.core.arithmetic import DataArrayGroupbyArithmetic, DatasetGroupbyArithmetic
 from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
 from xarray.core.concat import concat
-from xarray.core.coordinates import Coordinates
+from xarray.core.coordinates import Coordinates, _coordinates_from_variable
 from xarray.core.formatting import format_array_flat
 from xarray.core.indexes import (
-    PandasIndex,
     PandasMultiIndex,
     filter_indexes_from_coords,
 )
+from xarray.core.merge import merge_coords
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import (
     Dims,
@@ -851,7 +851,6 @@ def _flox_reduce(
         from flox.xarray import xarray_reduce
 
         from xarray.core.dataset import Dataset
-        from xarray.groupers import BinGrouper
 
         obj = self._original_obj
         variables = (
@@ -901,13 +900,6 @@ def _flox_reduce(
                 # set explicitly to avoid unnecessarily accumulating count
                 kwargs["min_count"] = 0
 
-        unindexed_dims: tuple[Hashable, ...] = tuple(
-            grouper.name
-            for grouper in self.groupers
-            if isinstance(grouper.group, _DummyGroup)
-            and not isinstance(grouper.grouper, BinGrouper)
-        )
-
         parsed_dim: tuple[Hashable, ...]
         if isinstance(dim, str):
             parsed_dim = (dim,)
@@ -963,26 +955,28 @@ def _flox_reduce(
         # we did end up reducing over dimension(s) that are
         # in the grouped variable
         group_dims = set(grouper.group.dims)
-        new_coords = {}
+        new_coords = []
         if group_dims.issubset(set(parsed_dim)):
-            new_indexes = {}
             for grouper in self.groupers:
                 output_index = grouper.full_index
                 if isinstance(output_index, pd.RangeIndex):
+                    # flox always assigns an index so we must drop it here if we don't need it.
+                    result = result.drop_vars(grouper.name)
                     continue
-                name = grouper.name
-                new_coords[name] = IndexVariable(
-                    dims=name, data=np.array(output_index), attrs=grouper.codes.attrs
+                new_coords.append(
+                    # Using IndexVariable here ensures we reconstruct PandasMultiIndex with
+                    # all associated levels properly.
+                    _coordinates_from_variable(
+                        IndexVariable(
+                            dims=grouper.name,
+                            data=output_index,
+                            attrs=grouper.codes.attrs,
+                        )
+                    )
                 )
-                index_cls = (
-                    PandasIndex
-                    if not isinstance(output_index, pd.MultiIndex)
-                    else PandasMultiIndex
-                )
-                new_indexes[name] = index_cls(output_index, dim=name)
             result = result.assign_coords(
-                Coordinates(new_coords, new_indexes)
-            ).drop_vars(unindexed_dims)
+                Coordinates._construct_direct(*merge_coords(new_coords))
+            )
 
         # broadcast any non-dim coord variables that don't
         # share all dimensions with the grouper

xarray/groupers.py

@@ -16,7 +16,7 @@
 
 from xarray.coding.cftime_offsets import BaseCFTimeOffset, _new_to_legacy_freq
 from xarray.core import duck_array_ops
-from xarray.core.coordinates import Coordinates
+from xarray.core.coordinates import Coordinates, _coordinates_from_variable
 from xarray.core.dataarray import DataArray
 from xarray.core.groupby import T_Group, _DummyGroup
 from xarray.core.indexes import safe_cast_to_index
@@ -42,17 +42,6 @@
 RESAMPLE_DIM = "__resample_dim__"
 
 
-def _coordinates_from_variable(variable: Variable) -> Coordinates:
-    from xarray.core.indexes import create_default_index_implicit
-
-    (name,) = variable.dims
-    new_index, index_vars = create_default_index_implicit(variable)
-    indexes = {k: new_index for k in index_vars}
-    new_vars = new_index.create_variables()
-    new_vars[name].attrs = variable.attrs
-    return Coordinates(new_vars, indexes)
-
-
 @dataclass(init=False)
 class EncodedGroups:
     """

xarray/tests/test_groupby.py

@@ -145,6 +145,25 @@ def test_multi_index_groupby_sum() -> None:
     assert_equal(expected, actual)
 
 
+def test_multi_index_propagation():
+    # regression test for GH9648
+    times = pd.date_range("2023-01-01", periods=4)
+    locations = ["A", "B"]
+    data = [[0.5, 0.7], [0.6, 0.5], [0.4, 0.6], [0.4, 0.9]]
+
+    da = xr.DataArray(
+        data, dims=["time", "location"], coords={"time": times, "location": locations}
+    )
+    da = da.stack(multiindex=["time", "location"])
+    grouped = da.groupby("multiindex")
+
+    with xr.set_options(use_flox=True):
+        actual = grouped.sum()
+    with xr.set_options(use_flox=False):
+        expected = grouped.first()
+    assert_identical(actual, expected)
+
+
 def test_groupby_da_datetime() -> None:
     # test groupby with a DataArray of dtype datetime for GH1132
     # create test data