BUG: GroupBy aggregation of DataFrame with MultiIndex columns breaks with custom function

doc/source/whatsnew/v1.0.2.rst

@@ -19,6 +19,7 @@ Fixed regressions
 - Fixed regression in :meth:`Series.align` when ``other`` is a DataFrame and ``method`` is not None (:issue:`31785`)
 - Fixed regression in :meth:`pandas.core.groupby.RollingGroupby.apply` where the ``raw`` parameter was ignored (:issue:`31754`)
 - Fixed regression in :meth:`rolling(..).corr() <pandas.core.window.Rolling.corr>` when using a time offset (:issue:`31789`)
+- Fixed regression in :meth:`Groupby.aggregate` which was failing on frames with MultiIndex columns and a custom function (:issue:`31777`)
 - Fixed regression in :meth:`DataFrameGroupBy.nunique` which was modifying the original values if ``NaN`` values were present (:issue:`31950`)
 - Fixed regression where :func:`read_pickle` raised a ``UnicodeDecodeError`` when reading a py27 pickle with :class:`MultiIndex` column (:issue:`31988`).
 - Fixed regression in :class:`DataFrame` arithmetic operations with mis-matched columns (:issue:`31623`)

pandas/core/groupby/generic.py

@@ -955,9 +955,11 @@ def aggregate(self, func=None, *args, **kwargs):
                         raise
                     result = self._aggregate_frame(func)
                 else:
-                    result.columns = Index(
-                        result.columns.levels[0], name=self._selected_obj.columns.name
-                    )
+                    # select everything except for the last level, which is the one
+                    # containing the name of the function(s), see GH 32040
+                    result.columns = result.columns.rename(
+                        [self._selected_obj.columns.name] * result.columns.nlevels
+                    ).droplevel(-1)
 
         if not self.as_index:
             self._insert_inaxis_grouper_inplace(result)

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -691,6 +691,28 @@ def test_agg_relabel_multiindex_duplicates():
     tm.assert_frame_equal(result, expected)
 
 
+@pytest.mark.parametrize(
+    "func, expected_values",
+    [
+        (lambda s: s.mean(), [[3, 2], [5.5, 8.0], [1.5, 3.0], [6.0, 5.5]]),
+        (np.mean, [[3.0, 2.0], [5.5, 8.0], [1.5, 3.0], [6.0, 5.5]]),
+        (np.nanmean, [[3.0, 2.0], [5.5, 8.0], [1.5, 3.0], [6.0, 5.5]]),
+    ],
+)
+def test_multiindex_custom_func(func, expected_values):
+    # GH 31777
+    data = [[1, 4, 2, 8], [5, 7, 1, 4], [2, 8, 1, 4], [2, 8, 5, 7]]
+    df = pd.DataFrame(data, columns=pd.MultiIndex.from_product([[1, 2], [3, 4]]))
+    grp = df.groupby(np.r_[np.zeros(2), np.ones(2)])
+    result = grp.agg(func)
+    expected_keys = [(1, 3), (1, 4), (2, 3), (2, 4)]
+    expected = pd.DataFrame(
+        {key: value for key, value in zip(expected_keys, expected_values)},
+        index=Index([0.0, 1.0], dtype=float),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
 def myfunc(s):
     return np.percentile(s, q=0.90)