Repository: dask/dask

Issue:
Feature Request: add observed- keyword to groupby
Hi, dask developers! Thank you for this very nice package, I have been using extensively for a while.

I am missing Pandas keyword observed for groupby with categorical variables. Without it, the memory used by my processes is huge, as groupby by default creates all possible tuples of categories, no matter whether they are found in the dataset.

Minimal example of the expected behavior:

```
import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame({
    'c1': pd.Categorical(['v11', 'v12', 'v11'], ['v11', 'v12', 'v13']),
    'c2': pd.Categorical(['v21', 'v22', 'v23']),
    'c3': [1, 4, 7]})

ddf = dd.from_pandas(df, npartitions=1)
ddf.groupby(['c1', 'c2'], observed=True).sum().compute()
```

Desired output:

c1 | c2 | c3
-- | -- | --
v11 | v21 | 1
v11 | v23 | 7
v12 | v22 | 4



Discussion:
I am working on implementing it, will keep you posted

Patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -24,6 +24,7 @@
     raise_on_meta_error,
     is_series_like,
     is_dataframe_like,
+    PANDAS_GT_100,
 )
 from ..base import tokenize
 from ..utils import derived_from, M, funcname, itemgetter
@@ -161,24 +162,40 @@ def _groupby_raise_unaligned(df, **kwargs):
 
 
 def _groupby_slice_apply(
-    df, grouper, key, func, *args, group_keys=True, dropna=None, **kwargs
+    df,
+    grouper,
+    key,
+    func,
+    *args,
+    group_keys=True,
+    dropna=None,
+    observed=False,
+    **kwargs
 ):
     # No need to use raise if unaligned here - this is only called after
     # shuffling, which makes everything aligned already
     dropna = {"dropna": dropna} if dropna is not None else {}
-    g = df.groupby(grouper, group_keys=group_keys, **dropna)
+    g = df.groupby(grouper, group_keys=group_keys, observed=observed, **dropna)
     if key:
         g = g[key]
     return g.apply(func, *args, **kwargs)
 
 
 def _groupby_slice_transform(
-    df, grouper, key, func, *args, group_keys=True, dropna=None, **kwargs
+    df,
+    grouper,
+    key,
+    func,
+    *args,
+    group_keys=True,
+    dropna=None,
+    observed=False,
+    **kwargs
 ):
     # No need to use raise if unaligned here - this is only called after
     # shuffling, which makes everything aligned already
     dropna = {"dropna": dropna} if dropna is not None else {}
-    g = df.groupby(grouper, group_keys=group_keys, **dropna)
+    g = df.groupby(grouper, group_keys=group_keys, observed=observed, **dropna)
     if key:
         g = g[key]
 
@@ -273,18 +290,20 @@ def __init__(self, name, chunk, agg, finalize=None):
 
 
 def _groupby_aggregate(
-    df, aggfunc=None, levels=None, dropna=None, sort=False, **kwargs
+    df, aggfunc=None, levels=None, dropna=None, sort=False, observed=False, **kwargs
 ):
     dropna = {"dropna": dropna} if dropna is not None else {}
-    return aggfunc(df.groupby(level=levels, sort=sort, **dropna), **kwargs)
+    if not PANDAS_GT_100 and observed:
+        raise NotImplementedError("``observed`` is only supported for pandas >= 1.0.0")
+    grouped = df.groupby(level=levels, sort=sort, observed=observed, **dropna)
+    return aggfunc(grouped, **kwargs)
 
 
-def _apply_chunk(df, *index, dropna=None, **kwargs):
+def _apply_chunk(df, *index, dropna=None, observed=False, **kwargs):
     func = kwargs.pop("chunk")
     columns = kwargs.pop("columns")
     dropna = {"dropna": dropna} if dropna is not None else {}
-    g = _groupby_raise_unaligned(df, by=index, **dropna)
-
+    g = _groupby_raise_unaligned(df, by=index, observed=observed, **dropna)
     if is_series_like(df) or columns is None:
         return func(g, **kwargs)
     else:
@@ -1029,10 +1048,21 @@ class _GroupBy(object):
     sort: bool, defult None
         Passed along to aggregation methods. If allowed,
         the output aggregation will have sorted keys.
