Repository: pyjanitor-devs/pyjanitor

Issue:
groupby_agg doesn't handle multiple columns well
It occurred to me that the current implementation of `groupby_agg` only handles the case where the DataFrame has n+1 columns, where n is the number of levels of groupings. 

```python
In [1]: 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],
   ...:         }
   ...:     )
   ...:

In [2]: df.groupby_agg(by=["date", "user_id"],
   ...:         new_column_name="date_average",
   ...:         agg_column_name="values",
   ...:         agg=pd.np.count_nonzero,
   ...:     ) # works
In [4]: df.groupby_agg(by=["date"],
   ...:         new_column_name="date_average",
   ...:         agg_column_name="values",
   ...:         agg=pd.np.count_nonzero,
   ...:     ) # Throws ValueError: Wrong number of items passed 2, placement implies 1
```

I can make a quick PR to fix it, though it might change the API a bit since the user should specify which column(s) to aggregate.


Discussion:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -3015,10 +3015,10 @@ def to_datetime(df: pd.DataFrame, column_name, **kwargs) -> pd.DataFrame:
 @deprecated_alias(new_column="new_column_name", agg_column="agg_column_name")
 def groupby_agg(
     df: pd.DataFrame,
-    by: str,
-    new_column_name,
-    agg_column_name,
-    agg: Union[Callable, str, List, Dict],
+    by: Union[List, str],
+    new_column_name: str,
+    agg_column_name: str,
+    agg: Union[Callable, str],
 ) -> pd.DataFrame:
     """
     Shortcut for assigning a groupby-transform to a new column.
@@ -3038,9 +3038,10 @@ def groupby_agg(
         import pandas as pd
         import janitor
         df = pd.DataFrame(...).groupby_agg(df,
-                                           by='col1',
+                                           by='group',
                                            agg='mean',
-                                           new_column_name='col1_mean')
+                                           agg_column_name="col1"
+                                           new_column_name='col1_mean_by_group')
 
     :param df: A pandas DataFrame.
     :param by: Column(s) to groupby on, either a `str` or
@@ -3051,7 +3052,7 @@ def groupby_agg(
     :param axis: Split along rows (0) or columns (1).
     :returns: A pandas DataFrame.
     """
-    new_col = df.groupby(by).transform(agg)
+    new_col = df.groupby(by)[agg_column_name].transform(agg)
     df_new = df.assign(**{new_column_name: new_col})
     return df_new
 


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
@@ -59,3 +59,36 @@ def test_groupby_agg_multi():
     expected_agg = np.array([1, 1, 1, 1, 1, 1])
 
     np.testing.assert_equal(df_new["date_average"], expected_agg)
+
+
+@pytest.mark.functions
+def test_groupby_agg_multi_column():
+    """
+    Test for the case when we want to groupby one column and agg
+    on another, while leaving out other columns.
+    """
+
+    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"],
+        new_column_name="values_avg",
+        agg_column_name="values",
+        agg="mean",
+    )
+
+    expected_agg = np.array([1.5, 1.5, 3.5, 3.5, 5.5, 5.5])
+    np.testing.assert_equal(df_new["values_avg"], expected_agg)