Repository: pyjanitor-devs/pyjanitor

Issue:
Feature enhancement: collapsing MultiIndex columns
For example, after aggregation with multiple functions.

For `df` with columns `['group', 'category', 'value']`:

```python
stats_df = (
    df.groupby(['group', 'category'])
    .agg(['mean', 'median'])
    .reset_index()
)
```

produces `stats_df` with a `MultiIndex` `.columns` attribute where `{'mean', 'median'}` are second level column names under `value`. It would be nice if `.columns` was just an `Index` instead, for some applications.

Now, to flatten the `MultiIndex` out into an `Index` by concatenating the different levels with an underscore:

`stats_df.columns.values` is `array([('group', ''), ('category', ''), ('value', 'mean'), ('value', 'median')], dtype=object)`

`stats_df.columns = ['_'.join(tup) if tup[1] != '' else tup[0] for tup in stats_df.columns.values]`

`stats_df.columns` is now `['group', 'category', 'value_mean', 'value_median']`


Discussion:
I'll work on generalizing this into a `.collapse_levels(df, sep='_')`-style hook for a `DataFrame` with an arbitrary number of column levels.
@zbarry thanks for putting this in! I'm sure it'll be a valuable addition. 

I did a second thought on the name of the function, and it may or may not be the most inclusive/accessible one for non-native English speakers. At first glance, I can't think of a better name, but if I can think of one, I'll ping back on the PR to suggest it. Otherwise, please feel free to run with `collapse_levels()`.
Sounds good. I was basing the name on [droplevel](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.MultiIndex.droplevel.html). Maybe a different word than collapse, then?
Maybe so. For now, `collapse_levels` should be fine, if paired with a good docstring. But don't sweat it for now. :smile:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -140,7 +140,7 @@ def _remove_special(col):
 
     # Store the original column names, if enabled by user
     if preserve_original_columns:
-        df.__dict__['original_columns'] = original_column_names
+        df.__dict__["original_columns"] = original_column_names
     return df
 
 
@@ -545,16 +545,12 @@ def fill_empty(df, columns, value):
         for col in columns:
             assert (
                 col in df.columns
-            ), "{col} missing from dataframe columns!".format(
-                col=col
-            )
+            ), "{col} missing from dataframe columns!".format(col=col)
             df[col] = df[col].fillna(value)
     else:
         assert (
             columns in df.columns
-        ), "{col} missing from dataframe columns!".format(
-            col=columns
-        )
+        ), "{col} missing from dataframe columns!".format(col=columns)
         df[columns] = df[columns].fillna(value)
 
     return df
