Fix reordering of dataarray dimensions inside dataset (#289)

* Fix groupby behaviour for dataarray dimensions inside dataset

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

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

---------

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

Author: eendebakpt

flox/xarray.py

@@ -28,10 +28,12 @@
     Dims = Union[str, Iterable[Hashable], None]
 
 
-def _restore_dim_order(result, obj, by):
+def _restore_dim_order(result, obj, by, no_groupby_reorder=False):
     def lookup_order(dimension):
         if dimension == by.name and by.ndim == 1:
             (dimension,) = by.dims
+            if no_groupby_reorder:
+                return -1e6  # some arbitrarily low value
         if dimension in obj.dims:
             axis = obj.get_axis_num(dimension)
         else:
@@ -491,7 +493,12 @@ def wrapper(array, *by, func, skipna, core_dims, **kwargs):
                 template = obj
 
             if actual[var].ndim > 1:
-                actual[var] = _restore_dim_order(actual[var], template, by_da[0])
+                no_groupby_reorder = isinstance(
+                    obj, xr.Dataset
+                )  # do not re-order dataarrays inside datasets
+                actual[var] = _restore_dim_order(
+                    actual[var], template, by_da[0], no_groupby_reorder=no_groupby_reorder
+                )
 
     if missing_dim:
         for k, v in missing_dim.items():

tests/test_xarray.py

@@ -598,3 +598,34 @@ def test_fill_value_xarray_binning():
         actual = data_array.groupby_bins("y", bins=4).mean()
 
     xr.testing.assert_identical(expected, actual)
+
+
+def test_groupby_2d_dataset():
+    d = {
+        "coords": {
+            "bit_index": {"dims": ("bit_index",), "attrs": {"name": "bit_index"}, "data": [0, 1]},
+            "index": {"dims": ("index",), "data": [0, 6, 8, 10, 14]},
+            "clifford": {"dims": ("index",), "attrs": {}, "data": [1, 1, 4, 10, 4]},
+        },
+        "dims": {"bit_index": 2, "index": 5},
+        "data_vars": {
+            "counts": {
+                "dims": ("bit_index", "index"),
+                "attrs": {
+                    "name": "counts",
+                },
+                "data": [[18, 30, 45, 70, 38], [382, 370, 355, 330, 362]],
+            }
+        },
+    }
+
+    ds = xr.Dataset.from_dict(d)
+
+    with xr.set_options(use_flox=False):
+        expected = ds.groupby("clifford").mean()
+    with xr.set_options(use_flox=True):
+        actual = ds.groupby("clifford").mean()
+    assert (
+        expected.counts.dims == actual.counts.dims
+    )  # https://github.com/pydata/xarray/issues/8292
+    xr.testing.assert_identical(expected, actual)