Repository: dask/dask

Issue:
Add dropna keyword to groupby aggregations
Cudf recently started supporting a `dropna=` keyword to groupby aggregations.  See https://github.com/rapidsai/cudf/issues/2627

Pandas considered this in https://github.com/pandas-dev/pandas/pull/21669 , where there was some discussion around the API.

Would it be reasonable to add this to dask dataframe to support cudf workloads, even if there is not yet support in Pandas?

cc @TomAugspurger 


Discussion:
Yep, that's reasonable.

On Wed, Nov 6, 2019 at 2:20 PM Matthew Rocklin <notifications@github.com>
wrote:

> Cudf recently started supporting a dropna= keyword to groupby
> aggregations. See rapidsai/cudf#2627
> <https://github.com/rapidsai/cudf/issues/2627>
>
> Pandas considered this in pandas-dev/pandas#21669
> <https://github.com/pandas-dev/pandas/pull/21669> , where there was some
> discussion around the API.
>
> Would it be reasonable to add this to dask dataframe to support cudf
> workloads, even if there is not yet support in Pandas?
>
> cc @TomAugspurger <https://github.com/TomAugspurger>
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/5565?email_source=notifications&email_token=AAKAOIV5S6CVPDKJTAKCJTLQSMKGVA5CNFSM4JJ3IFXKYY3PNVWWK3TUL52HS4DFUVEXG43VMWVGG33NNVSW45C7NFSM4HXKOJBQ>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAKAOITES7BUFFTK55NLGMLQSMKGVANCNFSM4JJ3IFXA>
> .
>

@rjzamora , this is low priority, but if you want to take a look that would be welcome.

Resolving patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -162,22 +162,25 @@ def _groupby_raise_unaligned(df, **kwargs):
     return df.groupby(**kwargs)
 
 
-def _groupby_slice_apply(df, grouper, key, func, *args, **kwargs):
+def _groupby_slice_apply(
+    df, grouper, key, func, *args, group_keys=True, dropna=None, **kwargs
+):
     # No need to use raise if unaligned here - this is only called after
     # shuffling, which makes everything aligned already
-    group_keys = kwargs.pop("group_keys", True)
-    g = df.groupby(grouper, group_keys=group_keys)
+    dropna = {"dropna": dropna} if dropna is not None else {}
+    g = df.groupby(grouper, group_keys=group_keys, **dropna)
     if key:
         g = g[key]
     return g.apply(func, *args, **kwargs)
 
 
-def _groupby_slice_transform(df, grouper, key, func, *args, **kwargs):
+def _groupby_slice_transform(
+    df, grouper, key, func, *args, group_keys=True, dropna=None, **kwargs
+):
     # No need to use raise if unaligned here - this is only called after
     # shuffling, which makes everything aligned already
-    group_keys = kwargs.pop("group_keys", True)
-
-    g = df.groupby(grouper, group_keys=group_keys)
+    dropna = {"dropna": dropna} if dropna is not None else {}
+    g = df.groupby(grouper, group_keys=group_keys, **dropna)
     if key:
         g = g[key]
 
@@ -271,15 +274,16 @@ def __init__(self, name, chunk, agg, finalize=None):
         self.__name__ = name
 
 
-def _groupby_aggregate(df, aggfunc=None, levels=None, **kwargs):
-    return aggfunc(df.groupby(level=levels, sort=False), **kwargs)
+def _groupby_aggregate(df, aggfunc=None, levels=None, dropna=None, **kwargs):
+    dropna = {"dropna": dropna} if dropna is not None else {}
+    return aggfunc(df.groupby(level=levels, sort=False, **dropna), **kwargs)
 
 
-def _apply_chunk(df, *index, **kwargs):
+def _apply_chunk(df, *index, dropna=None, **kwargs):
     func = kwargs.pop("chunk")
     columns = kwargs.pop("columns")
-
-    g = _groupby_raise_unaligned(df, by=index)
+    dropna = {"dropna": dropna} if dropna is not None else {}
+    g = _groupby_raise_unaligned(df, by=index, **dropna)
 
     if is_series_like(df) or columns is None:
         return func(g, **kwargs)
@@ -980,9 +984,11 @@ class _GroupBy(object):
         The slice keys applied to GroupBy result
     group_keys: bool
         Passed to pandas.DataFrame.groupby()
