Repository: dask/dask
Pull request: Add groupby shift method
- [x] Closes #7095
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

I have to implemented the shift `method` following the `transform` and `apply` methods. Let me know if you'd like any changes/improvements.

Problem statement:
Groupby shift
I'd like to be able to use `shift` on groupby Series and DataFrame objects to create per-group lag or lead columns. 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.shift(1))
ddf.groupby("a").b.shift(1)
0    NaN
1    NaN
2    0.0
3    1.0
4    3.0
5    2.0
Name: b, dtype: float64
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-17-f89d506bbdcd> in <module>
     11 
     12 print(pdf.groupby("a").b.shift(1))
---> 13 ddf.groupby("a").b.shift(1)

AttributeError: 'SeriesGroupBy' object has no attribute 'shift'
```

```
conda list | grep "pandas\|dask"
dask                      2021.1.0           pyhd8ed1ab_0    conda-forge
pandas                    1.1.5            py37hdc94413_0    conda-forge
```

Hints from issue discussion:
That seems like a reasonable feature request. Are you interested in working on it? In the meantime I wonder if `map_overlap` could suit your needs.
Thanks for the suggestion! I'll take a look at map_overlap.

I don't have the bandwidth to work on the implementation at the moment, but may in the future.
FYI, cuDF recently refactored and implemented groupby.shfit: https://github.com/rapidsai/cudf/pull/8131
Hi, I'd like to do the same thing with map_overlap, and I did the following to compute groupby+shift
`ddf.groupby(['a']).map_overlap(lambda x: x.shift(1), 1, 1).compute()`
However, I got:
`AttributeError: 'Column not found: map_overlap'`
I wonder if there is a way for groupby + (other operation like shift)
@beckernick @quasiben 
You should take a look at https://docs.dask.org/en/latest/dataframe-groupby.html#aggregate for tips on how to write custom aggregations. Or feel free to open a PR to add `shift` :)
I would love to work on this issue. 

This seems  similar to how `transform` or `apply `methods are implemented. My initial approach is to write a `_groupby_slice_shift` function and calling `shift` in it instead of `apply` or `transform`. I think we can then use it inside `map_partitions`.

Commit message: remove groupby shift example that fails doctest

provide meta for doc example

Patch to review:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index efe9206010f..4042e0d63c1 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -188,6 +188,19 @@ def _groupby_slice_transform(
     return g.transform(func, *args, **kwargs)
 
 
+def _groupby_slice_shift(
+    df, grouper, key, group_keys=True, dropna=None, observed=None, **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 {}
+    observed = {"observed": observed} if observed is not None else {}
+    g = df.groupby(grouper, group_keys=group_keys, **observed, **dropna)
+    if key:
+        g = g[key]
+    return g.shift(**kwargs)
+
+
 def _groupby_get_group(df, by_key, get_key, columns):
     # SeriesGroupBy may pass df which includes group key
     grouped = _groupby_raise_unaligned(df, by=by_key)
@@ -1800,6 +1813,94 @@ def transform(self, func, *args, **kwargs):
 
         return df3
 
+    @insert_meta_param_description(pad=12)
+    def shift(self, periods=1, freq=None, axis=0, fill_value=None, meta=no_default):
+        """Parallel version of pandas GroupBy.shift
+
+        This mimics the pandas version except for the following:
+
+        If the grouper does not align with the index then this causes a full
+        shuffle.  The order of rows within each group may not be preserved.
+
+        Parameters
+        ----------
+        periods : Delayed, Scalar or int, default 1
+            Number of periods to shift.
+        freq : Delayed, Scalar or str, optional
+            Frequency string.
+        axis : axis to shift, default 0
+            Shift direction.
+        fill_value : Scalar, Delayed or object, optional
+            The scalar value to use for newly introduced missing values.
+        $META
+
+        Returns
+        -------
+        shifted : Series or DataFrame shifted within each group.
+
+        Examples
+        --------
+        >>> import dask
+        >>> ddf = dask.datasets.timeseries(freq="1H")
+        >>> result = ddf.groupby("name").shift(1, meta={"id": int, "x": float, "y": float})
+        """
+        if meta is no_default:
+            with raise_on_meta_error("groupby.shift()", udf=False):
+                meta_kwargs = _extract_meta(
+                    {
+                        "periods": periods,
+                        "freq": freq,
+                        "axis": axis,
+                        "fill_value": fill_value,
+                    },
+                    nonempty=True,
+                )
+                meta = self._meta_nonempty.shift(**meta_kwargs)
+
+            msg = (
+                "`meta` is not specified, inferred from partial data. "
+                "Please provide `meta` if the result is unexpected.\n"
+                "  Before: .shift(1)\n"
+                "  After:  .shift(1, meta={'x': 'f8', 'y': 'f8'}) for dataframe result\n"
+                "  or:     .shift(1, meta=('x', 'f8'))            for series result"
+            )
+            warnings.warn(msg, stacklevel=2)
+
+        meta = make_meta(meta, parent_meta=self._meta.obj)
+
+        # Validate self.by
+        if isinstance(self.by, list) and any(
+            isinstance(item, Series) for item in self.by
+        ):
+            raise NotImplementedError(
+                "groupby-transform with a multiple Series is currently not supported"
+            )
+        df = self.obj
+        should_shuffle = not (df.known_divisions and df._contains_index_name(self.by))
+
+        if should_shuffle:
+            df2, by = self._shuffle(meta)
+        else:
+            df2 = df
+            by = self.by
+
+        # Perform embarrassingly parallel groupby-shift
+        result = map_partitions(
+            _groupby_slice_shift,
+            df2,
+            by,
+            self._slice,
+            periods=periods,
+            freq=freq,
+            axis=axis,
+            fill_value=fill_value,
+            group_keys=self.group_keys,
+            meta=meta,
+            **self.observed,
+            **self.dropna,
+        )
+        return result
+
     def rolling(self, window, min_periods=None, center=False, win_type=None, axis=0):
         """Provides rolling transformations.
 
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index 84b06dd1f46..4efe529bc94 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -2040,6 +2040,87 @@ def test_groupby_value_counts():
     assert_eq(dd_gb, pd_gb)
 
 
+@pytest.mark.parametrize("npartitions", [1, 2, 5])
+@pytest.mark.parametrize("period", [1, -1, 10])
+@pytest.mark.parametrize("axis", [0, 1])
+def test_groupby_shift_basic_input(npartitions, period, axis):
+    pdf = pd.DataFrame(
+        {
+            "a": [0, 0, 1, 1, 2, 2, 3, 3, 3],
+            "b": [4, 5, 6, 3, 2, 1, 0, 0, 0],
+            "c": [0, 0, 0, 0, 0, 1, 1, 1, 1],
+        },
+    )
+    ddf = dd.from_pandas(pdf, npartitions=npartitions)
+    with pytest.warns(UserWarning):
+        assert_eq(
+            pdf.groupby(["a", "c"]).shift(period, axis=axis),
+            ddf.groupby(["a", "c"]).shift(period, axis=axis),
+        )
+    with pytest.warns(UserWarning):
+        assert_eq(
+            pdf.groupby(["a"]).shift(period, axis=axis),
+            ddf.groupby(["a"]).shift(period, axis=axis),
+        )
+    with pytest.warns(UserWarning):
+        assert_eq(
+            pdf.groupby(pdf.c).shift(period, axis=axis),
+            ddf.groupby(ddf.c).shift(period, axis=axis),
+        )
+
+
+def test_groupby_shift_series():
+    pdf = pd.DataFrame(
+        {
+            "a": [0, 0, 1, 1, 2, 2, 3, 3, 3],
+            "b": [4, 5, 6, 3, 2, 1, 0, 0, 0],
+        },
+    )
+    ddf = dd.from_pandas(pdf, npartitions=3)
+    with pytest.warns(UserWarning):
+        assert_eq(
+            pdf.groupby("a")["b"].shift(periods=2),
+            ddf.groupby("a")["b"].shift(periods=2),
+        )
+
+
+def test_groupby_shift_lazy_input():
+    pdf = pd.DataFrame(
+        {
+            "a": [0, 0, 1, 1, 2, 2, 3, 3, 3],
+            "b": [4, 5, 6, 3, 2, 1, 0, 0, 0],
+            "c": [0, 0, 0, 0, 0, 1, 1, 1, 1],
+        },
+    )
+    delayed_periods = dask.delayed(lambda: 1)()
+    ddf = dd.from_pandas(pdf, npartitions=3)
+    assert_eq(
+        pdf.groupby(pdf.c).shift(periods=1),
+        ddf.groupby(ddf.c).shift(periods=delayed_periods, meta={"a": int, "b": int}),
+    )
+    with pytest.warns(UserWarning):
+        assert_eq(
+            pdf.groupby(pdf.c).shift(periods=1, fill_value=pdf.b.max()),
+            ddf.groupby(ddf.c).shift(periods=1, fill_value=ddf.b.max()),
+        )
+
+
+def test_groupby_shift_with_freq():
+    pdf = pd.DataFrame(
+        dict(a=[1, 2, 3, 4, 5, 6], b=[0, 0, 0, 1, 1, 1]),
+        index=pd.date_range(start="20100101", periods=6),
+    )
+    ddf = dd.from_pandas(pdf, npartitions=3)
+
+    # just pass the pandas result as meta for convenience
+    df_result = pdf.groupby(pdf.index).shift(periods=-2, freq="D")
+    assert_eq(
+        df_result, ddf.groupby(ddf.index).shift(periods=-2, freq="D", meta=df_result)
+    )
+    df_result = pdf.groupby("b").shift(periods=-2, freq="D")
+    assert_eq(df_result, ddf.groupby("b").shift(periods=-2, freq="D", meta=df_result))
+
+
 @pytest.mark.parametrize(
     "transformation", [lambda x: x.sum(), np.sum, "sum", pd.Series.rank]
 )


Code review comments:
diff_hunk: @@ -1800,6 +1813,94 @@ def transform(self, func, *args, **kwargs):
 
         return df3
 
+    @insert_meta_param_description(pad=12)
+    def shift(self, periods=1, freq=None, axis=0, fill_value=None, meta=no_default):
+        """Parallel version of pandas GroupBy.shift
+
+        This mimics the pandas version except for the following:
+
+        If the grouper does not align with the index then this causes a full
+        shuffle.  The order of rows within each group may not be preserved.
+
+        Parameters
+        ----------
+        periods : Delayed, Scalar or int, default 1
+            Number of periods to shift.
+        freq : Delayed, Scalar or str, optional
+            Frequency string.
+        axis : axis to shift, default 0
+            Shift direction.
+        fill_value : Scalar, Delayed or object, optional
+            The scalar value to use for newly introduced missing values.
+        $META
+
+        Returns
+        -------
+        shifted : Series or DataFrame shifted within each group.
+
+        Examples
+        --------
+        >>> import dask
+        >>> ddf = dask.datasets.timeseries(freq="1H")
+        >>> result = ddf.groupby("name").shift(1, meta={"id": int, "x": float, "y": float})
+        """
+        if meta is no_default:
+            with raise_on_meta_error("groupby.shift()", udf=False):
+                meta_kwargs = _extract_meta(
+                    {
+                        "periods": periods,
+                        "freq": freq,
+                        "axis": axis,
+                        "fill_value": fill_value,
+                    },
+                    nonempty=True,
+                )
+                meta = self._meta_nonempty.shift(**meta_kwargs)
+
+            msg = (
+                "`meta` is not specified, inferred from partial data. "
+                "Please provide `meta` if the result is unexpected.\n"
+                "  Before: .shift(1)\n"
+                "  After:  .shift(1, meta={'x': 'f8', 'y': 'f8'}) for dataframe result\n"
+                "  or:     .shift(1, meta=('x', 'f8'))            for series result"
+            )
+            warnings.warn(msg, stacklevel=2)
+
+        meta = make_meta(meta, parent_meta=self._meta.obj)
+
+        # Validate self.by
+        if isinstance(self.by, list) and any(
+            isinstance(item, Series) for item in self.by
+        ):
+            raise NotImplementedError(
+                "groupby-transform with a multiple Series is currently not supported"
path: dask/dataframe/groupby.py
text: @user1:
copy-paste holdover :)
```suggestion
                "groupby-shift with a multiple Series is currently not supported"
```