Keep MultiIndex columns in GroupBy.agg for MultiIndex-column sources

python/cudf/cudf/core/groupby/groupby.py

@@ -1314,7 +1314,7 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
             # RangeIndex) columns.
             data = ColumnAccessor(
                 data,
-                multiindex=False,
+                multiindex=self.obj._data.multiindex,
                 level_names=self.obj._data.level_names,
                 rangeindex=self.obj._data.rangeindex,
                 label_dtype=self.obj._data.label_dtype,
@@ -1325,12 +1325,30 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
             and self.obj.ndim == 2
             and self.obj._data.level_names != (None,)
         ):
-            data = ColumnAccessor(
-                data,
-                multiindex=False,
-                level_names=self.obj._data.level_names,
-                label_dtype=self.obj._data.label_dtype,
-            )
+            if self.obj._data.multiindex and all(
+                isinstance(label, tuple)
+                and len(label) == self.obj._data.nlevels
+                for label in data
+            ):
+                # the aggregation kept the source's tuple labels: preserve
+                # the MultiIndex columns and their per-level metadata
+                data = ColumnAccessor(
+                    data,
+                    multiindex=True,
+                    level_names=self.obj._data.level_names,
+                    label_dtype=self.obj._data.label_dtype,
+                    level_dtypes=self.obj._data.level_dtypes,
+                )
+            else:
+                # relabeling aggregations (``agg(new=(col, func))``) emit
+                # new flat labels: the source's multi-level metadata does
+                # not describe them
+                data = ColumnAccessor(
+                    data,
+                    multiindex=False,
+                    level_names=self.obj._data.level_names,
+                    label_dtype=self.obj._data.label_dtype,
+                )
         else:
             data = ColumnAccessor(data, multiindex=multilevel)
         if not multilevel and len(data) > 0:

python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py

@@ -1766,7 +1766,6 @@ def pytest_unconfigure(config):
     "tests/groupby/aggregate/test_aggregate.py::test_agg_str_with_kwarg_axis_1_raises[nunique]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_directory[size]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_empty_key_empty_return": "TODO: Add a reason for failure",
-    "tests/groupby/aggregate/test_aggregate.py::test_multiindex_custom_func[<lambda>0]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_aggregate.py::test_order_aggregate_multiple_funcs": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-False]": "TODO: Add a reason for failure",
     "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-True]": "TODO: Add a reason for failure",
@@ -1886,7 +1885,6 @@ def pytest_unconfigure(config):
     "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[ns]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[s]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[us]": "TODO: Add a reason for failure",
-    "tests/groupby/test_groupby.py::test_groupby_with_hier_columns": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_groups_repr_truncates[1-{0: [0], ...}]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_groups_repr_truncates[4-{0: [0], 1: [1], 2: [2], 3: [3], ...}]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_groups_repr_truncates[5-{0: [0], 1: [1], 2: [2], 3: [3], 4: [4]}]": "TODO: Add a reason for failure",
@@ -1895,7 +1893,6 @@ def pytest_unconfigure(config):
     "tests/groupby/test_groupby.py::test_ops_not_as_index[idxmin]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_ops_not_as_index[size]": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby.py::test_single_element_listlike_level_grouping[level_arg0-False]": "AssertionError: assert ['x', 'y'] == [('x',), ('y',)]",
-    "tests/groupby/test_groupby.py::test_wrap_aggregated_output_multindex": "TODO: Add a reason for failure",
     "tests/groupby/test_groupby_dropna.py::test_groupby_nan_included": "GroupBy.indices returns cupy arrays nested in a dict that cudf.pandas does not wrap, so assert_numpy_array_equal sees mismatched array classes",
     "tests/groupby/test_groupby_subclass.py::test_groupby_preserves_metadata": "proxy pd.Index passed to the raw SubclassedSeries constructor loses its name: pandas' maybe_extract_name checks isinstance against the concrete Index class, which proxies cannot satisfy",
     "tests/groupby/test_grouping.py::TestGetGroup::test_get_group_grouped_by_tuple": "TODO: Add a reason for failure",

python/cudf/cudf/tests/groupby/test_agg.py

@@ -804,3 +804,32 @@ def test_groupby_idxminmax_all_na_group_raises(op):
         getattr(pdf.groupby("key"), op)()
     with pytest.raises(ValueError):
         getattr(gdf.groupby("key"), op)()
+
+
+def test_agg_multiindex_columns_preserved():
+    # aggregating a MultiIndex-column frame keeps hierarchical columns
+    pdf = pd.DataFrame(
+        [[1, 2, 3], [1, 5, 6], [2, 8, 9]],
+        columns=pd.MultiIndex.from_tuples(
+            [("k", ""), ("x", "a"), ("x", "b")], names=["l0", "l1"]
+        ),
+    )
+    gdf = cudf.DataFrame(pdf)
+
+    expect = pdf.groupby(("k", "")).agg("sum")
+    got = gdf.groupby(("k", "")).agg("sum")
+    assert_eq(expect, got)
+
+
+def test_agg_relabel_flat_columns_from_multiindex():
+    # relabeling aggregations emit new flat labels; the source's
+    # multi-level column metadata must not be attached to them
+    pdf = pd.DataFrame(
+        [[1, 2], [1, 5], [2, 8]],
+        columns=pd.MultiIndex.from_tuples([("k", ""), ("x", "a")]),
+    )
+    gdf = cudf.DataFrame(pdf)
+
+    expect = pdf.groupby(("k", "")).agg(total=(("x", "a"), "sum"))
+    got = gdf.groupby(("k", "")).agg(total=(("x", "a"), "sum"))
+    assert_eq(expect, got)