Preserve multiindex (#216)

* test based on example

* fix

* Update tests/test_xarray.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Deepak Cherian <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

flox/xarray.py

@@ -432,7 +432,7 @@ def wrapper(array, *by, func, skipna, core_dims, **kwargs):
 
     # restore non-dim coord variables without the core dimension
     # TODO: shouldn't apply_ufunc handle this?
-    for var in set(ds_broad.variables) - set(ds_broad.dims):
+    for var in set(ds_broad.variables) - set(ds_broad.xindexes) - set(ds_broad.dims):
         if all(d not in ds_broad[var].dims for d in dim_tuple):
             actual[var] = ds_broad[var]
 

tests/test_xarray.py

@@ -594,3 +594,34 @@ def test_dtype_accumulation(use_flox, chunk):
         assert np.issubdtype(actual.dtype, np.float64)
         assert np.issubdtype(actual.compute().dtype, np.float64)
         xr.testing.assert_allclose(expected, actual, **tolerance64)
+
+
+def test_preserve_multiindex():
+    """Regression test for GH issue #215"""
+
+    vort = xr.DataArray(
+        name="vort",
+        data=np.random.uniform(size=(4, 2)),
+        dims=["i", "face"],
+        coords={"i": ("i", np.arange(4)), "face": ("face", np.arange(2))},
+    )
+
+    vort = (
+        vort.coarsen(i=2)
+        .construct(i=("i_region_coarse", "i_region"))
+        .stack(region=["face", "i_region_coarse"])
+    )
+
+    bins = [np.linspace(0, 1, 10)]
+    bin_intervals = tuple(pd.IntervalIndex.from_breaks(b) for b in bins)
+
+    hist = xarray_reduce(
+        xr.DataArray(1),  # weights
+        vort,  # variables we want to bin
+        func="count",  # count occurrences falling in bins
+        expected_groups=bin_intervals,  # bins for each variable
+        dim=["i_region"],  # broadcast dimensions
+        fill_value=0,  # fill empty bins with 0 counts
+    )
+
+    assert "region" in hist.coords