Repository: databricks/koalas

Issue:
GroupBy.shift
Standard Spark functionality is missing but implemented in the inner _shift function. Adding part_cols will allow this to easily scale.


Discussion:
Had an offline discussion and found the request is to support `groupby(...).shift()`.

Resolving patch:
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -1387,6 +1387,65 @@ def ffill(self, limit=None):
 
     pad = ffill
 
+    def shift(self, periods=1, fill_value=None):
+        """
+        Shift each group by periods observations.
+
+        Parameters
+        ----------
+        periods : integer, default 1
+            number of periods to shift
+        fill_value : optional
+
+        Returns
+        -------
+        Series or DataFrame
+            Object shifted within each group.
+
+        Examples
+        --------
+
+        >>> df = ks.DataFrame({
+        ...     'a': [1, 1, 1, 2, 2, 2, 3, 3, 3],
+        ...     'b': [1, 2, 2, 2, 3, 3, 3, 4, 4]}, columns=['a', 'b'])
+        >>> df
+           a  b
+        0  1  1
+        1  1  2
+        2  1  2
+        3  2  2
+        4  2  3
+        5  2  3
+        6  3  3
+        7  3  4
+        8  3  4
+
+        >>> df.groupby('a').shift().sort_index()  # doctest: +SKIP
+             b
+        0  NaN
+        1  1.0
+        2  2.0
+        3  NaN
+        4  2.0
+        5  3.0
+        6  NaN
+        7  3.0
+        8  4.0
+
+        >>> df.groupby('a').shift(periods=-1, fill_value=0).sort_index()  # doctest: +SKIP
+           b
+        0  2
+        1  2
+        2  0
+        3  3
+        4  3
+        5  0
+        6  4
+        7  4
+        8  0
+        """
+        return self._shift(periods, fill_value)
+
     def transform(self, func):
         """
         Apply function column-by-column to the GroupBy object.
@@ -1736,6 +1795,19 @@ def _fillna(self, *args, **kwargs):
                                       column_index=[c._internal.column_index[0] for c in applied])
         return DataFrame(internal)
 
+    def _shift(self, periods, fill_value):
+        applied = []
+        kdf = self._kdf
+
+        for column in self._agg_columns:
+            applied.append(column.groupby(self._groupkeys).shift(periods, fill_value))
+
+        sdf = kdf._sdf.select(kdf._internal.index_scols + [c._scol for c in applied])
+        internal = kdf._internal.copy(sdf=sdf,
+                                      data_columns=[c._internal.data_columns[0] for c in applied],
+                                      column_index=[c._internal.column_index[0] for c in applied])
+        return DataFrame(internal)
+
 
 class SeriesGroupBy(GroupBy):
 
@@ -1772,6 +1844,10 @@ def _fillna(self, *args, **kwargs):
         groupkey_scols = [s._scol for s in self._groupkeys]
         return Series._fillna(self._ks, *args, **kwargs, part_cols=groupkey_scols)
 
+    def _shift(self, periods, fill_value):
+        groupkey_scols = [s._scol for s in self._groupkeys]
+        return Series._shift(self._ks, periods, fill_value, part_cols=groupkey_scols)
+
     @property
     def _kdf(self) -> DataFrame:
         return self._ks._kdf
diff --git a/databricks/koalas/missing/groupby.py b/databricks/koalas/missing/groupby.py
--- a/databricks/koalas/missing/groupby.py
+++ b/databricks/koalas/missing/groupby.py
@@ -62,7 +62,6 @@ class _MissingPandasLikeDataFrameGroupBy(object):
     prod = unsupported_function('prod')
     resample = unsupported_function('resample')
     sem = unsupported_function('sem')
-    shift = unsupported_function('shift')
     tail = unsupported_function('tail')
 
 
@@ -105,5 +104,4 @@ class _MissingPandasLikeSeriesGroupBy(object):
     prod = unsupported_function('prod')
     resample = unsupported_function('resample')
     sem = unsupported_function('sem')
-    shift = unsupported_function('shift')
     tail = unsupported_function('tail')


Test patch:
diff --git a/databricks/koalas/tests/test_groupby.py b/databricks/koalas/tests/test_groupby.py
--- a/databricks/koalas/tests/test_groupby.py
+++ b/databricks/koalas/tests/test_groupby.py
@@ -561,6 +561,39 @@ def test_bfill(self):
             self.assert_eq(kdf.groupby(("X", "A")).bfill(),
                            pdf.groupby(("X", "A")).bfill())
 
+    @unittest.skipIf(pd.__version__ < '0.24.0', "not supported before pandas 0.24.0")
+    def test_shift(self):
+        pdf = pd.DataFrame({'a': [1, 1, 2, 2, 3, 3],
+                            'b': [1, 1, 2, 2, 3, 4],
+                            'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c'])
+        kdf = koalas.DataFrame(pdf)
+
+        self.assert_eq(kdf.groupby('a').shift().sort_index(),
+                       pdf.groupby('a').shift().sort_index())
+        # TODO: seems like a pandas' bug when fill_value is not None?
+        # self.assert_eq(kdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index(),
+        #                pdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index())
+        self.assert_eq(kdf.groupby(['b'])['a'].shift().sort_index(),
+                       pdf.groupby(['b'])['a'].shift().sort_index(), almost=True)
+        self.assert_eq(kdf.groupby(['a', 'b'])['c'].shift().sort_index(),
+                       pdf.groupby(['a', 'b'])['c'].shift().sort_index(), almost=True)
+        self.assert_eq(kdf.groupby(['b'])[['a', 'c']].shift(periods=-1, fill_value=0).sort_index(),
+                       pdf.groupby(['b'])[['a', 'c']].shift(periods=-1, fill_value=0).sort_index(),
+                       almost=True)
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')])
+        pdf.columns = columns
+        kdf.columns = columns
+
+        self.assert_eq(kdf.groupby(('x', 'a')).shift().sort_index(),
+                       pdf.groupby(('x', 'a')).shift().sort_index())
+        # TODO: seems like a pandas' bug when fill_value is not None?
+        # self.assert_eq(kdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1,
+        #                                                            fill_value=0).sort_index(),
+        #                pdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1,
+        #                                                            fill_value=0).sort_index())
+
     def test_apply(self):
         pdf = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6],
                             'b': [1, 1, 2, 3, 5, 8],