flox causes groupby on multiindex to remove the multiindex indexes

[max-sixty]

What happened?

When applying .sum() on an xarray.DataArray grouped by a multiindex, the DataArray loses the individual components of the multiindex. If flox is disabled, or we run .first, it retains the individual indexes.

What did you expect to happen?

Running through flox wouldn't lose the individual indexes.

Minimal Complete Verifiable Example

import xarray as xr
import pandas as pd

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})

# Stack into a MultiIndex
da = da.stack(multiindex=["time", "location"])

# Groupby the multiindex
grouped = da.groupby('multiindex')

# Apply sum and first
sum_result = grouped.sum()
first_result = grouped.first()

print(sum_result)
print(first_result)

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

<xarray.DataArray (multiindex: 8)>
array([0.5488135 , 0.71518937, 0.60276338, 0.54488318, 0.4236548 ,
       0.64589411, 0.43758721, 0.891773  ])
Coordinates:
  * multiindex  (multiindex) object MultiIndex

# Missing the individual index items

<xarray.DataArray (multiindex: 8)>
array([0.5488135 , 0.71518937, 0.60276338, 0.54488318, 0.4236548 ,
       0.64589411, 0.43758721, 0.891773  ])
Coordinates:
  * multiindex  (multiindex) object MultiIndex
  * time        (multiindex) datetime64[ns] 2023-01-01 ... 2023-01-04
  * location    (multiindex) <U1 'A' 'B' 'A' 'B' 'A' 'B' 'A' 'B'

Anything else we need to know?

No response

Environment

INSTALLED VERSIONS

commit: b9780e7
python: 3.11.10 (main, Sep 7 2024, 01:03:31) [Clang 15.0.0 (clang-1500.3.9.4)]
python-bits: 64
OS: Darwin
OS-release: 23.6.0
machine: arm64
processor: arm
byteorder: little
LC_ALL: en_US.UTF-8
LANG: None
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.9.1.dev32+gece582dd
pandas: 2.2.2
numpy: 2.0.2
scipy: 1.14.1
netCDF4: 1.7.1.post2
pydap: None
h5netcdf: 1.3.0
h5py: 3.11.0
zarr: 2.18.3
cftime: 1.6.4
nc_time_axis: 1.4.1
iris: None
bottleneck: 1.4.0
dask: 2024.8.2
distributed: 2024.8.2
matplotlib: 3.9.2
cartopy: None
seaborn: 0.13.2
numbagg: 0.8.1
fsspec: 2024.9.0
cupy: None
pint: None
sparse: None
flox: 0.9.12
numpy_groupies: 0.11.2
setuptools: 69.2.0
pip: 24.0
conda: None
pytest: 8.3.3
mypy: 1.11.2
IPython: 8.24.0
sphinx: None