Repository: dask/dask
Pull request: Add tail and head to SeriesGroupby (#7934)
- [x] Closes #7934 
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask` / `isort dask`


Problem statement:
Series GroupBy head and tail 
I'd like to be able to use head and tail on group-by Series. Today, I can do this in pandas but not dask.
```python
import pandas as pd
import dask.dataframe as dd

pdf = pd.DataFrame({
    "a": [0,1,0,1,1,0],
    "b": range(6),
    "c": ["a","b","c","d","e","f"]
})
ddf = dd.from_pandas(pdf, 2)

print(pdf.groupby("a").b.head(1))
ddf.groupby("a").b.head(1)
```
**Output**
```
0    0
1    1
Name: b, dtype: int64

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-1-6339dbe399f3> in <module>
     10 
     11 print(pdf.groupby("a").b.head(1))
---> 12 ddf.groupby("a").b.head(1)

AttributeError: 'SeriesGroupBy' object has no attribute 'head'
```

The same applies to `tail`. The pandas version is:
```python
dask.__version__
'2021.06.0'
```

Hints from issue discussion:
Looks like the `head` and `tail` methods are defined in the `_Frame` class, but the `SeriesGroupBy` class inherits only from `_GroupBy` and so doesn't know anything about them. 

https://github.com/dask/dask/blob/1fab96cb763e7273079083d0093a55aad096549e/dask/dataframe/core.py#L1057-L1121
Ah, I see you have a WIP open. Let us know when you'd like someone to look over it.

Commit message: fix issue with index levels

Patch to review:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index a845c8ada63..cbc53287536 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -1140,6 +1140,7 @@ def _aca_agg(
         token,
         func,
         aggfunc=None,
+        metafunc=None,
         split_every=None,
         split_out=1,
         chunk_kwargs={},
@@ -1148,7 +1149,10 @@ def _aca_agg(
         if aggfunc is None:
             aggfunc = func
 
-        meta = func(self._meta_nonempty)
+        if metafunc is None:
+            metafunc = func
+
+        meta = metafunc(self._meta_nonempty)
         columns = meta.name if is_series_like(meta) else meta.columns
 
         token = self._token_prefix + token
@@ -1804,7 +1808,6 @@ def transform(self, func, *args, **kwargs):
 
 
 class DataFrameGroupBy(_GroupBy):
-
     _token_prefix = "dataframe-groupby-"
 
     def __getitem__(self, key):
@@ -1849,7 +1852,6 @@ def agg(self, arg, split_every=None, split_out=1):
 
 
 class SeriesGroupBy(_GroupBy):
-
     _token_prefix = "series-groupby-"
 
     def __init__(self, df, by=None, slice=None, observed=None, **kwargs):
@@ -1949,6 +1951,34 @@ def unique(self, split_every=None, split_out=1):
             split_out=split_out,
         )
 
+    @derived_from(pd.core.groupby.SeriesGroupBy)
+    def tail(self, n=5, split_every=None, split_out=1):
+        index_levels = len(self.index) if isinstance(self.index, list) else 1
+        return self._aca_agg(
+            token="tail",
+            func=_tail_chunk,
+            aggfunc=_tail_aggregate,
+            metafunc=M.tail,
+            chunk_kwargs={"n": n},
+            aggregate_kwargs={"n": n, "index_levels": index_levels},
+            split_every=split_every,
+            split_out=split_out,
+        )
+
+    @derived_from(pd.core.groupby.SeriesGroupBy)
+    def head(self, n=5, split_every=None, split_out=1):
+        index_levels = len(self.index) if isinstance(self.index, list) else 1
+        return self._aca_agg(
+            token="head",
+            func=_head_chunk,
+            aggfunc=_head_aggregate,
+            metafunc=M.head,
+            chunk_kwargs={"n": n},
+            aggregate_kwargs={"n": n, "index_levels": index_levels},
+            split_every=split_every,
+            split_out=split_out,
+        )
+
 
 def _unique_aggregate(series_gb, name=None):
     ret = type(series_gb.obj)(
@@ -1969,3 +1999,23 @@ def _value_counts_aggregate(series_gb):
     to_concat = {k: v.groupby(level=1).sum() for k, v in series_gb}
     names = list(series_gb.obj.index.names)
     return pd.Series(pd.concat(to_concat, names=names))
+
+
+def _tail_chunk(series_gb, **kwargs):
+    keys, groups = zip(*series_gb) if len(series_gb) else ((True,), (series_gb,))
+    return pd.concat([group.tail(**kwargs) for group in groups], keys=keys)
+
+
+def _tail_aggregate(series_gb, **kwargs):
+    levels = kwargs.pop("index_levels")
+    return series_gb.tail(**kwargs).droplevel(list(range(levels)))
+
+
+def _head_chunk(series_gb, **kwargs):
+    keys, groups = zip(*series_gb) if len(series_gb) else ((True,), (series_gb,))
+    return pd.concat([group.head(**kwargs) for group in groups], keys=keys)
+
+
+def _head_aggregate(series_gb, **kwargs):
+    levels = kwargs.pop("index_levels")
+    return series_gb.head(**kwargs).droplevel(list(range(levels)))
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index 9132bf95796..bf7bf650134 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -1,5 +1,6 @@
 import collections
 import warnings
+from operator import methodcaller
 
 import numpy as np
 import pandas as pd
@@ -524,6 +525,8 @@ def test_split_apply_combine_on_series(empty):
         assert_eq(ddf.groupby(ddkey).a.size(), pdf.groupby(pdkey).a.size())
         assert_eq(ddf.groupby(ddkey).a.first(), pdf.groupby(pdkey).a.first())
         assert_eq(ddf.groupby(ddkey).a.last(), pdf.groupby(pdkey).a.last())
+        assert_eq(ddf.groupby(ddkey).a.tail(), pdf.groupby(pdkey).a.tail())
+        assert_eq(ddf.groupby(ddkey).a.head(), pdf.groupby(pdkey).a.head())
         for ddof in ddofs:
             assert_eq(ddf.groupby(ddkey).a.var(ddof), pdf.groupby(pdkey).a.var(ddof))
             assert_eq(ddf.groupby(ddkey).a.std(ddof), pdf.groupby(pdkey).a.std(ddof))
@@ -598,6 +601,8 @@ def test_split_apply_combine_on_series(empty):
         assert_eq(ddf.groupby(ddf.b > i).a.size(), pdf.groupby(pdf.b > i).a.size())
         assert_eq(ddf.groupby(ddf.b > i).a.first(), pdf.groupby(pdf.b > i).a.first())
         assert_eq(ddf.groupby(ddf.b > i).a.last(), pdf.groupby(pdf.b > i).a.last())
+        assert_eq(ddf.groupby(ddf.b > i).a.tail(), pdf.groupby(pdf.b > i).a.tail())
+        assert_eq(ddf.groupby(ddf.b > i).a.head(), pdf.groupby(pdf.b > i).a.head())
         assert_eq(ddf.groupby(ddf.b > i).a.prod(), pdf.groupby(pdf.b > i).a.prod())
 
         assert_eq(ddf.groupby(ddf.a > i).b.sum(), pdf.groupby(pdf.a > i).b.sum())
@@ -611,6 +616,8 @@ def test_split_apply_combine_on_series(empty):
         assert_eq(ddf.groupby(ddf.b > i).b.size(), pdf.groupby(pdf.b > i).b.size())
         assert_eq(ddf.groupby(ddf.b > i).b.first(), pdf.groupby(pdf.b > i).b.first())
         assert_eq(ddf.groupby(ddf.b > i).b.last(), pdf.groupby(pdf.b > i).b.last())
+        assert_eq(ddf.groupby(ddf.b > i).b.tail(), pdf.groupby(pdf.b > i).b.tail())
+        assert_eq(ddf.groupby(ddf.b > i).b.head(), pdf.groupby(pdf.b > i).b.head())
         assert_eq(ddf.groupby(ddf.b > i).b.prod(), pdf.groupby(pdf.b > i).b.prod())
 
         assert_eq(ddf.groupby(ddf.b > i).sum(), pdf.groupby(pdf.b > i).sum())
@@ -694,6 +701,8 @@ def test_split_apply_combine_on_series(empty):
     assert_dask_graph(ddf.groupby("b").a.cov(), "series-groupby-cov")
     assert_dask_graph(ddf.groupby("b").a.first(), "series-groupby-first")
     assert_dask_graph(ddf.groupby("b").a.last(), "series-groupby-last")
+    assert_dask_graph(ddf.groupby("b").a.tail(), "series-groupby-tail")
+    assert_dask_graph(ddf.groupby("b").a.head(), "series-groupby-head")
     assert_dask_graph(ddf.groupby("b").a.prod(), "series-groupby-prod")
     # mean consists from sum and count operations
     assert_dask_graph(ddf.groupby("b").a.mean(), "series-groupby-sum")
@@ -2413,3 +2422,76 @@ def test_groupby_with_pd_grouper():
         ddf.groupby(pd.Grouper(key="key1"))
     with pytest.raises(NotImplementedError):
         ddf.groupby(["key1", pd.Grouper(key="key2")])
+
+
+@pytest.mark.parametrize("operation", ["head", "tail"])
+def test_empty_partitions_with_rows_operation(operation):
+
+    df = pd.DataFrame(
+        data=[
+            ["a1", "b1"],
+            ["a1", None],
+            ["a1", "b1"],
+            [None, None],
+            [None, None],
+            [None, None],
+            ["a3", "b3"],
+            ["a3", "b3"],
+            ["a5", "b5"],
+        ],
+        columns=["A", "B"],
+    )
+
+    caller = methodcaller(operation, 1)
+    expected = caller(df.groupby("A")["B"])
+    ddf = dd.from_pandas(df, npartitions=3)
+    actual = caller(ddf.groupby("A")["B"])
+    assert_eq(expected, actual)
+
+
+@pytest.mark.parametrize("operation", ["head", "tail"])
+def test_groupby_with_row_operations(operation):
+    df = pd.DataFrame(
+        data=[
+            ["a0", "b1"],
+            ["a0", "b2"],
+            ["a1", "b1"],
+            ["a3", "b3"],
+            ["a3", "b3"],
+            ["a5", "b5"],
+            ["a1", "b1"],
+            ["a1", "b1"],
+            ["a1", "b1"],
+        ],
+        columns=["A", "B"],
+    )
+
+    caller = methodcaller(operation)
+    expected = caller(df.groupby("A")["B"])
+    ddf = dd.from_pandas(df, npartitions=3)
+    actual = caller(ddf.groupby("A")["B"])
+    assert_eq(expected, actual)
+
+
+@pytest.mark.parametrize("operation", ["head", "tail"])
+def test_groupby_with_row_operations(operation):
+    df = pd.DataFrame(
+        data=[
+            ["a0", "b1"],
+            ["a0", "b2"],
+            ["a1", "b1"],
+            ["a3", "b3"],
+            ["a3", "b3"],
+            ["a5", "b5"],
+            ["a1", "b1"],
+            ["a1", "b1"],
+            ["a1", "b1"],
+        ],
+        columns=["A", "B"],
+    )
+
+    caller = methodcaller(operation)
+    expected = caller(df.groupby(["A", df["A"].eq("a1")])["B"])
+    ddf = dd.from_pandas(df, npartitions=3)
+    actual = caller(ddf.groupby(["A", ddf["A"].eq("a1")])["B"])
+    assert_eq(expected, actual)


Code review comments:
diff_hunk: @@ -1143,6 +1143,7 @@ def _aca_agg(
         token,
         func,
         aggfunc=None,
+        metafunc=None,
path: dask/dataframe/groupby.py
text: @user1:
It would probably be more palatable to include an optional `meta` kwarg that accepts a precomputed `meta` rather than `metafunc`.