+    observed: bool, default False
+        This only applies if any of the groupers are Categoricals.
+        If True: only show observed values for categorical groupers.
+        If False: show all values for categorical groupers.
     """
 
     def __init__(
-        self, df, by=None, slice=None, group_keys=True, dropna=None, sort=None
+        self,
+        df,
+        by=None,
+        slice=None,
+        group_keys=True,
+        dropna=None,
+        sort=None,
+        observed=False,
     ):
 
         assert isinstance(df, (DataFrame, Series))
@@ -1076,8 +1106,10 @@ def __init__(
         if dropna is not None:
             self.dropna["dropna"] = dropna
 
+        self.observed = observed
+
         self._meta = self.obj._meta.groupby(
-            index_meta, group_keys=group_keys, **self.dropna
+            index_meta, group_keys=group_keys, observed=observed, **self.dropna
         )
 
     @property
@@ -1099,7 +1131,12 @@ def _meta_nonempty(self):
         else:
             index_meta = self.index
 
-        grouped = sample.groupby(index_meta, group_keys=self.group_keys, **self.dropna)
+        grouped = sample.groupby(
+            index_meta,
+            group_keys=self.group_keys,
+            observed=self.observed,
+            **self.dropna,
+        )
         return _maybe_slice(grouped, self._slice)
 
     def _aca_agg(
@@ -1127,14 +1164,22 @@ def _aca_agg(
             else [self.obj] + self.index,
             chunk=_apply_chunk,
             chunk_kwargs=dict(
-                chunk=func, columns=columns, **chunk_kwargs, **self.dropna
+                chunk=func,
+                columns=columns,
+                observed=self.observed,
+                **chunk_kwargs,
+                **self.dropna,
             ),
             aggregate=_groupby_aggregate,
             meta=meta,
             token=token,
             split_every=split_every,
             aggregate_kwargs=dict(
-                aggfunc=aggfunc, levels=levels, **aggregate_kwargs, **self.dropna
+                aggfunc=aggfunc,
+                levels=levels,
+                observed=self.observed,
+                **aggregate_kwargs,
+                **self.dropna,
             ),
             split_out=split_out,
             split_out_setup=split_out_on_index,
@@ -1161,7 +1206,7 @@ def _cum_agg(self, token, chunk, aggregate, initial):
             columns=columns,
             token=name_part,
             meta=meta,
-            **self.dropna
+            **self.dropna,
         )
 
         cumpart_raw_frame = (
@@ -1190,7 +1235,7 @@ def _cum_agg(self, token, chunk, aggregate, initial):
             chunk=M.last,
             meta=meta,
             token=name_last,
-            **self.dropna
+            **self.dropna,
         )
 
         # aggregate cumulated partitions and its previous last element
@@ -1653,8 +1698,9 @@ def apply(self, func, *args, **kwargs):
             token=funcname(func),
             *args,
             group_keys=self.group_keys,
+            observed=self.observed,
             **self.dropna,
-            **kwargs
+            **kwargs,
         )
 
         return df3
@@ -1741,8 +1787,9 @@ def transform(self, func, *args, **kwargs):
             token=funcname(func),
             *args,
             group_keys=self.group_keys,
+            observed=self.observed,
             **self.dropna,
-            **kwargs
+            **kwargs,
         )
 
         return df3
@@ -1797,7 +1844,7 @@ class SeriesGroupBy(_GroupBy):
 
     _token_prefix = "series-groupby-"
 
-    def __init__(self, df, by=None, slice=None, **kwargs):
+    def __init__(self, df, by=None, slice=None, observed=False, **kwargs):
         # for any non series object, raise pandas-compat error message
 
         if isinstance(df, Series):
@@ -1809,12 +1856,12 @@ def __init__(self, df, by=None, slice=None, **kwargs):
 
                 non_series_items = [item for item in by if not isinstance(item, Series)]
                 # raise error from pandas, if applicable
-                df._meta.groupby(non_series_items)
+                df._meta.groupby(non_series_items, observed=observed)
             else:
                 # raise error from pandas, if applicable
-                df._meta.groupby(by)
+                df._meta.groupby(by, observed=observed)
 
-        super().__init__(df, by=by, slice=slice, **kwargs)
+        super().__init__(df, by=by, slice=slice, observed=observed, **kwargs)
 
     @derived_from(pd.core.groupby.SeriesGroupBy)
     def nunique(self, split_every=None, split_out=1):


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
@@ -2420,3 +2420,44 @@ def test_groupby_sort_true_split_out():
     with pytest.raises(NotImplementedError):
         # Cannot use sort=True with split_out>1 (for now)
         M.sum(ddf.groupby("x", sort=True), split_out=2)
+
+
+@pytest.mark.skipif(
+    not dd._compat.PANDAS_GT_100, reason="observed only supported for newer pandas"
+)
+@pytest.mark.parametrize("known_cats", [True, False])
+@pytest.mark.parametrize("ordered_cats", [True, False])
+@pytest.mark.parametrize("groupby", ["cat_1", ["cat_1", "cat_2"]])
+@pytest.mark.parametrize("observed", [True, False])
+def test_groupby_aggregate_categorical_observed(
+    known_cats, ordered_cats, agg_func, groupby, observed
+):
+    if agg_func in ["cov", "corr", "nunique"]:
+        pytest.skip("Not implemented for DataFrameGroupBy yet.")
+    if agg_func in ["sum", "count", "prod"] and groupby != "cat_1":
+        pytest.skip("Gives zeros rather than nans.")
+    if agg_func in ["std", "var"] and observed:
+        pytest.skip("Can't calculate observed with all nans")
+
+    pdf = pd.DataFrame(
+        {
+            "cat_1": pd.Categorical(
+                list("AB"), categories=list("ABCDE"), ordered=ordered_cats
+            ),
+            "cat_2": pd.Categorical([1, 2], categories=[1, 2, 3], ordered=ordered_cats),
+            "value_1": np.random.uniform(size=2),
+        }
+    )
+    ddf = dd.from_pandas(pdf, 2)
+
+    if not known_cats:
+        ddf["cat_1"] = ddf["cat_1"].cat.as_unknown()
+        ddf["cat_2"] = ddf["cat_2"].cat.as_unknown()
+
+    def agg(grp, **kwargs):
+        return getattr(grp, agg_func)(**kwargs)
+
+    assert_eq(
+        agg(pdf.groupby(groupby, observed=observed)),
+        agg(ddf.groupby(groupby, observed=observed)),
+    )