@@ -1533,6 +1529,59 @@ def min_max_scale(
     return df
 
 
+@pf.register_dataframe_method
+def collapse_levels(df: pd.DataFrame, sep: str = "_"):
+    """
+    Given a `DataFrame` containing multi-level columns, flatten to single-
+    level by string-joining the column labels in each level.
+
+    After a `groupby` / `aggregate` operation where `.agg()` is passed a
+    list of multiple aggregation functions, a multi-level `DataFrame` is
+    returned with the name of the function applied in the second level.
+    It is sometimes convenient for later indexing to flatten out this
+    multi-level configuration back into a single level. This function does
+    this through a simple string-joining of all the names across different
+    levels in a single column.
+
+    Method chaining example given two value columns `['var1', 'var2']`:
+
+    .. code-block:: python
+
+        df = (
+            pd.DataFrame(...)
+            .groupby('mygroup')
+            .agg(['mean', 'median'])
+            .collapse_levels(sep='_')
+        )
+
+    Before applying `.collapse_levels`, the `.agg` operation returns a
+    multi-level column `DataFrame` whose columns are (level 1, level 2):
+
+    `[('mygroup', ''), ('var1', 'mean'), ('var1', 'median'), ('var2', 'mean'),
+    ('var2', 'median')]`
+
+    `.collapse_levels` then flattens the column names to:
+
+    `['mygroup', 'var1_mean', 'var1_median', 'var2_mean', 'var2_median']`
+
+    :param df: A pandas DataFrame.
+    :param sep: String separator used to join the column level names
+    :returns: df
+    """
+
+    check("sep", sep, [str])
+
+    # if already single-level, just return the DataFrame
+    if not isinstance(df.columns.values[0], tuple):
+        return df
+
+    df.columns = [
+        sep.join([str(el) for el in tup if str(el) != ""])
+        for tup in df.columns.values
+    ]
+    return df
+
+
 def check(varname: str, value, expected_types: list):
     """
     One-liner syntactic sugar for checking types.


Test patch:
diff --git a/tests/test_functions.py b/tests/test_functions.py
--- a/tests/test_functions.py
+++ b/tests/test_functions.py
@@ -57,6 +57,32 @@ def multiindex_dataframe():
     return df
 
 
+@pytest.fixture
+def multiindex_with_missing_dataframe():
+    data = {
+        ("a", ""): [1, 2, 3],
+        ("", "Normal  Distribution"): [1, 2, 3],
+        ("decorated-elephant", "r.i.p-rhino :'("): [1, 2, 3],
+    }
+    df = pd.DataFrame(data)
+    return df
+
+
+@pytest.fixture
+def multiindex_with_missing_3level_dataframe():
+    data = {
+        ("a", "", ""): [1, 2, 3],
+        ("", "Normal  Distribution", "Hypercuboid (???)"): [1, 2, 3],
+        ("decorated-elephant", "r.i.p-rhino :'(", "deadly__flamingo"): [
+            1,
+            2,
+            3,
+        ],
+    }
+    df = pd.DataFrame(data)
+    return df
+
+
 def test_clean_names_functional(dataframe):
     df = clean_names(dataframe)
     expected_columns = [
@@ -656,7 +682,7 @@ def test_add_column_raise_error(dataframe):
         dataframe.add_column("cities", 1)
 
 
-def test_add_column_iterator_repeat(dataframe):
+def test_add_column_iterator_repeat_subtraction(dataframe):
     df = dataframe.add_column("city_pop", dataframe.a - dataframe.a)
     assert df.city_pop.sum() == 0
     assert df.city_pop.iloc[0] == 0
@@ -665,7 +691,7 @@ def test_add_column_iterator_repeat(dataframe):
 def test_row_to_names(dataframe):
     df = dataframe.row_to_names(2)
     assert df.columns[0] == 3
-    assert df.columns[1] == 3.2346125
+    assert df.columns[1] == 3.234_612_5
     assert df.columns[2] == 3
     assert df.columns[3] == "lion"
     assert df.columns[4] == "Basel"
@@ -674,7 +700,7 @@ def test_row_to_names(dataframe):
 def test_row_to_names_delete_this_row(dataframe):
     df = dataframe.row_to_names(2, remove_row=True)
     assert df.iloc[2, 0] == 1
-    assert df.iloc[2, 1] == 1.23452345
+    assert df.iloc[2, 1] == 1.234_523_45
     assert df.iloc[2, 2] == 1
     assert df.iloc[2, 3] == "rabbit"
     assert df.iloc[2, 4] == "Cambridge"
@@ -683,7 +709,7 @@ def test_row_to_names_delete_this_row(dataframe):
 def test_row_to_names_delete_above(dataframe):
     df = dataframe.row_to_names(2, remove_rows_above=True)
     assert df.iloc[0, 0] == 3
-    assert df.iloc[0, 1] == 3.2346125
+    assert df.iloc[0, 1] == 3.234_612_5
     assert df.iloc[0, 2] == 3
     assert df.iloc[0, 3] == "lion"
     assert df.iloc[0, 4] == "Basel"
@@ -729,3 +755,63 @@ def test_min_max_old_min_max_errors(dataframe):
 def test_min_max_new_min_max_errors(dataframe):
     with pytest.raises(ValueError):
         df = dataframe.min_max_scale(col_name="a", new_min=10, new_max=0)
+
+
+def test_collapse_levels_sanity(multiindex_with_missing_dataframe):
+    with pytest.raises(TypeError):
+        multiindex_with_missing_dataframe.collapse_levels(sep=3)
+
+
+def test_collapse_levels_non_multilevel(multiindex_with_missing_dataframe):
+    # an already single-level DataFrame is not distorted
+    pd.testing.assert_frame_equal(
+        multiindex_with_missing_dataframe.copy().collapse_levels(),
+        multiindex_with_missing_dataframe.collapse_levels().collapse_levels(),
+    )
+
+
+def test_collapse_levels_functionality_2level(
+    multiindex_with_missing_dataframe
+):
+
+    assert all(
+        multiindex_with_missing_dataframe.copy()
+        .collapse_levels()
+        .columns.values
+        == ["a", "Normal  Distribution", "decorated-elephant_r.i.p-rhino :'("]
+    )
+    assert all(
+        multiindex_with_missing_dataframe.copy()
+        .collapse_levels(sep="AsDf")
+        .columns.values
+        == [
+            "a",
+            "Normal  Distribution",
+            "decorated-elephantAsDfr.i.p-rhino :'(",
+        ]
+    )
+
+
+def test_collapse_levels_functionality_3level(
+    multiindex_with_missing_3level_dataframe
+):
+    assert all(
+        multiindex_with_missing_3level_dataframe.copy()
+        .collapse_levels()
+        .columns.values
+        == [
+            "a",
+            "Normal  Distribution_Hypercuboid (???)",
+            "decorated-elephant_r.i.p-rhino :'(_deadly__flamingo",
+        ]
+    )
+    assert all(
+        multiindex_with_missing_3level_dataframe.copy()
+        .collapse_levels(sep="AsDf")
+        .columns.values
+        == [
+            "a",
+            "Normal  DistributionAsDfHypercuboid (???)",
+            "decorated-elephantAsDfr.i.p-rhino :'(AsDfdeadly__flamingo",
+        ]
+    )