Repository: pyjanitor-devs/pyjanitor (python)

Issue:
`summarise` fails for MultiIndex
# Brief Description

`summarise` does not handle multiindex columns

# System Information


- Operating system: macOS

- Python version (required):3.12

# Minimally Reproducible Code

```
def dfmi():
    """Create a MultiIndex DataFrame"""

    # https://pandas.pydata.org/docs/user_guide/advanced.html#using-slicers
    def mklbl(prefix, n):
        return ["%s%s" % (prefix, i) for i in range(n)]

    miindex = pd.MultiIndex.from_product(
        [mklbl("A", 4), mklbl("B", 2), mklbl("C", 4), mklbl("D", 2)]
    )

    micolumns = pd.MultiIndex.from_tuples(
        [("a", "foo"), ("a", "bar"), ("b", "foo"), ("b", "bah")],
        names=["lvl0", "lvl1"],
    )
    dfmi = (
        pd.DataFrame(
            np.arange(len(miindex) * len(micolumns)).reshape(
                (len(miindex), len(micolumns))
            ),
            index=miindex,
            columns=micolumns,
        )
        .sort_index()
        .sort_index(axis=1)
    )
    dfmi.index.names = list("ABCD")
    return dfmi


df = dfmi()
df.summarise(('a',['min','max']),by='a')
```

# Error Messages

```
ValueError: Cannot subset columns with a tuple with more than one element. Use a list instead.
```


Pull request:
[ENH] Fix `summarise` for MultiIndex
# PR Description

Please describe the changes proposed in the pull request:

- Fix error when MultiIndex is used with `summarise`


```
def dfmi():
    """Create a MultiIndex DataFrame"""

    # https://pandas.pydata.org/docs/user_guide/advanced.html#using-slicers
    def mklbl(prefix, n):
        return ["%s%s" % (prefix, i) for i in range(n)]

    miindex = pd.MultiIndex.from_product(
        [mklbl("A", 4), mklbl("B", 2), mklbl("C", 4), mklbl("D", 2)]
    )

    micolumns = pd.MultiIndex.from_tuples(
        [("a", "foo"), ("a", "bar"), ("b", "foo"), ("b", "bah")],
        names=["lvl0", "lvl1"],
    )
    dfmi = (
        pd.DataFrame(
            np.arange(len(miindex) * len(micolumns)).reshape(
                (len(miindex), len(micolumns))
            ),
            index=miindex,
            columns=micolumns,
        )
        .sort_index()
        .sort_index(axis=1)
    )
    dfmi.index.names = list("ABCD")
    return dfmi


df = dfmi()
df.summarise(('a',['min','max']),by=('a','bar'))

lvl0        a               
lvl1      bar       foo     
          min  max  min  max
(a, bar)                    
1           1    1    0    0
5           5    5    4    4
9           9    9    8    8
13         13   13   12   12
17         17   17   16   16
...       ...  ...  ...  ...
237       237  237  236  236
241       241  241  240  240
245       245  245  244  244
249       249  249  248  248
253       253  253  252  252

[64 rows x 4 columns]

```


**This PR resolves #1459 .**

<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [x] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.md`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.md` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

# Automatic checks

There will be automatic checks run on the PR. These include:

