Repository: pyjanitor-devs/pyjanitor

Issue:
[ENH] make collapse_levels applicable to multilevel index  
collapse levels of dataframes with multilevel index  


Discussion:
@mdini would you like to tackle this one? Let me know, and I'll tag it as unavailable for hacking!
@ericmjl yes please, I would like to. 

Resolving patch:
diff --git a/janitor/functions/collapse_levels.py b/janitor/functions/collapse_levels.py
--- a/janitor/functions/collapse_levels.py
+++ b/janitor/functions/collapse_levels.py
@@ -1,5 +1,7 @@
 """Implementation of the `collapse_levels` function."""
 
+from typing import Union
+
 import pandas as pd
 import pandas_flavor as pf
 from pandas.api.types import is_string_dtype
@@ -8,13 +10,18 @@
 
 
 @pf.register_dataframe_method
-def collapse_levels(df: pd.DataFrame, sep: str = "_") -> pd.DataFrame:
-    """Flatten multi-level column dataframe to a single level.
+def collapse_levels(
+    df: pd.DataFrame,
+    sep: Union[str, None] = None,
+    glue: Union[str, None] = None,
+    axis="columns",
+) -> pd.DataFrame:
+    """Flatten multi-level index/column dataframe to a single level.
 
     This method does not mutate the original DataFrame.
 
-    Given a DataFrame containing multi-level columns, flatten to single-level
-    by string-joining the column labels in each level.
+    Given a DataFrame containing multi-level index/columns, flatten to single-level
+    by string-joining the labels in each level.
 
     After a `groupby` / `aggregate` operation where `.agg()` is passed a
     list of multiple aggregation functions, a multi-level DataFrame is
@@ -53,59 +60,98 @@ class  max_speed    type
         bird        267.333333             389.0
         mammal       50.500000              50.5
 
-    Before applying `.collapse_levels`, the `.agg` operation returns a
-    multi-level column DataFrame whose columns are `(level 1, level 2)`:
+        Before applying `.collapse_levels`, the `.agg` operation returns a
+        multi-level column DataFrame whose columns are `(level 1, level 2)`:
+
+        ```python
+        [("max_speed", "mean"), ("max_speed", "median")]
+        ```
+
+        `.collapse_levels` then flattens the column MultiIndex into a single
+        level index with names:
+
+        ```python
+        ["max_speed_mean", "max_speed_median"]
+        ```
+
+        For more control, a `glue` specification can be passed,
+        where the names of the levels are used to control the output of the
+        flattened index:
+        >>> (grouped_df
+        ...  .rename_axis(columns=['column_name', 'agg_name'])
+        ...  .collapse_levels(glue="{agg_name}_{column_name}")
+        ... )
+                mean_max_speed  median_max_speed
+        class
+        bird        267.333333             389.0
+        mammal       50.500000              50.5
 
-    ```python
-    [("max_speed", "mean"), ("max_speed", "median")]
-    ```
+        Note that for `glue` to work, the keyword arguments
+        in the glue specification
+        should be the names of the levels in the MultiIndex.
 
-    `.collapse_levels` then flattens the column MultiIndex into a single
-    level index with names:
+    !!! abstract "Version Changed"
 
-    ```python
-    ["max_speed_mean", "max_speed_median"]
-    ```
+        - 0.27.0
+            - Added `glue` and `axis` parameters.
 
     Args:
         df: A pandas DataFrame.
         sep: String separator used to join the column level names.
