ENH: Add dropna in groupby to allow NaN in keys

doc/source/whatsnew/v1.0.0.rst

@@ -191,6 +191,16 @@ method on a :func:`pandas.api.indexers.BaseIndexer` subclass that will generate
 indices used for each window during the rolling aggregation. For more details and example usage, see
 the :ref:`custom window rolling documentation <stats.custom_rolling_window>`
 
+.. _whatsnew_1000.groupby_key:
+
+Allow NaN in groupby key
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+We've added a ``dropna`` keyword to :meth:`DataFrame.groupby` and :meth:`Series.groupby` in order to
+allow ``NaN`` values in group keys. Users can define ``dropna`` to ``False`` if they want to include
+``NaN`` values in groupby keys. The default is set to ``True`` for ``dropna`` to keep backwards
+compatibility (:issue:`3729`)
+
 .. _whatsnew_1000.enhancements.other:
 
 Other enhancements

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: bool = True,
 ) -> 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 not dropna 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/frame.py

@@ -1988,7 +1988,7 @@ def to_feather(self, path):
     @Substitution(klass="DataFrame")
     @Appender(_shared_docs["to_markdown"])
     def to_markdown(
-        self, buf: Optional[IO[str]] = None, mode: Optional[str] = None, **kwargs,
+        self, buf: Optional[IO[str]] = None, mode: Optional[str] = None, **kwargs
     ) -> Optional[str]:
         kwargs.setdefault("headers", "keys")
         kwargs.setdefault("tablefmt", "pipe")
@@ -5648,6 +5648,41 @@ def update(
 Type
 Captive      210.0
 Wild         185.0
+
+We can also choose to include NaN in group keys or not by defining
+`dropna` parameter:
+
+>>> l = [[1, 2, 3], [1, None, 4], [2, 1, 3], [1, 2, 2]]
+>>> df = pd.DataFrame(l, columns=["a", "b", "c"])
+
+>>> df.groupby(by=["b"]).sum()
+    a   c
+b
+1.0 2   3
+2.0 2   5
+
+>>> df.groupby(by=["b"], dropna=False).sum()
+    a   c
+b
+1.0 2   3
+2.0 2   5
+NaN 1   4
+
+>>> l = [["a", 12, 12], [None, 12.3, 33.], ["b", 12.3, 123], ["a", 1, 1]]
+>>> df = pd.DataFrame(l, columns=["a", "b", "c"])
+
+>>> df.groupby(by="a").sum()
+    b     c
+a
+a   13.0   13.0
+b   12.3  123.0
+
+>>> df.groupby(by="a", dropna=False).sum()
+    b     c
+a
+a   13.0   13.0
+b   12.3  123.0
+NaN 12.3   33.0
 """
     )
     @Appender(_shared_docs["groupby"] % _shared_doc_kwargs)
@@ -5661,6 +5696,7 @@ def groupby(
         group_keys: bool = True,
         squeeze: bool = False,
         observed: bool = False,
+        dropna: bool = True,
     ) -> "groupby_generic.DataFrameGroupBy":
 
         if level is None and by is None:
@@ -5677,6 +5713,7 @@ def groupby(
             group_keys=group_keys,
             squeeze=squeeze,
             observed=observed,
+            dropna=dropna,
         )
 
     _shared_docs[

pandas/core/generic.py

@@ -7346,6 +7346,12 @@ def clip(
             If False: show all values for categorical groupers.
 
             .. versionadded:: 0.23.0
+        dropna : bool, default True
+            If 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
         -------

pandas/core/groupby/groupby.py

@@ -373,6 +373,7 @@ def __init__(
         squeeze: bool = False,
         observed: bool = False,
         mutated: bool = False,
+        dropna: bool = True,
     ):
 
         self._selection = selection
@@ -396,6 +397,8 @@ def __init__(
         self.observed = observed
         self.mutated = mutated
 
+        self.dropna = dropna
+
         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: bool = True,
 ) -> GroupBy:
 
     klass: Type[GroupBy]
@@ -2571,4 +2576,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=True
+    ):
         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
 
     @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: bool = True,
     ):
         self.name = name
         self.level = level
@@ -261,6 +266,8 @@ def __init__(
         self.observed = observed
         self.in_axis = in_axis
 
+        self.dropna = dropna
+
         # 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: bool = True,
 ) -> "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/core/series.py

@@ -1433,7 +1433,7 @@ def to_string(
     @Substitution(klass="Series")
     @Appender(generic._shared_docs["to_markdown"])
     def to_markdown(
-        self, buf: Optional[IO[str]] = None, mode: Optional[str] = None, **kwargs,
+        self, buf: Optional[IO[str]] = None, mode: Optional[str] = None, **kwargs
     ) -> Optional[str]:
         return self.to_frame().to_markdown(buf, mode, **kwargs)
 
@@ -1620,6 +1620,34 @@ def _set_name(self, name, inplace=False):
 Captive    210.0
 Wild       185.0
 Name: Max Speed, dtype: float64
+
+We can also choose to include NaN in group keys or not by defining
+`dropna` parameter:
+
+>>> ser = pd.Series([1, 2, 3, 3], index=["a", 'a', 'b', np.nan])
+>>> ser.groupby(level=0).sum()
+a    3
+b    3
+dtype: int64
+
+>>> ser.groupby(level=0, dropna=False).sum()
+a    3
+b    3
+NaN  3
+dtype: int64
+
+>>> arrays = ['Falcon', 'Falcon', 'Parrot', 'Parrot']
+>>> ser = pd.Series([390., 350., 30., 20.], index=arrays, name="Max Speed")
+>>> ser.groupby(["a", "b", "a", np.nan]).mean()
+a    210.0
+b    350.0
+Name: Max Speed, dtype: float64
+
+>>> ser.groupby(["a", "b", "a", np.nan], dropna=False).mean()
+a    210.0
+b    350.0
+NaN   20.0
+Name: Max Speed, dtype: float64
 """
     )
     @Appender(generic._shared_docs["groupby"] % _shared_doc_kwargs)
@@ -1633,6 +1661,7 @@ def groupby(
         group_keys: bool = True,
         squeeze: bool = False,
         observed: bool = False,
+        dropna: bool = True,
     ) -> "groupby_generic.SeriesGroupBy":
 
         if level is None and by is None:
@@ -1649,6 +1678,7 @@ def groupby(
             group_keys=group_keys,
             squeeze=squeeze,
             observed=observed,
+            dropna=dropna,
         )
 
     # ----------------------------------------------------------------------
@@ -4478,9 +4508,7 @@ def to_period(self, freq=None, copy=True):
     hist = pandas.plotting.hist_series
 
 
-Series._setup_axes(
-    ["index"], docs={"index": "The index (axis labels) of the Series."},
-)
+Series._setup_axes(["index"], docs={"index": "The index (axis labels) of the Series."})
 Series._add_numeric_operations()
 Series._add_series_or_dataframe_operations()