Repository: ericmjl/pyjanitor (python)

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
```

Pull request description:
[ENH] Groupby_agg function
# PR Description

Please describe the changes proposed in the pull request:

- Updates groupby_agg to cater for add_count method
- Covers aggregation for null entries, pending when Pyjanitor's minimum version is 1.1, which natively supports keeping null values

<!-- Doing so provides maintainers with context on what the PR is, and can help us more effectively review your PR. -->

<!-- Please also identify below which issue that has been raised that you are going to close. -->

**This PR resolves #623 **

<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [x] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.rst`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.rst` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

## Quick Check

To do a very quick check that everything is correct, follow these steps below:

- [ ] Run the command `make check` from pyjanitor's top-level directory. This will automatically run:
    - black formatting
    - flake8 checking
    - running the test suite
    - docs build

Once done, please check off the check-box above.

If `make check` does not work for you, you can execute the commands listed in the Makefile individually.

## Code Changes

<!-- If you have not made code changes, please feel free to delete this section. -->

If you are adding code changes, please ensure the following:

- [ ] Ensure that you have added tests.
- [ ] Run all tests (`$ pytest .`) locally on your machine.
    - [ ] Check to ensure that test coverage covers the lines of code that you have added.
    - [ ] Ensure that all tests pass.

## Documentation Changes

<!-- If you have not made documentation changes, please feel free to delete this section. -->

If you are adding documentation changes, please ensure the following:

- [ ] Build the docs locally.
- [ ] View the docs to check that it renders correctly.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/.travis.yml b/.travis.yml
index a207fd8..6bfd474 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -11,7 +11,7 @@ install:
   # We do this conditionally because it saves us some downloading if the
   # version is the same.
   - wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
-  - bash miniconda.sh -b -p $HOME/miniconda
+  - bash miniconda.sh -b -p $HOME/miniconda && rm miniconda.sh
   - export PATH="$HOME/miniconda/bin:$PATH"
   - hash -r
   - conda config --set always_yes yes --set changeps1 no
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 235d797..43b9566 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,5 +1,6 @@
 new version (on deck)
 =====================
+- [ENH] Updated groupby_agg function to account for null entries in the ``by`` argument. @samukweku
 
 
 v0.20.8
diff --git a/examples/notebooks/groupby_agg.ipynb b/examples/notebooks/groupby_agg.ipynb
index 2e9a589..33ba403 100644
--- a/examples/notebooks/groupby_agg.ipynb
+++ b/examples/notebooks/groupby_agg.ipynb
@@ -18,7 +18,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This notebook serves to show a brief and simple example of how to use the groupby_agg method from pyjanitor's general functions submodule.\n",
+    "This notebook serves to show how to use the groupby_agg method from pyjanitor's general functions submodule.\n",
     "\n",
     "The groupby_agg method allows us to add the result of an aggregation from a grouping, as a new column, back to the dataframe. \n",
     "\n",
@@ -32,93 +32,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "# load modules\n",
     "import pandas as pd\n",
+    "import numpy as np\n",
     "from janitor import groupby_agg"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
+     "output_type": "execute_result",
      "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>item</th>\n",
-       "      <th>MRP</th>\n",
-       "      <th>number_sold</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>220</td>\n",
-       "      <td>100</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>450</td>\n",
-       "      <td>40</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>bag</td>\n",
-       "      <td>320</td>\n",
-       "      <td>56</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>200</td>\n",
-       "      <td>38</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>bag</td>\n",
-       "      <td>305</td>\n",
-       "      <td>25</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "   item  MRP  number_sold\n",
-       "0  shoe  220          100\n",
-       "1  shoe  450           40\n",
-       "2   bag  320           56\n",
-       "3  shoe  200           38\n",
-       "4   bag  305           25"
-      ]
+      "text/plain": "   item  MRP  number_sold\n0  shoe  220          100\n1  shoe  450           40\n2   bag  320           56\n3  shoe  200           38\n4   bag  305           25",
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>item</th>\n      <th>MRP</th>\n      <th>number_sold</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>shoe</td>\n      <td>220</td>\n      <td>100</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>shoe</td>\n      <td>450</td>\n      <td>40</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>bag</td>\n      <td>320</td>\n      <td>56</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>shoe</td>\n      <td>200</td>\n      <td>38</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>bag</td>\n      <td>305</td>\n      <td>25</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
      },
-     "execution_count": 3,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 2
     }
    ],
    "source": [
@@ -142,88 +78,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
+     "output_type": "execute_result",
      "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>item</th>\n",
-       "      <th>MRP</th>\n",
-       "      <th>number_sold</th>\n",
-       "      <th>Avg_MRP</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>220</td>\n",
-       "      <td>100</td>\n",
-       "      <td>290.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>450</td>\n",
-       "      <td>40</td>\n",
-       "      <td>290.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>bag</td>\n",
-       "      <td>320</td>\n",
-       "      <td>56</td>\n",
-       "      <td>312.5</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>shoe</td>\n",
-       "      <td>200</td>\n",
-       "      <td>38</td>\n",
-       "      <td>290.0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>bag</td>\n",
-       "      <td>305</td>\n",
-       "      <td>25</td>\n",
-       "      <td>312.5</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "   item  MRP  number_sold  Avg_MRP\n",
-       "0  shoe  220          100    290.0\n",
-       "1  shoe  450           40    290.0\n",
-       "2   bag  320           56    312.5\n",
-       "3  shoe  200           38    290.0\n",
-       "4   bag  305           25    312.5"
-      ]
+      "text/plain": "   item  MRP  number_sold  Avg_MRP\n0  shoe  220          100    290.0\n1  shoe  450           40    290.0\n2   bag  320           56    312.5\n3  shoe  200           38    290.0\n4   bag  305           25    312.5",
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>item</th>\n      <th>MRP</th>\n      <th>number_sold</th>\n      <th>Avg_MRP</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>shoe</td>\n      <td>220</td>\n      <td>100</td>\n      <td>290.0</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>shoe</td>\n      <td>450</td>\n      <td>40</td>\n      <td>290.0</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>bag</td>\n      <td>320</td>\n      <td>56</td>\n      <td>312.5</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>shoe</td>\n      <td>200</td>\n      <td>38</td>\n      <td>290.0</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>bag</td>\n      <td>305</td>\n      <td>25</td>\n      <td>312.5</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
      },
-     "execution_count": 4,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 3
     }
    ],
    "source": [
@@ -236,6 +101,104 @@
     "\n",
     "df"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Null cells are retained as well : "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": "    name   type  num  nulls\n0  black  chair    4    1.0\n1  black  chair    5    1.0\n2  black   sofa   12    NaN\n3    red   sofa    4    NaN\n4    red  plate    3    3.0",
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>name</th>\n      <th>type</th>\n      <th>num</th>\n      <th>nulls</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>4</td>\n      <td>1.0</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>5</td>\n      <td>1.0</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>black</td>\n      <td>sofa</td>\n      <td>12</td>\n      <td>NaN</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>red</td>\n      <td>sofa</td>\n      <td>4</td>\n      <td>NaN</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>red</td>\n      <td>plate</td>\n      <td>3</td>\n      <td>3.0</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
+     },
+     "metadata": {},
+     "execution_count": 4
+    }
+   ],
+   "source": [
+    "df = pd.DataFrame(\n",
+    "        {\n",
+    "            \"name\": (\"black\", \"black\", \"black\", \"red\", \"red\"),\n",
+    "            \"type\": (\"chair\", \"chair\", \"sofa\", \"sofa\", \"plate\"),\n",
+    "            \"num\": (4, 5, 12, 4, 3),\n",
+    "            \"nulls\": (1, 1, np.nan, np.nan, 3),\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": "    name   type  num  nulls  counter\n0  black  chair    4    1.0        2\n1  black  chair    5    1.0        2\n2  black   sofa   12    NaN        2\n3    red   sofa    4    NaN        2\n4    red  plate    3    3.0        1",
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>name</th>\n      <th>type</th>\n      <th>num</th>\n      <th>nulls</th>\n      <th>counter</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>4</td>\n      <td>1.0</td>\n      <td>2</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>5</td>\n      <td>1.0</td>\n      <td>2</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>black</td>\n      <td>sofa</td>\n      <td>12</td>\n      <td>NaN</td>\n      <td>2</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>red</td>\n      <td>sofa</td>\n      <td>4</td>\n      <td>NaN</td>\n      <td>2</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>red</td>\n      <td>plate</td>\n      <td>3</td>\n      <td>3.0</td>\n      <td>1</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
+     },
+     "metadata": {},
+     "execution_count": 5
+    }
+   ],
+   "source": [
+    "filtered_df = df.groupby_agg(by=['nulls'],\n",
+    "                             agg='size',\n",
+    "                             agg_column_name='type',\n",
+    "                             new_column_name='counter'\n",
+    "                             )\n",
+    "filtered_df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The groupby_agg method can be extended for different purposes. One of these is groupwise filtering, where only groups that meet a condition are retained. Let's explore this with an example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": "    name   type  num  nulls  counter\n0  black  chair    4    1.0        2\n1  black  chair    5    1.0        2",
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>name</th>\n      <th>type</th>\n      <th>num</th>\n      <th>nulls</th>\n      <th>counter</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>4</td>\n      <td>1.0</td>\n      <td>2</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>black</td>\n      <td>chair</td>\n      <td>5</td>\n      <td>1.0</td>\n      <td>2</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
+     },
+     "metadata": {},
+     "execution_count": 6
+    }
+   ],
+   "source": [
+    "filtered_df = (df.groupby_agg(by=['name', 'type'],\n",
+    "                             agg='size',\n",
+    "                             agg_column_name='type',\n",
+    "                             new_column_name='counter'\n",
+    "                             )\n",
+    "                .query('counter > 1'))\n",
+    "filtered_df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -259,4 +222,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/janitor/functions.py b/janitor/functions.py
index d7487b3..b4668a8 100644
--- 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
index 9c1f3ca..e3c74c9 100644
--- 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,
+    )