ENH: Add dropna in groupby to allow NaN in keys

doc/source/whatsnew/v1.0.0.rst

@@ -549,6 +549,7 @@ Other API changes
   Supplying anything else than ``how`` to ``**kwargs`` raised a ``TypeError`` previously (:issue:`29388`)
 - When testing pandas, the new minimum required version of pytest is 5.0.1 (:issue:`29664`)
 - :meth:`Series.str.__iter__` was deprecated and will be removed in future releases (:issue:`28277`).
+- :meth:`DataFrame.groupby` and :meth:`Series.groupby` have gained ``dropna`` argument in order to allow ``NaN`` values in group keys (:issue:`3729`)
 
 
 .. _whatsnew_1000.api.documentation:

pandas/core/algorithms.py

@@ -596,7 +596,11 @@ def _factorize_array(
 )
 @Appender(_shared_docs["factorize"])
 def factorize(
-    values, sort: bool = False, na_sentinel: int = -1, size_hint: Optional[int] = None
+    values,
+    sort: bool = False,
+    na_sentinel: int = -1,
+    size_hint: Optional[int] = None,
+    dropna: Optional[bool] = None,
 ) -> Tuple[np.ndarray, Union[np.ndarray, ABCIndex]]:
     # Implementation notes: This method is responsible for 3 things
     # 1.) coercing data to array-like (ndarray, Index, extension array)
