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

- [x] closes #31777 
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry


Labels:
name: Groupby
name: MultiIndex

Comments:
pls merge master as well.
@WillAyd @jreback thanks for your reviews.

> Rather than parametrize this test can you just construct the expectation from literal values?

Sure, done. Simplified test a bit, but it still reproduces the error:

```
$ git checkout upstream/master -- pandas/core
$ pytest pandas/tests/groupby/aggregate/test_aggregate.py::test_multiindex_custom_func
```

<details><summary>Full output</summary>

```
========================================================================================= test session starts =========================================================================================
platform linux -- Python 3.7.6, pytest-5.3.4, py-1.8.0, pluggy-0.13.0
rootdir: /home/SERILOCAL/m.gorelli/pandas, inifile: setup.cfg
plugins: forked-1.1.2, hypothesis-5.3.0, asyncio-0.10.0, cov-2.8.1, xdist-1.31.0
collected 3 items                                                                                                                                                                                     

pandas/tests/groupby/aggregate/test_aggregate.py FFF                                                                                                                                            [100%]

============================================================================================== FAILURES ===============================================================================================
_______________________________________________________________________________ test_multiindex_custom_func[<lambda>0] ________________________________________________________________________________

func = <function <lambda> at 0x7fb76ca12170>

    @pytest.mark.parametrize(
        "func", [lambda s: s.mean(), lambda s: np.mean(s), lambda s: np.nanmean(s)]
    )
    def test_multiindex_custom_func(func):
        # GH 31777
        data = [[1, 4, 2], [5, 7, 1]]
        df = pd.DataFrame(data, columns=pd.MultiIndex.from_arrays([[1, 1, 2], [3, 4, 3]]))
>       result = df.groupby(np.array([0, 1])).agg(func)

pandas/tests/groupby/aggregate/test_aggregate.py:701: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pandas/core/groupby/generic.py:959: in aggregate
    result.columns.levels[0], name=self._selected_obj.columns.name
pandas/core/generic.py:5178: in __setattr__
    return object.__setattr__(self, name, value)
pandas/_libs/properties.pyx:67: in pandas._libs.properties.AxisProperty.__set__
    obj._set_axis(self.axis, value)
pandas/core/generic.py:565: in _set_axis
    self._data.set_axis(axis, labels)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = BlockManager
Items: MultiIndex([(1, 3, '<lambda>'),
            (1, 4, '<lambda>'),
            (2, 3, '<lambda>')],
 ... 1, 1), 1 x 2, dtype: int64
IntBlock: slice(1, 2, 1), 1 x 2, dtype: int64
IntBlock: slice(2, 3, 1), 1 x 2, dtype: int64
axis = 0, new_labels = Int64Index([1, 2], dtype='int64')

    def set_axis(self, axis: int, new_labels: Index) -> None:
        # Caller is responsible for ensuring we have an Index object.
        old_len = len(self.axes[axis])
        new_len = len(new_labels)
    
        if new_len != old_len:
            raise ValueError(
>               f"Length mismatch: Expected axis has {old_len} elements, new "
                f"values have {new_len} elements"
            )
E           ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements

pandas/core/internals/managers.py:223: ValueError
_______________________________________________________________________________ test_multiindex_custom_func[<lambda>1] ________________________________________________________________________________

func = <function <lambda> at 0x7fb76ca12200>

    @pytest.mark.parametrize(
        "func", [lambda s: s.mean(), lambda s: np.mean(s), lambda s: np.nanmean(s)]
    )
    def test_multiindex_custom_func(func):
        # GH 31777
        data = [[1, 4, 2], [5, 7, 1]]
        df = pd.DataFrame(data, columns=pd.MultiIndex.from_arrays([[1, 1, 2], [3, 4, 3]]))
>       result = df.groupby(np.array([0, 1])).agg(func)

pandas/tests/groupby/aggregate/test_aggregate.py:701: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pandas/core/groupby/generic.py:959: in aggregate
    result.columns.levels[0], name=self._selected_obj.columns.name
pandas/core/generic.py:5178: in __setattr__
    return object.__setattr__(self, name, value)
pandas/_libs/properties.pyx:67: in pandas._libs.properties.AxisProperty.__set__
    obj._set_axis(self.axis, value)
pandas/core/generic.py:565: in _set_axis
    self._data.set_axis(axis, labels)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = BlockManager
Items: MultiIndex([(1, 3, '<lambda>'),
            (1, 4, '<lambda>'),
            (2, 3, '<lambda>')],
 ... 1, 1), 1 x 2, dtype: int64
IntBlock: slice(1, 2, 1), 1 x 2, dtype: int64
IntBlock: slice(2, 3, 1), 1 x 2, dtype: int64
axis = 0, new_labels = Int64Index([1, 2], dtype='int64')

    def set_axis(self, axis: int, new_labels: Index) -> None:
        # Caller is responsible for ensuring we have an Index object.
        old_len = len(self.axes[axis])
        new_len = len(new_labels)
    
        if new_len != old_len:
            raise ValueError(
>               f"Length mismatch: Expected axis has {old_len} elements, new "
                f"values have {new_len} elements"
            )
E           ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements

pandas/core/internals/managers.py:223: ValueError
_______________________________________________________________________________ test_multiindex_custom_func[<lambda>2] ________________________________________________________________________________

func = <function <lambda> at 0x7fb76ca12290>

    @pytest.mark.parametrize(
        "func", [lambda s: s.mean(), lambda s: np.mean(s), lambda s: np.nanmean(s)]
    )
    def test_multiindex_custom_func(func):
        # GH 31777
        data = [[1, 4, 2], [5, 7, 1]]
        df = pd.DataFrame(data, columns=pd.MultiIndex.from_arrays([[1, 1, 2], [3, 4, 3]]))
>       result = df.groupby(np.array([0, 1])).agg(func)

pandas/tests/groupby/aggregate/test_aggregate.py:701: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pandas/core/groupby/generic.py:959: in aggregate
    result.columns.levels[0], name=self._selected_obj.columns.name
pandas/core/generic.py:5178: in __setattr__
    return object.__setattr__(self, name, value)
pandas/_libs/properties.pyx:67: in pandas._libs.properties.AxisProperty.__set__
    obj._set_axis(self.axis, value)
pandas/core/generic.py:565: in _set_axis
    self._data.set_axis(axis, labels)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = BlockManager
Items: MultiIndex([(1, 3, '<lambda>'),
            (1, 4, '<lambda>'),
            (2, 3, '<lambda>')],
 ... 1, 1), 1 x 2, dtype: int64
IntBlock: slice(1, 2, 1), 1 x 2, dtype: int64
IntBlock: slice(2, 3, 1), 1 x 2, dtype: int64
axis = 0, new_labels = Int64Index([1, 2], dtype='int64')

    def set_axis(self, axis: int, new_labels: Index) -> None:
        # Caller is responsible for ensuring we have an Index object.
        old_len = len(self.axes[axis])
        new_len = len(new_labels)
    
        if new_len != old_len:
            raise ValueError(
>               f"Length mismatch: Expected axis has {old_len} elements, new "
                f"values have {new_len} elements"
            )
E           ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements

pandas/core/internals/managers.py:223: ValueError
```