+    dropna: bool
+        Whether to drop null values from groupby index
     """
 
-    def __init__(self, df, by=None, slice=None, group_keys=True):
+    def __init__(self, df, by=None, slice=None, group_keys=True, dropna=None):
 
         assert isinstance(df, (DataFrame, Series))
         self.group_keys = group_keys
@@ -1020,7 +1026,13 @@ def __init__(self, df, by=None, slice=None, group_keys=True):
         else:
             index_meta = self.index
 
-        self._meta = self.obj._meta.groupby(index_meta, group_keys=group_keys)
+        self.dropna = {}
+        if dropna is not None:
+            self.dropna["dropna"] = dropna
+
+        self._meta = self.obj._meta.groupby(
+            index_meta, group_keys=group_keys, **self.dropna
+        )
 
     @property
     def _meta_nonempty(self):
@@ -1041,7 +1053,7 @@ def _meta_nonempty(self):
         else:
             index_meta = self.index
 
-        grouped = sample.groupby(index_meta, group_keys=self.group_keys)
+        grouped = sample.groupby(index_meta, group_keys=self.group_keys, **self.dropna)
         return _maybe_slice(grouped, self._slice)
 
     def _aca_agg(
@@ -1068,12 +1080,16 @@ def _aca_agg(
             if not isinstance(self.index, list)
             else [self.obj] + self.index,
             chunk=_apply_chunk,
-            chunk_kwargs=dict(chunk=func, columns=columns, **chunk_kwargs),
+            chunk_kwargs=dict(
+                chunk=func, columns=columns, **chunk_kwargs, **self.dropna
+            ),
             aggregate=_groupby_aggregate,
             meta=meta,
             token=token,
             split_every=split_every,
-            aggregate_kwargs=dict(aggfunc=aggfunc, levels=levels, **aggregate_kwargs),
+            aggregate_kwargs=dict(
+                aggfunc=aggfunc, levels=levels, **aggregate_kwargs, **self.dropna
+            ),
             split_out=split_out,
             split_out_setup=split_out_on_index,
         )
@@ -1097,7 +1113,8 @@ def _cum_agg(self, token, chunk, aggregate, initial):
             chunk=chunk,
             columns=columns,
             token=name_part,
-            meta=meta
+            meta=meta,
+            **self.dropna
         )
 
         cumpart_raw_frame = (
@@ -1125,7 +1142,8 @@ def _cum_agg(self, token, chunk, aggregate, initial):
             columns=0 if columns is None else columns,
             chunk=M.last,
             meta=meta,
-            token=name_last
+            token=name_last,
+            **self.dropna
         )
 
         # aggregate cumulated partitions and its previous last element
@@ -1585,6 +1603,7 @@ def apply(self, func, *args, **kwargs):
             token=funcname(func),
             *args,
             group_keys=self.group_keys,
+            **self.dropna,
             **kwargs
         )
 
@@ -1672,6 +1691,7 @@ def transform(self, func, *args, **kwargs):
             token=funcname(func),
             *args,
             group_keys=self.group_keys,
+            **self.dropna,
             **kwargs
         )
 
@@ -1684,9 +1704,9 @@ class DataFrameGroupBy(_GroupBy):
 
     def __getitem__(self, key):
         if isinstance(key, list):
-            g = DataFrameGroupBy(self.obj, by=self.index, slice=key)
+            g = DataFrameGroupBy(self.obj, by=self.index, slice=key, **self.dropna)
         else:
-            g = SeriesGroupBy(self.obj, by=self.index, slice=key)
+            g = SeriesGroupBy(self.obj, by=self.index, slice=key, **self.dropna)
 
         # error is raised from pandas
         g._meta = g._meta[key]


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -2187,3 +2187,60 @@ def test_groupby_transform_ufunc_partitioning(npartitions, indexed):
                 lambda series: series - series.mean()
             ),
         )
+
+
+@pytest.mark.xfail(reason="dropna kwarg not supported in pandas groupby.")
+@pytest.mark.parametrize("dropna", [False, True])
+def test_groupby_dropna_pandas(dropna):
+
+    # The `dropna` arg is not currently supported by pandas
+    # (See #https://github.com/pandas-dev/pandas/pull/21669)
+    # Dask supports the argument for the cudf backend,
+    # but passing it to the pandas backend will fail.
+
+    # TODO: Expand test when `dropna` is supported in pandas.
+    #       (See: `test_groupby_dropna_cudf`)
+
+    df = pd.DataFrame(
+        {"a": [1, 2, 3, 4, None, None, 7, 8], "e": [4, 5, 6, 3, 2, 1, 0, 0]}
+    )
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    dask_result = ddf.groupby("a", dropna=dropna)
+    pd_result = df.groupby("a", dropna=dropna)
+    assert_eq(dask_result, pd_result)
+
+
+@pytest.mark.parametrize("dropna", [False, True, None])
+@pytest.mark.parametrize("by", ["a", "c", "d", ["a", "b"], ["a", "c"], ["a", "d"]])
+def test_groupby_dropna_cudf(dropna, by):
+
+    # NOTE: This test requires cudf/dask_cudf, and will
+    # be skipped by non-GPU CI
+
+    cudf = pytest.importorskip("cudf")
+    dask_cudf = pytest.importorskip("dask_cudf")
+
+    df = cudf.DataFrame(
+        {
+            "a": [1, 2, 3, 4, None, None, 7, 8],
+            "b": [1, 0] * 4,
+            "c": ["a", "b", None, None, "e", "f", "g", "h"],
+            "e": [4, 5, 6, 3, 2, 1, 0, 0],
+        }
+    )
+    df["d"] = df["c"].astype("category")
+    ddf = dask_cudf.from_cudf(df, npartitions=3)
+
+    if dropna is None:
+        dask_result = ddf.groupby(by).e.sum()
+        cudf_result = df.groupby(by).e.sum()
+    else:
+        dask_result = ddf.groupby(by, dropna=dropna).e.sum()
+        cudf_result = df.groupby(by, dropna=dropna).e.sum()
+    if by in ["c", "d"]:
+        # Loose string/category index name in cudf...
+        dask_result = dask_result.compute()
+        dask_result.index.name = cudf_result.index.name
+
+    assert_eq(dask_result, cudf_result)