- Building a preview of the docs on Netlify
- Automatically linting the code
- Making sure the code is documented
- Making sure that all tests are passed
- Making sure that code coverage doesn't go down.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/janitor/functions/mutate.py b/janitor/functions/mutate.py
index 38fab38..5ef865a 100644
--- a/janitor/functions/mutate.py
+++ b/janitor/functions/mutate.py
@@ -177,6 +177,8 @@ def mutate(
         A pandas DataFrame or Series with aggregated columns.
     """  # noqa: E501
     check("copy", copy, [bool])
+    if copy:
+        df = df.copy(deep=None)
     if by is not None:
         if isinstance(by, DataFrameGroupBy):
             # it is assumed that by is created from df
@@ -188,8 +190,7 @@ def mutate(
             if is_scalar(by):
                 by = [by]
             by = df.groupby(by, sort=False, observed=True)
-    if copy:
-        df = df.copy(deep=None)
+
     for arg in args:
         df = _mutator(arg, df=df, by=by)
     return df
@@ -226,11 +227,9 @@ def _(arg, df, by):
     for column_name, mutator in arg.items():
         if isinstance(mutator, tuple):
             column, func = mutator
-            column = _process_within_dict(mutator=func, obj=val[column])
+            column = _apply_func_to_obj(mutator=func, obj=val[column])
         else:
-            column = _process_within_dict(
-                mutator=mutator, obj=val[column_name]
-            )
+            column = _apply_func_to_obj(mutator=mutator, obj=val[column_name])
         df[column_name] = column
     return df
 
@@ -262,7 +261,7 @@ def _process_maybe_string(func: str, obj):
     return obj.transform(func)
 
 
-def _process_within_dict(mutator, obj):
+def _apply_func_to_obj(mutator, obj):
     """Handle str/callables within a dictionary"""
     if isinstance(mutator, str):
         return _process_maybe_string(func=mutator, obj=obj)
diff --git a/janitor/functions/summarise.py b/janitor/functions/summarise.py
index 980cf6d..b2da04e 100644
--- a/janitor/functions/summarise.py
+++ b/janitor/functions/summarise.py
@@ -19,7 +19,7 @@ def summarise(
     df: pd.DataFrame,
     *args: tuple[dict | tuple],
     by: Any = None,
-) -> pd.DataFrame | pd.Series:
+) -> pd.DataFrame:
     """
 
     !!! info "New in version 0.31.0"
@@ -42,13 +42,15 @@ def summarise(
     - **dictionary argument**:
     If the argument is a dictionary,
     the value in the `{key:value}` pairing
-    should be either a string, a callable or a tuple.
+    should be either a string, a callable, or a tuple.
 
         - If the value in the dictionary
         is a string or a callable,
         the key of the dictionary
         should be an existing column name.
 
+        The function is applied on the `df[column_name]` series.
+
         !!!note
 
             - If the value is a string,
@@ -57,28 +59,24 @@ def summarise(
 
         - If the value of the dictionary is a tuple,
         it should be of length 2, and of the form
-        `(column_name, mutation_func)`,
+        `(column_name, aggfunc)`,
         where `column_name` should exist in the DataFrame,
-        and `mutation_func` should be either a string or a callable.
-
-        !!!note
+        and `aggfunc` should be either a string or a callable.
 
-            - If `mutation_func` is a string,
-            the string should be a pandas string function,
-            e.g "sum", "mean", etc.
+        This option allows for custom renaming of the aggregation output,
+        where the key in the dictionary can be a new column name.
 
-        The key in the dictionary can be a new column name.
 
     - **tuple argument**:
     If the argument is a tuple, it should be of length 2,
     and of the form
-    `(column_name, mutation_func)`,
+    `(column_name, aggfunc)`,
     where column_name should exist in the DataFrame,
-    and `mutation_func` should be either a string or a callable.
+    and `aggfunc` should be either a string or a callable.
 
         !!!note
 
-            - if `mutation_func` is a string,
+            - if `aggfunc` is a string,
             the string should be a pandas string function,
             e.g "sum", "mean", etc.
 
@@ -89,6 +87,7 @@ def summarise(
             as such multiple columns can be processed here -
             they will be processed individually.
 
+
     - **callable argument**:
     If the argument is a callable, the callable is applied
     on the DataFrame or GroupBy object.
@@ -169,7 +168,7 @@ def summarise(
         ValueError: If a tuple is passed and the length is not 2.
 
     Returns:
-        A pandas DataFrame or Series with aggregated columns.
+        A pandas DataFrame with aggregated columns.
 
     """  # noqa: E501
 
@@ -184,18 +183,52 @@ def summarise(
             if is_scalar(by):
                 by = [by]
             by = df.groupby(by, sort=False, observed=True)
-    dictionary = {}
+    contents = []
     for arg in args:
-        aggregate = _mutator(arg, df=df, by=by)
-        dictionary.update(aggregate)
-    values = map(is_scalar, dictionary.values())
-    if all(values):
-        return pd.Series(dictionary)
-    return pd.concat(dictionary, axis="columns", sort=False, copy=False)
+        aggregate = _aggfunc(arg, df=df, by=by)
+        contents.extend(aggregate)
+    counts = 0
+    for entry in contents:
+        if isinstance(entry, pd.DataFrame):
+            length = entry.columns.nlevels
+        elif isinstance(entry.name, tuple):
+            length = len(entry.name)
+        else:
+            length = 1
+        counts = max(counts, length)
+    contents_ = []
+    for entry in contents:
+        if isinstance(entry, pd.DataFrame):
+            length_ = entry.columns.nlevels
+            length = counts - length_
+            if length:
+                patch = [""] * length
+                columns = [
+                    entry.columns.get_level_values(n) for n in range(length_)
+                ]
+                columns.append(patch)
+                names = [*entry.columns.names]
+                names.extend([None] * length)
+                columns = pd.MultiIndex.from_arrays(columns, names=names)
+                entry.columns = columns
+        elif is_scalar(entry.name):
+            length = counts - 1
+            if length:
+                patch = [""] * length
+                name = (entry.name, *patch)
+                entry.name = name
+        elif isinstance(entry.name, tuple):
+            length = counts - len(entry.name)
+            if length:
+                patch = [""] * length
+                name = (*entry.name, *patch)
+                entry.name = name
+        contents_.append(entry)
+    return pd.concat(contents_, axis=1, copy=False, sort=False)
 
 
 @singledispatch
-def _mutator(arg, df, by):
+def _aggfunc(arg, df, by):
     if by is None:
         val = df
     else:
@@ -204,12 +237,26 @@ def _mutator(arg, df, by):
     if isinstance(outcome, pd.Series):
         if not outcome.name:
             raise ValueError("Ensure the pandas Series object has a name")
-        return {outcome.name: outcome}
-    # assumption: a mapping - DataFrame/dictionary/...
-    return {**outcome}
+        return [outcome]
+    if isinstance(outcome, pd.DataFrame):
+        return [outcome]
+    raise TypeError(
+        "The output from the aggregation should be a named Series or a DataFrame"
+    )
 
 
-@_mutator.register(dict)
+@_aggfunc.register(tuple)
+def _(arg, df, by):
+    """Dispatch function for tuple"""
+    if len(arg) != 2:
+        raise ValueError("the tuple has to be a length of 2")
+    column_name, aggfunc = arg
+    column_names = get_index_labels(arg=[column_name], df=df, axis="columns")
+    mapping = {column_name: aggfunc for column_name in column_names}
+    return _aggfunc(mapping, df=df, by=by)
+
+
+@_aggfunc.register(dict)
 def _(arg, df, by):
     """Dispatch function for dictionary"""
     if by is None:
@@ -217,28 +264,41 @@ def _(arg, df, by):
     else:
         val = by
 
-    dictionary = {}
-    for column_name, mutator in arg.items():
-        if isinstance(mutator, tuple):
-            column, func = mutator
-            column = _process_within_dict(mutator=func, obj=val[column])
+    contents = []
+    for column_name, aggfunc in arg.items():
+        if isinstance(aggfunc, tuple):
+            if len(aggfunc) != 2:
+                raise ValueError("the tuple has to be a length of 2")
+            column, func = aggfunc
+            column_ = _handle_tuple_groupby_selection(by=by, column=column)
+            column = _apply_func_to_obj(aggfunc=func, obj=val[column_])
+            if isinstance(column, pd.DataFrame) and column.shape[-1] == 1:
+                column = column.squeeze()
+            column = _convert_obj_to_named_series(
+                obj=column,
+                column_name=column_name,
+                function=func,
+            )
+            if not isinstance(column, pd.Series):
+                raise TypeError(
+                    "Expected a pandas Series object; "
+                    f"instead got {type(column)}"
+                )
         else:
-            column = _process_within_dict(
-                mutator=mutator, obj=val[column_name]
+            column_ = _handle_tuple_groupby_selection(
+                by=by, column=column_name
             )
-        dictionary[column_name] = column
-    return dictionary
-
-
-@_mutator.register(tuple)
-def _(arg, df, by):
-    """Dispatch function for tuple"""
-    if len(arg) != 2:
-        raise ValueError("the tuple has to be a length of 2")
-    column_names, mutator = arg
-    column_names = get_index_labels(arg=[column_names], df=df, axis="columns")
-    mapping = {column_name: mutator for column_name in column_names}
-    return _mutator(mapping, df=df, by=by)
+            column = _apply_func_to_obj(aggfunc=aggfunc, obj=val[column_])
+            column = _convert_obj_to_named_series(
+                obj=column,
+                column_name=column_name,
+                function=aggfunc,
+            )
+        column = _rename_column_in_by(
+            column=column, column_name=column_name, by=by
+        )
+        contents.append(column)
+    return contents
 
 
 def _process_maybe_callable(func: callable, obj):
@@ -257,8 +317,39 @@ def _process_maybe_string(func: str, obj):
     return obj.agg(func)
 
 
-def _process_within_dict(mutator, obj):
+def _apply_func_to_obj(aggfunc, obj):
     """Handle str/callables within a dictionary"""
-    if isinstance(mutator, str):
-        return _process_maybe_string(func=mutator, obj=obj)
-    return _process_maybe_callable(func=mutator, obj=obj)
+    if isinstance(aggfunc, str):
+        return _process_maybe_string(func=aggfunc, obj=obj)
+    return _process_maybe_callable(func=aggfunc, obj=obj)
+
+
+def _handle_tuple_groupby_selection(by: Any, column: Any):
+    """
+    Properly handle a tuple column selection in the presence of a groupby
+    """
+    if (by is not None) and isinstance(column, tuple):
+        return [column]
+    return column
+
+
+def _convert_obj_to_named_series(obj, function: Any, column_name: Any):
+    if isinstance(obj, pd.Series):
+        obj.name = column_name
+        return obj
+    if not is_scalar(obj):
+        return obj
+    if isinstance(function, str):
+        function_name = function
+    else:
+        function_name = function.__name__
+    return pd.Series(data=obj, index=[function_name], name=column_name)
+
+
+def _rename_column_in_by(column, column_name, by):
+    if by is None:
+        return column
+    elif isinstance(column, pd.DataFrame) and is_scalar(column_name):
+        columns = pd.MultiIndex.from_product([[column_name], column.columns])
+        column.columns = columns
+    return column


Test patch:
diff --git a/tests/functions/test_summarise.py b/tests/functions/test_summarise.py
index 59e899e..0210162 100644
--- a/tests/functions/test_summarise.py
+++ b/tests/functions/test_summarise.py
@@ -1,6 +1,7 @@
+import numpy as np
 import pandas as pd
 import pytest
-from pandas.testing import assert_frame_equal, assert_series_equal
+from pandas.testing import assert_frame_equal
 
 
 @pytest.fixture
@@ -13,7 +14,38 @@ def df_summarise():
     return pd.DataFrame(data)
 
 
-def test_mutate_callable_series_unnamed(df_summarise):
+@pytest.fixture
+def dfmi():
+    """Create a MultiIndex DataFrame"""
+
+    # https://pandas.pydata.org/docs/user_guide/advanced.html#using-slicers
+    def mklbl(prefix, n):
+        return ["%s%s" % (prefix, i) for i in range(n)]
+
+    miindex = pd.MultiIndex.from_product(
+        [mklbl("A", 4), mklbl("B", 2), mklbl("C", 4), mklbl("D", 2)]
+    )
+
+    micolumns = pd.MultiIndex.from_tuples(
+        [("a", "foo"), ("a", "bar"), ("b", "foo"), ("b", "bah")],
+        names=["lvl0", "lvl1"],
+    )
+    dfmi = (
+        pd.DataFrame(
+            np.arange(len(miindex) * len(micolumns)).reshape(
+                (len(miindex), len(micolumns))
+            ),
+            index=miindex,
+            columns=micolumns,
+        )
+        .sort_index()
+        .sort_index(axis=1)
+    )
+    dfmi.index.names = list("ABCD")
+    return dfmi
+
+
+def test_summarise_callable_series_unnamed(df_summarise):
     """Test output for callable"""
     with pytest.raises(
         ValueError, match="Ensure the pandas Series object has a name"
@@ -21,7 +53,16 @@ def test_mutate_callable_series_unnamed(df_summarise):
         df_summarise.summarise(lambda df: df.mean())
 
 
-def test_mutate_callable_series(df_summarise):
+def test_summarise_callable_not_pandas_object(df_summarise):
+    """Raise if output is not a pandas Series/DataFrame"""
+    with pytest.raises(
+        TypeError,
+        match="The output from the aggregation should be a named Series or a DataFrame",
+    ):
+        df_summarise.summarise(lambda df: df.avg_jump.mean())
+
+
+def test_summarise_callable_series(df_summarise):
     """Test output for callable"""
     expected = df_summarise.summarise(lambda df: df.mean().rename("mean"))
     actual = df_summarise.mean().to_frame("mean")
@@ -39,8 +80,8 @@ def test_summarise_by_callable_grp(df_summarise):
 def test_summarise_dict_df_str(df_summarise):
     """Test output for a dictionary"""
     actual = df_summarise.summarise({"avg_run": "mean"})
-    expected = df_summarise.agg({"avg_run": "mean"})
-    assert_series_equal(actual, expected)
+    expected = df_summarise.agg({"avg_run": ["mean"]})
+    assert_frame_equal(actual, expected)
 
 
 def test_summarise_dict_by_str(df_summarise):
@@ -55,8 +96,8 @@ def test_summarise_dict_by_str(df_summarise):
 def test_summarise_dict_df_callable(df_summarise):
     """Test output for a dictionary"""
     actual = df_summarise.summarise({"avg_run": lambda df: df.sum()})
-    expected = df_summarise.agg({"avg_run": "sum"})
-    assert_series_equal(actual, expected)
+    expected = df_summarise.agg({"avg_run": [lambda df: df.sum()]})
+    assert_frame_equal(actual, expected)
 
 
 def test_summarise_dict_by_callable(df_summarise):
@@ -79,12 +120,10 @@ def test_summarise_dict_by_agg_callable(df_summarise):
 
 def test_summarise_dict_df_tuple(df_summarise):
     """Test output for a dictionary"""
-    actual = (
-        df_summarise.summarise({"avg_run_sum": ("avg_run", "sum")})
-        .rename("avg_run")
-        .to_frame()
+    actual = df_summarise.summarise({"avg_run_sum": ("avg_run", "sum")})
+    expected = df_summarise.agg(avg_run_sum=("avg_run", "sum")).T.rename(
+        index={"avg_run": "sum"}
     )
-    expected = df_summarise.agg(avg_run_sum=("avg_run", "sum"))
     assert_frame_equal(actual, expected)
 
 
@@ -99,6 +138,12 @@ def test_summarise_dict_by_tuple(df_summarise):
     assert_frame_equal(actual, expected)
 
 
+def test_summarise_dict_tuple_count_not_eq_2(df_summarise):
+    """Raise error if length of tuple is not 2"""
+    with pytest.raises(ValueError, match="the tuple has to be a length of 2"):
+        df_summarise.summarise({"avg_run_sum": ("avg_run",)})
+
+
 def test_summarise_tuple_count_not_eq_2(df_summarise):
     """Raise error if length of tuple is not 2"""
     with pytest.raises(ValueError, match="the tuple has to be a length of 2"):
@@ -108,8 +153,8 @@ def test_summarise_tuple_count_not_eq_2(df_summarise):
 def test_summarise_df_tuple(df_summarise):
     "Test output for a tuple"
     actual = df_summarise.summarise(("avg_run", "sum"))
-    expected = df_summarise.agg({"avg_run": "sum"})
-    assert_series_equal(actual, expected)
+    expected = df_summarise.agg({"avg_run": ["sum"]})
+    assert_frame_equal(actual, expected)
 
 
 def test_summarise_by_tuple(df_summarise):
@@ -122,8 +167,8 @@ def test_summarise_by_tuple(df_summarise):
 def test_summarise_tuple_df_callable(df_summarise):
     """Test output for a tuple"""
     actual = df_summarise.summarise(("avg_run", lambda df: df.sum()))
-    expected = df_summarise.agg({"avg_run": "sum"})
-    assert_series_equal(actual, expected)
+    expected = df_summarise.agg({"avg_run": [lambda df: df.sum()]})
+    assert_frame_equal(actual, expected)
 
 
 def test_summarise_tuple_by_callable(df_summarise):
@@ -154,3 +199,84 @@ def test_summarise_tuple_grouped_object(df_summarise):
     )
     expected = grp.agg({"avg_run": ["sum", "mean"]})
     assert_frame_equal(actual, expected)
+
+
+def test_summarise_MI(dfmi):
+    """Test summarise on a MultiIndex"""
+    actual = dfmi.summarise(("a", "min"))
+    expected = dfmi.agg(
+        {("a", "bar"): "min", ("a", "foo"): ["min"]}
+    ).rename_axis(columns=[None, None])
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_MI_different_levels(dfmi):
+    """Test summarise on a MultiIndex"""
+    actual = dfmi.summarise(
+        {("a", "bar"): "sum", ("rar", "bla", "gah"): (("a", "foo"), "mean")},
+        ("b", "min"),
+        by={"level": "A"},
+    )
+    actual.columns = ["A", "B", "C", "D"]
+    grp = dfmi.groupby(level="A")
+    expected = grp.agg(
+        {
+            ("a", "bar"): "sum",
+            ("a", "foo"): "mean",
+            ("b", "bah"): "min",
+            ("b", "foo"): "min",
+        }
+    )
+    expected.columns = ["A", "B", "C", "D"]
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_MI_different_levels_dict(dfmi):
+    """Test summarise on a MultiIndex"""
+    actual = dfmi.summarise(
+        {("a", "bar"): "sum", "rar": (("a", "foo"), "mean")},
+        ("b", "min"),
+        by={"level": "A"},
+    )
+    actual.columns = ["A", "B", "C", "D"]
+    grp = dfmi.groupby(level="A")
+    expected = grp.agg(
+        {
+            ("a", "bar"): "sum",
+            ("a", "foo"): "mean",
+            ("b", "bah"): "min",
+            ("b", "foo"): "min",
+        }
+    )
+    expected.columns = ["A", "B", "C", "D"]
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_MI_different_levels_tuple(dfmi):
+    """Test summarise on a MultiIndex"""
+    actual = dfmi.summarise(
+        {("a", "bar"): "sum", ("rar",): (("a", "foo"), "mean")},
+        ("b", "min"),
+        by={"level": "A"},
+    )
+    actual.columns = ["A", "B", "C", "D"]
+    grp = dfmi.groupby(level="A")
+    expected = grp.agg(
+        {
+            ("a", "bar"): "sum",
+            ("a", "foo"): "mean",
+            ("b", "bah"): "min",
+            ("b", "foo"): "min",
+        }
+    )
+    expected.columns = ["A", "B", "C", "D"]
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_MI_different_levels_dataframe(dfmi):
+    """raise if dictionary value is a tuple and the returned aggregate is a DataFrame"""
+    with pytest.raises(TypeError, match="Expected a pandas Series object;.+"):
+        dfmi.summarise(
+            {"rar": (("a", "foo"), ["min", "mean"])},
+            by={"level": "A"},
+        )