+        glue: A specification on how the column levels should be combined.
+            It allows for a more granular composition,
+            and serves as an alternative to `sep`.
+        axis: Determines whether to collapse the
+            levels on the index or columns.
 
     Returns:
         A pandas DataFrame with single-level column index.
     """  # noqa: E501
-    check("sep", sep, [str])
-
-    # if already single-level, just return the DataFrame
-    if not isinstance(df.columns, pd.MultiIndex):
+    if (sep is not None) and (glue is not None):
+        raise ValueError("Only one of sep or glue should be provided.")
+    if sep is not None:
+        check("sep", sep, [str])
+    if glue is not None:
+        check("glue", glue, [str])
+    check("axis", axis, [str])
+    if axis not in {"index", "columns"}:
+        raise ValueError(
+            "axis argument should be either 'index' or 'columns'."
+        )
+
+    if not isinstance(getattr(df, axis), pd.MultiIndex):
         return df
 
     # TODO: Pyarrow offers faster string computations
     # future work should take this into consideration,
     # which would require a different route from python's string.join
-
     # since work is only on the columns
     # it is safe, and more efficient to slice/view the dataframe
     # plus Pandas creates a new Index altogether
     # as such, the original dataframe is not modified
     df = df[:]
-    new_columns = df.columns
-    levels = [
-        new_columns.get_level_values(num) for num in range(new_columns.nlevels)
-    ]
+    new_index = getattr(df, axis)
+    if glue is not None:
+        new_index = [dict(zip(new_index.names, entry)) for entry in new_index]
+        new_index = [glue.format_map(mapping) for mapping in new_index]
+        setattr(df, axis, new_index)
+        return df
+    sep = "_" if sep is None else sep
+    levels = [level for level in new_index.levels]
     all_strings = all(map(is_string_dtype, levels))
     if all_strings:
         no_empty_string = all((entry != "").all() for entry in levels)
         if no_empty_string:
-            df.columns = new_columns.map(sep.join)
+            new_index = new_index.map(sep.join)
+            setattr(df, axis, new_index)
             return df
-    new_columns = (map(str, entry) for entry in new_columns)
-    new_columns = [
+    new_index = (map(str, entry) for entry in new_index)
+    new_index = [
         # faster to use a list comprehension within string.join
         # compared to a generator
         # https://stackoverflow.com/a/37782238
         sep.join([entry for entry in word if entry])
-        for word in new_columns
+        for word in new_index
     ]
-    df.columns = new_columns
+    setattr(df, axis, new_index)
     return df
diff --git a/janitor/functions/conditional_join.py b/janitor/functions/conditional_join.py
--- a/janitor/functions/conditional_join.py
+++ b/janitor/functions/conditional_join.py
@@ -763,7 +763,8 @@ def _multiple_conditional_join_eq(
             grp = right_df.groupby(equi_col, sort=False)
             non_equi_col = right_columns[1]
             # groupby.is_monotonic_increasing uses apply under the hood
-            # this circumvents the Series creation (which isn't required here)
+            # the approach used below circumvents the Series creation
+            # (which isn't required here)
             # and just gets a sequence of booleans, before calling `all`
             # to get a single True or False.
             right_is_sorted = all(


Test patch:
diff --git a/tests/functions/test_collapse_levels.py b/tests/functions/test_collapse_levels.py
--- a/tests/functions/test_collapse_levels.py
+++ b/tests/functions/test_collapse_levels.py
@@ -1,3 +1,7 @@
+from functools import reduce
+
+import numpy as np
+import pandas as pd
 import pytest
 from pandas.testing import assert_frame_equal
 
@@ -63,3 +67,75 @@ def test_collapse_levels_functionality_3level(
             "decorated-elephantAsDfr.i.p-rhino :'(AsDfdeadly__flamingo",
         ]
     )
+
+
+@pytest.fixture
+def mi_index():
+    """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_sep_and_glue(mi_index):
+    """raise if sep and glue are provided."""
+    msg = "Only one of sep or glue should be provided."
+    with pytest.raises(ValueError, match=msg):
+        mi_index.collapse_levels(sep="_", glue="_")
+
+
+def test_glue_type(mi_index):
+    """raise if glue is not str."""
+    with pytest.raises(TypeError, match="glue should be one of.+"):
+        mi_index.collapse_levels(glue=3)
+
+
+def test_axis_type(mi_index):
+    """raise if axis is not str."""
+    with pytest.raises(TypeError, match="axis should be one of.+"):
+        mi_index.collapse_levels(axis=3)
+
+
+def test_axis_value(mi_index):
+    """raise if axis is not index/columns."""
+    msg = "axis argument should be either 'index' or 'columns'."
+    with pytest.raises(ValueError, match=msg):
+        mi_index.collapse_levels(axis="INDEX")
+
+
+def test_glue_output(mi_index):
+    """test output if glue is provided."""
+    expected = mi_index.collapse_levels(glue="{A}{B}{C}{D}", axis="index")
+    index = reduce(
+        lambda x, y: x + y,
+        [
+            mi_index.index.get_level_values(num)
+            for num in range(mi_index.index.nlevels)
+        ],
+    )
+    actual = mi_index.copy()
+    actual.index = index
+    assert_frame_equal(expected, actual)