fix setting of index

Author: MarcoGorelli

doc/source/whatsnew/v1.0.2.rst

@@ -19,7 +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`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/groupby/generic.py

@@ -945,7 +945,6 @@ def aggregate(self, func=None, *args, **kwargs):
             return result
 
         if result is None:
-
             # grouper specific aggregations
             if self.grouper.nkeys > 1:
                 return self._python_agg_general(func, *args, **kwargs)
@@ -968,7 +967,8 @@ def aggregate(self, func=None, *args, **kwargs):
                     result = self._aggregate_frame(func)
                 else:
                     result.columns = Index(
-                        result.columns.levels[0], name=self._selected_obj.columns.name
+                        [i[:-1] if len(i) > 2 else i[0] for i in result.columns],
+                        name=self._selected_obj.columns.name,
                     )
 
         if not self.as_index:

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -679,6 +679,17 @@ def test_agg_relabel_multiindex_duplicates():
     tm.assert_frame_equal(result, expected)
 
 
+def test_multiindex_custom_func():
+    # GH 31777
+    df = pd.DataFrame(
+        np.random.rand(10, 4), columns=pd.MultiIndex.from_product([[1, 2], [3, 4]])
+    )
+    grp = df.groupby(np.r_[np.ones(5), np.zeros(5)])
+    result = grp.agg(lambda s: s.mean())
+    expected = grp.agg("mean")
+    tm.assert_frame_equal(result, expected)
+
+
 def myfunc(s):
     return np.percentile(s, q=0.90)