</details>

> when is nlevels actually > 1 here?

Always.

E.g. if our columns are multiindices: if we run
```
In [1]: import pandas as pd                                                                                                                                                                            

In [2]:     data = [[1, 4, 2], [5, 7, 1]] 
   ...:     df = pd.DataFrame(data, columns=pd.MultiIndex.from_arrays([[1, 1, 2], [3, 4, 3]])) 
   ...:                                                                                                                                                                                                

In [3]: import numpy as np                                                                                                                                                                             

In [4]: df.groupby(np.array([0, 1])).agg(lambda x: x.mean()) 
```
then when we get here, `result.columns` is
```
MultiIndex([(1, 3, '<lambda>'),
            (1, 4, '<lambda>'),
            (2, 3, '<lambda>')],
           )
```

E.g. if our columns are not multiindices: if we run
```
In [5]: df = pd.DataFrame(data, columns=['col1', 'col2', 'col3']) 

In [6]: df.groupby(np.array([0, 1])).agg(lambda x: x.mean()) 
```
then when we get here, `result.columns` is
```
MultiIndex([('col1', '<lambda>'),
            ('col2', '<lambda>'),
            ('col3', '<lambda>')],
           )
```

Furthermore, I tried adding in
```
assert result.columns.nlevels > 1
```
and all the tests in `pandas/tests/groupby` still passed
@MarcoGorelli there's a merge conflict in 1.0.2.rst if you could update.
@TomAugspurger sure, done (I presume the CI failure is unrelated?)
thanks @MarcoGorelli 
@meeseeksdev backport 1.0.x
Something went wrong ... Please have  a look at my logs.