Repository: pyjanitor-devs/pyjanitor

Issue:
groupby_agg doesn't support multi-index groupby objects
Elaborating a bit on the [existing unit test](https://github.com/ericmjl/pyjanitor/blob/17f8a4b30fe8eb610a9955af43057e635f56e04c/tests/functions/test_groupby_agg.py#L6) for this function, I want to add a simple column and group on both

``` python
df = pd.DataFrame(
    {
        "date": [
            "20190101",
            "20190101",
            "20190102",
            "20190102",
            "20190304",
            "20190304",
        ],
        "user_id": [1, 1, 1, 2, 2, 2],
        "values": [1, 1, 2, 2, 3, 3],
    }
)

df_new = df.groupby_agg(
    by=["date", "user_id"],
    new_column="date_average",
    agg_column="values",
    agg=pd.np.mean,
)
```

Running yields

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-40f7fe6f4a1a> in <module>
     19     new_column="date_average",
     20     agg_column="values",
---> 21     agg=pd.np.mean,
     22 )
     23 assert df.shape[0] == df_new.shape[0]

...
...

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_indexer()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.lookup()

TypeError: unhashable type: 'list'
```

I think we just want to change [this section](https://github.com/ericmjl/pyjanitor/blob/17f8a4b30fe8eb610a9955af43057e635f56e04c/janitor/functions.py#L2721) to account for one or multiple `groupby` conditions


Discussion:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -2721,7 +2721,8 @@ def groupby_agg(
                                            new_column='col1_mean')
 
     :param df: A pandas DataFrame.
-    :param by: Column to groupby on.
+    :param by: Column(s) to groupby on, either a `str` or
+               a `list` of `str`
     :param new_column: Name of the aggregation output column.
     :param agg_column: Name of the column to aggregate over.
     :param agg: How to aggregate.
@@ -2734,7 +2735,12 @@ def groupby_agg(
         .agg(agg, axis=axis)
         .reset_index()
         .rename(columns={agg_column: new_column})
-    )[[by, new_column]]
+    )
+
+    if isinstance(by, list) or isinstance(by, tuple):
+        df_grp = df_grp[[*by, new_column]]
+    else:
+        df_grp = df_grp[[by, new_column]]
 
     df = df.merge(df_grp, on=by)
 


Test patch:
diff --git a/tests/functions/test_groupby_agg.py b/tests/functions/test_groupby_agg.py
--- a/tests/functions/test_groupby_agg.py
+++ b/tests/functions/test_groupby_agg.py
@@ -1,4 +1,5 @@
 import pandas as pd
+import numpy as np
 import pytest
 
 
@@ -28,3 +29,33 @@ def test_groupby_agg():
     assert df.shape[0] == df_new.shape[0]
     assert "date_average" in df_new.columns
     assert df_new["date_average"].iloc[0] == 1
+
+
+@pytest.mark.functions
+def test_groupby_agg_multi():
+
+    df = pd.DataFrame(
+        {
+            "date": [
+                "20190101",
+                "20190101",
+                "20190102",
+                "20190102",
+                "20190304",
+                "20190304",
+            ],
+            "user_id": [1, 2, 1, 2, 1, 2],
+            "values": [1, 2, 3, 4, 5, 6],
+        }
+    )
+
+    df_new = df.groupby_agg(
+        by=["date", "user_id"],
+        new_column="date_average",
+        agg_column="values",
+        agg=pd.np.count_nonzero,
+    )
+
+    expected_agg = np.array([1, 1, 1, 1, 1, 1])
+
+    np.testing.assert_equal(df_new["date_average"], expected_agg)