Repository: pyjanitor-devs/pyjanitor

Issue:
Add add_count and add_total functions
I know there is functionality with groupby_agg here, so I have been going back and forth on this one.  But like the tidyverse there are helper functions the are more direct like tally or add_count, where you could use summarize() and count() or n() together to get the same output.  

I am not deeply invested in this suggestion, so if you think that groupby_agg can work here alone, I am fine with that.  

Here are the two proposed functions I could possibly add:

add_count function
```
@pf.register_dataframe_method
def add_count(df: pd.DataFrame(), count_column: str, new_column_name: str, include_na: bool):
    """Add countby column on target column"""  
    
    if include_na == True:
        df[new_column_name] = df.groupby([count_column])[count_column].transform('count')
        num = df[new_column_name].isna().sum()
        df[new_column_name].fillna(num, inplace= True)
        return df
    
    else:
        df[new_column_name] = df.groupby([count_column])[count_column].transform('count')
        return df

```

add_total function
```
@pf.register_dataframe_method
def add_total(df: pd.DataFrame(), sum_columns: list, new_column_name: str, axis: int, skipna: bool):
    """Add countby column on target column"""  
    
    if axis == 1:
        df[new_column_name] = df[sum_columns].sum(axis= axis, skipna = skipna )
        return df

    else:
        df[new_column_name] = df[sum_columns].sum(axis= axis, skipna = skipna ).sum()
        return df
```


Discussion:
@jcvall-olympus this is a great suggestion! Aliases that make things easier are always good. :smile:. 

Would you like to work on this? If so, I’ll assign this one to you! In terms of what to put inside the PR, I’d suggest doing one function at a time, e.g. let’s add in a `tally` or `add_count` function first.
I plan on upload my functions.  I hope to do them in a few weeks.  Sorry for the delay.

<sub>Sent with <a href="http://githawk.com">GitHawk</a></sub>
stil working on this @jcvall ? if u are not, I would love to give it a go

Go ahead Sam.  Thanks for taking this one on. 

<sub>Sent with <a href="http://githawk.com">GitHawk</a></sub>
Pandas 1.1.0 has some updates that I think will be useful for ``` groupby_agg ```  and by extension ```count``` - a parameter for including or excluding null values in groupby has been added. Just waiting for the release and I will push off on some functions for this issue

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -3244,9 +3244,30 @@ def groupby_agg(
     :param axis: Split along rows (0) or columns (1).
     :returns: A pandas DataFrame.
     """
-    new_col = df.groupby(by)[agg_column_name].transform(agg)
-    df_new = df.assign(**{new_column_name: new_col})
-    return df_new
+
+    # convert to list
+    # needed when creating a mapping through the iteration
+    if isinstance(by, str):
+        by = [by]
+    # this is a temporary measure, till the minimum Pandas version is 1.1,
+    # which supports null values in the group by
+    # If any of the grouping columns has null values, we temporarily
+    # replace the values with some outrageous value, that should not exist
+    # in the column. Also, the hasnans property is significantly faster than
+    # .isnull().any()
+    if any(df[col].hasnans for col in by):
+
+        mapping = {
+            column: ".*^%s1ho1go1logoban?*&-|/\\gos1he()#_" for column in by
+        }
+
+        df[new_column_name] = (
+            df.fillna(mapping).groupby(by)[agg_column_name].transform(agg)
+        )
+
+    else:
+        df[new_column_name] = df.groupby(by)[agg_column_name].transform(agg)
+    return df
 
 
 @pf.register_dataframe_accessor("data_description")


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,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from pandas.testing import assert_frame_equal
 
 
 @pytest.mark.functions
@@ -24,7 +25,7 @@ def test_groupby_agg():
         by="date",
         new_column_name="date_average",
         agg_column_name="values",
-        agg=pd.np.mean,
+        agg=np.mean,
     )
     assert df.shape[0] == df_new.shape[0]
     assert "date_average" in df_new.columns
@@ -53,7 +54,7 @@ def test_groupby_agg_multi():
         by=["date", "user_id"],
         new_column_name="date_average",
         agg_column_name="values",
-        agg=pd.np.count_nonzero,
+        agg=np.count_nonzero,
     )
 
     expected_agg = np.array([1, 1, 1, 1, 1, 1])
@@ -92,3 +93,35 @@ def test_groupby_agg_multi_column():
 
     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)
+
+
+@pytest.fixture
+def df():
+    return pd.DataFrame(
+        {
+            "name": ("black", "black", "black", "red", "red"),
+            "type": ("chair", "chair", "sofa", "sofa", "plate"),
+            "num": (4, 5, 12, 4, 3),
+            "nulls": (1, 1, np.nan, np.nan, 3),
+        }
+    )
+
+
+def test_groupby_agg_null_columns(df):
+    """
+    Test that aggregation on null column retains the null entries and returns
+    the right aggregated value.
+    """
+    expected = df.copy()
+    expected["num_count"] = (
+        expected.fillna("rar").groupby(["nulls"])["type"].transform("size")
+    )
+    assert_frame_equal(
+        df.groupby_agg(
+            by="nulls",
+            new_column_name="num_count",
+            agg_column_name="num",
+            agg="size",
+        ),
+        expected,
+    )