@@ -630,6 +634,9 @@ def factorize(
         uniques, codes = safe_sort(
             uniques, codes, na_sentinel=na_sentinel, assume_unique=True, verify=False
         )
+    if dropna is False and (codes == na_sentinel).any():
+        uniques = np.append(uniques, [np.nan])
+        codes = np.where(codes == na_sentinel, len(uniques) - 1, codes)
 
     uniques = _reconstruct_data(uniques, dtype, original)
 

pandas/core/generic.py

@@ -7311,6 +7311,7 @@ def groupby(
         group_keys: bool_t = True,
         squeeze: bool_t = False,
         observed: bool_t = False,
+        dropna: Optional[bool_t] = None,
     ):
         """
         Group DataFrame or Series using a mapper or by a Series of columns.
@@ -7355,6 +7356,12 @@ def groupby(
             If False: show all values for categorical groupers.
 
             .. versionadded:: 0.23.0
+        dropna : bool or None, default None
+            If None or True, and if group keys contain NaN values, NaN values together
+            with row/column will be dropped.
+            If False, NaN values will also be treated as the key in groups
+
+            .. versionadded:: 1.0.0
 
         Returns
         -------
@@ -7433,6 +7440,7 @@ def groupby(
             group_keys=group_keys,
             squeeze=squeeze,
             observed=observed,
+            dropna=dropna,
         )
 
     def asfreq(

pandas/core/groupby/groupby.py

@@ -373,6 +373,7 @@ def __init__(
         squeeze: bool = False,
         observed: bool = False,
         mutated: bool = False,
+        dropna: Optional[bool] = None,
     ):
 
         self._selection = selection
@@ -396,6 +397,8 @@ def __init__(
         self.observed = observed
         self.mutated = mutated
 
+        self.dropna = dropna if dropna is not None else True
+
         if grouper is None:
             from pandas.core.groupby.grouper import get_grouper
 
@@ -407,6 +410,7 @@ def __init__(
                 sort=sort,
                 observed=observed,
                 mutated=self.mutated,
+                dropna=self.dropna,
             )
 
         self.obj = obj
@@ -2543,6 +2547,7 @@ def get_groupby(
     squeeze: bool = False,
     observed: bool = False,
     mutated: bool = False,
+    dropna: Optional[bool] = None,
 ) -> GroupBy:
 
     klass: Type[GroupBy]
@@ -2557,6 +2562,8 @@ def get_groupby(
     else:
         raise TypeError(f"invalid type: {obj}")
 
+    dropna = dropna if dropna is not None else True
+
     return klass(
         obj=obj,
         keys=by,
@@ -2571,4 +2578,5 @@ def get_groupby(
         squeeze=squeeze,
         observed=observed,
         mutated=mutated,
+        dropna=dropna,
     )

pandas/core/groupby/grouper.py

@@ -100,7 +100,9 @@ def __new__(cls, *args, **kwargs):
             cls = TimeGrouper
         return super().__new__(cls)
 
-    def __init__(self, key=None, level=None, freq=None, axis=0, sort=False):
+    def __init__(
+        self, key=None, level=None, freq=None, axis=0, sort=False, dropna=None
+    ):
         self.key = key
         self.level = level
         self.freq = freq
@@ -112,6 +114,7 @@ def __init__(self, key=None, level=None, freq=None, axis=0, sort=False):
         self.indexer = None
         self.binner = None
         self._grouper = None
+        self.dropna = dropna if dropna is not None else True
 
     @property
     def ax(self):
@@ -138,6 +141,7 @@ def _get_grouper(self, obj, validate: bool = True):
             level=self.level,
             sort=self.sort,
             validate=validate,
+            dropna=self.dropna,
         )
         return self.binner, self.grouper, self.obj
 
@@ -250,6 +254,7 @@ def __init__(
         sort: bool = True,
         observed: bool = False,
         in_axis: bool = False,
+        dropna: Optional[bool] = None,
     ):
         self.name = name
         self.level = level
@@ -261,6 +266,8 @@ def __init__(
         self.observed = observed
         self.in_axis = in_axis
 
+        self.dropna = dropna if dropna is not None else True
+
         # right place for this?
         if isinstance(grouper, (Series, Index)) and name is None:
             self.name = grouper.name
@@ -413,7 +420,9 @@ def _make_codes(self) -> None:
                 codes = self.grouper.codes_info
                 uniques = self.grouper.result_index
             else:
-                codes, uniques = algorithms.factorize(self.grouper, sort=self.sort)
+                codes, uniques = algorithms.factorize(
+                    self.grouper, sort=self.sort, dropna=self.dropna
+                )
                 uniques = Index(uniques, name=self.name)
             self._codes = codes
             self._group_index = uniques
@@ -432,6 +441,7 @@ def get_grouper(
     observed: bool = False,
     mutated: bool = False,
     validate: bool = True,
+    dropna: Optional[bool] = None,
 ) -> "Tuple[ops.BaseGrouper, List[Hashable], FrameOrSeries]":
     """
     Create and return a BaseGrouper, which is an internal
@@ -621,6 +631,7 @@ def is_in_obj(gpr) -> bool:
                 sort=sort,
                 observed=observed,
                 in_axis=in_axis,
+                dropna=dropna,
             )
             if not isinstance(gpr, Grouping)
             else gpr

pandas/tests/groupby/test_groupby.py

