Repository: modin-project/modin

Issue:
Implementation for `groupby.shift`
Similar to #1192 and #1287, which was done by @prutskov. This is a high priority.


Discussion:
+++

Resolving patch:
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -2113,6 +2113,10 @@ def shift(self, periods=1, freq=None, axis=0, fill_value=None):
                         new_frame.columns = self.columns.copy()
                     return new_frame
             else:
+                if not isinstance(self, DataFrame):
+                    raise ValueError(
+                        f"No axis named {axis} for object type {type(self)}"
+                    )
                 res_columns = self.columns
                 from .general import concat
 
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -22,6 +22,7 @@
 Manually add documentation for methods which are not presented in pandas.
 """
 
+import numpy as np
 import pandas
 import pandas.core.groupby
 from pandas.core.dtypes.common import is_list_like
@@ -30,6 +31,7 @@
 
 from modin.error_message import ErrorMessage
 from modin.utils import _inherit_docstrings, try_cast_to_pandas
+from modin.backends.base.query_compiler import BaseQueryCompiler
 from modin.config import IsExperimental
 from .series import Series
 
@@ -171,10 +173,64 @@ def idxmax(self):
     def ndim(self):
         return 2  # ndim is always 2 for DataFrames
 
-    def shift(self, periods=1, freq=None, axis=0):
-        return self._default_to_pandas(
-            lambda df: df.shift(periods=periods, freq=freq, axis=axis)
-        )
+    def shift(self, periods=1, freq=None, axis=0, fill_value=None):
+        def _shift(periods, freq, axis, fill_value, is_set_nan_rows=True):
+            from .dataframe import DataFrame
+
+            result = self._df.shift(periods, freq, axis, fill_value)
+
+            if (
+                is_set_nan_rows
+                and isinstance(self._by, BaseQueryCompiler)
+                and (
+                    # Check using `issubset` is effective only in case of MultiIndex
+                    set(self._by.columns).issubset(list(self._df.columns))
+                    if isinstance(self._by.columns, pandas.MultiIndex)
+                    else len(
+                        self._by.columns.unique()
+                        .sort_values()
+                        .difference(self._df.columns.unique().sort_values())
+                    )
+                    == 0
+                )
+                and DataFrame(query_compiler=self._by.isna()).any(axis=None)
+            ):
+                mask_nan_rows = self._df[self._by.columns].isna()
+                if (isinstance(mask_nan_rows, DataFrame)) and len(
+                    mask_nan_rows.columns
+                ) == 1:
+                    mask_nan_rows = mask_nan_rows.squeeze(axis=1)
+                idx_nan_rows = mask_nan_rows[
+                    mask_nan_rows.any(axis=1)
+                    if (isinstance(mask_nan_rows, DataFrame))
+                    else mask_nan_rows
+                ].index
+                result.loc[idx_nan_rows] = np.nan
+            return result
+
+        if freq is None and axis == 1 and self._axis == 0:
+            result = _shift(periods, freq, axis, fill_value)
+        elif (
+            freq is not None
+            and axis == 0
+            and self._axis == 0
+            and isinstance(self._by, BaseQueryCompiler)
+        ):
+            result = _shift(periods, freq, axis, fill_value, is_set_nan_rows=False)
+            new_idx_lvl_arrays = np.concatenate(
+                [self._df[self._by.columns].values.T, [list(result.index)]]
+            )
+            result.index = pandas.MultiIndex.from_arrays(
+                new_idx_lvl_arrays,
+                names=[col_name for col_name in self._by.columns] + [result.index.name],
+            )
+            result = result.dropna(subset=self._by.columns).sort_index()
+        else:
+            result = self._apply_agg_function(
+                lambda df: df.shift(periods, freq, axis, fill_value)
+            )
+            result.index.name = None
+        return result
 
     def nth(self, n, dropna=None):
         return self._default_to_pandas(lambda df: df.nth(n, dropna=dropna))


Test patch:
diff --git a/modin/pandas/test/test_groupby.py b/modin/pandas/test/test_groupby.py
--- a/modin/pandas/test/test_groupby.py
+++ b/modin/pandas/test/test_groupby.py
@@ -104,6 +104,7 @@ def test_mixed_dtypes_groupby(as_index):
             modin_df_almost_equals_pandas,
             is_default=True,
         )
+        eval_shift(modin_groupby, pandas_groupby)
         eval_mean(modin_groupby, pandas_groupby)
         eval_any(modin_groupby, pandas_groupby)
         eval_min(modin_groupby, pandas_groupby)
@@ -298,6 +299,7 @@ def maybe_get_columns(df, by):
 
     modin_groupby_equals_pandas(modin_groupby, pandas_groupby)
     eval_ngroups(modin_groupby, pandas_groupby)
+    eval_shift(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.ffill(), is_default=True)
     eval_general(
         modin_groupby,
@@ -437,6 +439,7 @@ def test_single_group_row_groupby():
 
     modin_groupby_equals_pandas(modin_groupby, pandas_groupby)
     eval_ngroups(modin_groupby, pandas_groupby)
+    eval_shift(modin_groupby, pandas_groupby)
     eval_skew(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.ffill(), is_default=True)
     eval_general(
@@ -552,6 +555,7 @@ def test_large_row_groupby(is_by_category):
 
     modin_groupby_equals_pandas(modin_groupby, pandas_groupby)
     eval_ngroups(modin_groupby, pandas_groupby)
+    eval_shift(modin_groupby, pandas_groupby)
     eval_skew(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.ffill(), is_default=True)
     eval_general(
@@ -666,6 +670,7 @@ def test_simple_col_groupby():
 
     modin_groupby_equals_pandas(modin_groupby, pandas_groupby)
     eval_ngroups(modin_groupby, pandas_groupby)
+    eval_shift(modin_groupby, pandas_groupby)
     eval_skew(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.ffill(), is_default=True)
     eval_general(
@@ -796,6 +801,7 @@ def test_series_groupby(by, as_index_series_or_dataframe):
 
         modin_groupby_equals_pandas(modin_groupby, pandas_groupby)
         eval_ngroups(modin_groupby, pandas_groupby)
+        eval_shift(modin_groupby, pandas_groupby)
         eval_general(
             modin_groupby, pandas_groupby, lambda df: df.ffill(), is_default=True
         )
@@ -1069,7 +1075,26 @@ def eval_groups(modin_groupby, pandas_groupby):
 
 
 def eval_shift(modin_groupby, pandas_groupby):
-    assert modin_groupby.groups == pandas_groupby.groups
+    eval_general(
+        modin_groupby,
+        pandas_groupby,
+        lambda groupby: groupby.shift(),
+    )
+    eval_general(
+        modin_groupby,
+        pandas_groupby,
+        lambda groupby: groupby.shift(periods=0),
+    )
+    eval_general(
+        modin_groupby,
+        pandas_groupby,
+        lambda groupby: groupby.shift(periods=-3),
+    )
+    eval_general(
+        modin_groupby,
+        pandas_groupby,
+        lambda groupby: groupby.shift(axis=1, fill_value=777),
+    )
 
 
 def test_groupby_on_index_values_with_loop():
@@ -1122,6 +1147,38 @@ def test_groupby_multiindex():
     df_equals(modin_df.groupby(by=by).count(), pandas_df.groupby(by=by).count())
 
 
+@pytest.mark.parametrize("groupby_axis", [0, 1])
+@pytest.mark.parametrize("shift_axis", [0, 1])
+def test_shift_freq(groupby_axis, shift_axis):
+    pandas_df = pandas.DataFrame(
+        {
+            "col1": [1, 0, 2, 3],
+            "col2": [4, 5, np.NaN, 7],
+            "col3": [np.NaN, np.NaN, 12, 10],
+            "col4": [17, 13, 16, 15],
+        }
+    )
+    modin_df = from_pandas(pandas_df)
+
+    new_index = pandas.date_range("1/12/2020", periods=4, freq="S")
+    if groupby_axis == 0 and shift_axis == 0:
+        pandas_df.index = modin_df.index = new_index
+        by = [["col2", "col3"], ["col2"], ["col4"], [0, 1, 0, 2]]
+    else:
+        pandas_df.index = modin_df.index = new_index
+        pandas_df.columns = modin_df.columns = new_index
+        by = [[0, 1, 0, 2]]
+
+    for _by in by:
+        pandas_groupby = pandas_df.groupby(by=_by, axis=groupby_axis)
+        modin_groupby = modin_df.groupby(by=_by, axis=groupby_axis)
+        eval_general(
+            modin_groupby,
+            pandas_groupby,
+            lambda groupby: groupby.shift(axis=shift_axis, freq="S"),
+        )
+
+
 def test_agg_func_None_rename():
     pandas_df = pandas.DataFrame(
         {
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -2903,6 +2903,7 @@ def test_shift(data):
     df_equals(modin_series.shift(fill_value=777), pandas_series.shift(fill_value=777))
     df_equals(modin_series.shift(periods=7), pandas_series.shift(periods=7))
     df_equals(modin_series.shift(periods=-3), pandas_series.shift(periods=-3))
+    eval_general(modin_series, pandas_series, lambda df: df.shift(axis=1))
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)