@@ -2025,3 +2025,137 @@ def test_groupby_crash_on_nunique(axis):
         expected = expected.T
 
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "dropna, tuples, outputs",
+    [
+        (
+            None,
+            [["A", "B"], ["B", "A"]],
+            {"c": [13.0, 123.23], "d": [13.0, 123.0], "e": [13.0, 1.0]},
+        ),
+        (
+            True,
+            [["A", "B"], ["B", "A"]],
+            {"c": [13.0, 123.23], "d": [13.0, 123.0], "e": [13.0, 1.0]},
+        ),
+        (
+            False,
+            [["A", "B"], ["A", np.nan], ["B", "A"]],
+            {
+                "c": [13.0, 12.3, 123.23],
+                "d": [13.0, 233.0, 123.0],
+                "e": [13.0, 12.0, 1.0],
+            },
+        ),
+    ],
+)
+def test_groupby_dropna_multi_index_dataframe(dropna, tuples, outputs):
+    # GH 3729
+    df_list = [
+        ["A", "B", 12, 12, 12],
+        ["A", None, 12.3, 233.0, 12],
+        ["B", "A", 123.23, 123, 1],
+        ["A", "B", 1, 1, 1.0],
+    ]
+    df = pd.DataFrame(df_list, columns=["a", "b", "c", "d", "e"])
+    grouped = df.groupby(["a", "b"], dropna=dropna).sum()
+
+    mi = pd.MultiIndex.from_tuples(tuples, names=list("ab"))
+    expected = pd.DataFrame(outputs, index=mi)
+
+    tm.assert_frame_equal(grouped, expected, check_index_type=False)
+
+
+@pytest.mark.parametrize(
+    "dropna, idx, outputs",
+    [
+        (None, ["A", "B"], {"b": [123.23, 13.0], "c": [123.0, 13.0], "d": [1.0, 13.0]}),
+        (True, ["A", "B"], {"b": [123.23, 13.0], "c": [123.0, 13.0], "d": [1.0, 13.0]}),
+        (
+            False,
+            ["A", "B", np.nan],
+            {
+                "b": [123.23, 13.0, 12.3],
+                "c": [123.0, 13.0, 233.0],
+                "d": [1.0, 13.0, 12.0],
+            },
+        ),
+    ],
+)
+def test_groupby_dropna_normal_index_dataframe(dropna, idx, outputs):
+    # GH 3729
+    df_list = [
+        ["B", 12, 12, 12],
+        [None, 12.3, 233.0, 12],
+        ["A", 123.23, 123, 1],
+        ["B", 1, 1, 1.0],
+    ]
+    df = pd.DataFrame(df_list, columns=["a", "b", "c", "d"])
+    grouped = df.groupby("a", dropna=dropna).sum()
+
+    expected = pd.DataFrame(outputs, index=pd.Index(idx, dtype="object", name="a"))
+
+    tm.assert_frame_equal(grouped, expected, check_index_type=False)
+
+
+@pytest.mark.parametrize(
+    "dropna, idx, expected",
+    [
+        (None, ["a", "a", "b", np.nan], pd.Series([3, 3], index=["a", "b"])),
+        (True, ["a", "a", "b", np.nan], pd.Series([3, 3], index=["a", "b"])),
+        (
+            False,
+            ["a", "a", "b", np.nan],
+            pd.Series([3, 3, 3], index=["a", "b", np.nan]),
+        ),
+    ],
+)
+def test_groupby_dropna_series(dropna, idx, expected):
+    ser = pd.Series([1, 2, 3, 3], index=idx)
+
+    result = ser.groupby(level=0, dropna=dropna).sum()
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "dropna, tuples, outputs",
+    [
+        (
+            None,
+            [["A", "B"], ["B", "A"]],
+            {"c": [13.0, 123.23], "d": [12.0, 123.0], "e": [1.0, 1.0]},
+        ),
+        (
+            True,
+            [["A", "B"], ["B", "A"]],
+            {"c": [13.0, 123.23], "d": [12.0, 123.0], "e": [1.0, 1.0]},
+        ),
+        (
+            False,
+            [["A", "B"], ["A", np.nan], ["B", "A"]],
+            {
+                "c": [13.0, 12.3, 123.23],
+                "d": [12.0, 233.0, 123.0],
+                "e": [1.0, 12.0, 1.0],
+            },
+        ),
+    ],
+)
+def test_groupby_dropna_multi_index_dataframe_agg(dropna, tuples, outputs):
+    # GH 3729
+    df_list = [
+        ["A", "B", 12, 12, 12],
+        ["A", None, 12.3, 233.0, 12],
+        ["B", "A", 123.23, 123, 1],
+        ["A", "B", 1, 1, 1.0],
+    ]
+    df = pd.DataFrame(df_list, columns=["a", "b", "c", "d", "e"])
+    agg_dict = {"c": sum, "d": max, "e": "min"}
+    grouped = df.groupby(["a", "b"], dropna=dropna).agg(agg_dict)
+
+    mi = pd.MultiIndex.from_tuples(tuples, names=list("ab"))
+    expected = pd.DataFrame(outputs, index=mi)
+
+    tm.assert_frame_equal(grouped, expected, check_index_type=False)