Add diff in Groupby (#622)

This PR adds rank GroupBy (both SeriesGroupBy and DataFrameGroupBy) by using existing diff logic in Series.

Author: RainFung

databricks/koalas/groupby.py

@@ -425,6 +425,63 @@ def size(self):
                                              for i, s in enumerate(groupkeys)])
         return _col(DataFrame(internal))
 
+    def diff(self, periods=1):
+        """
+        First discrete difference of element.
+
+        Calculates the difference of a DataFrame element compared with another element in the
+        DataFrame group (default is the element in the same column of the previous row).
+
+        Parameters
+        ----------
+        periods : int, default 1
+            Periods to shift for calculating difference, accepts negative values.
+
+        Returns
+        -------
+        diffed : DataFrame or Series
+
+        See Also
+        --------
+        databricks.koalas.Series.groupby
+        databricks.koalas.DataFrame.groupby
+
+        Examples
+        --------
+        >>> df = ks.DataFrame({'a': [1, 2, 3, 4, 5, 6],
+        ...                    'b': [1, 1, 2, 3, 5, 8],
+        ...                    'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c'])
+        >>> df
+           a  b   c
+        0  1  1   1
+        1  2  1   4
+        2  3  2   9
+        3  4  3  16
+        4  5  5  25
+        5  6  8  36
+
+        >>> df.groupby(['b']).diff().sort_index()
+             a    c
+        0  NaN  NaN
+        1  1.0  3.0
+        2  NaN  NaN
+        3  NaN  NaN
+        4  NaN  NaN
+        5  NaN  NaN
+
+        Difference with previous column in a group.
+
+        >>> df.groupby(['b'])['a'].diff().sort_index()
+        0    NaN
+        1    1.0
+        2    NaN
+        3    NaN
+        4    NaN
+        5    NaN
+        Name: a, dtype: float64
+        """
+        return self._diff(periods)
+
     def cummax(self):
         """
         Cumulative max for each group.
@@ -1164,6 +1221,19 @@ def __getitem__(self, item):
             return DataFrameGroupBy(self._kdf, self._groupkeys, as_index=self._as_index,
                                     agg_columns=item)
 
+    def _diff(self, *args, **kwargs):
+        applied = []
+        kdf = self._kdf
+        groupkey_columns = [s.name for s in self._groupkeys]
+
+        for column in kdf._internal.data_columns:
+            if column not in groupkey_columns:
+                applied.append(kdf[column].groupby(self._groupkeys)._diff(*args, **kwargs))
+
+        sdf = kdf._sdf.select(kdf._internal.index_scols + [c._scol for c in applied])
+        internal = kdf._internal.copy(sdf=sdf, data_columns=[c.name for c in applied])
+        return DataFrame(internal)
+
     def _rank(self, *args, **kwargs):
         applied = []
         kdf = self._kdf
@@ -1226,6 +1296,10 @@ def __getattr__(self, item: str) -> Any:
                 return partial(property_or_func, self)
         raise AttributeError(item)
 
+    def _diff(self, *args, **kwargs):
+        groupkey_scols = [s._scol for s in self._groupkeys]
+        return Series._diff(self._ks, *args, **kwargs, part_cols=groupkey_scols)
+
     def _cum(self, func):
         groupkey_scols = [s._scol for s in self._groupkeys]
         return Series._cum(self._ks, func, True, part_cols=groupkey_scols)

databricks/koalas/missing/groupby.py

@@ -33,7 +33,6 @@ class _MissingPandasLikeDataFrameGroupBy(object):
     corr = unsupported_property('corr')
     corrwith = unsupported_property('corrwith')
     cov = unsupported_property('cov')
-    diff = unsupported_property('diff')
     dtypes = unsupported_property('dtypes')
     fillna = unsupported_property('fillna')
     groups = unsupported_property('groups')
@@ -82,7 +81,6 @@ class _MissingPandasLikeSeriesGroupBy(object):
     # Properties
     corr = unsupported_property('corr')
     cov = unsupported_property('cov')
-    diff = unsupported_property('diff')
     dtype = unsupported_property('dtype')
     fillna = unsupported_property('fillna')
     groups = unsupported_property('groups')

databricks/koalas/series.py

@@ -2638,16 +2638,17 @@ def diff(self, periods=1):
         5     NaN
         Name: c, dtype: float64
         """
+        return self._diff(periods)
 
+    def _diff(self, periods, part_cols=()):
         if len(self._internal.index_columns) == 0:
             raise ValueError("Index must be set.")
-
         if not isinstance(periods, int):
             raise ValueError('periods should be an int; however, got [%s]' % type(periods))
-
-        col = self._scol
-        window = Window.orderBy(self._internal.index_scols).rowsBetween(-periods, -periods)
-        return self._with_new_scol(col - F.lag(col, periods).over(window)).alias(self.name)
+        window = Window.partitionBy(*part_cols).orderBy(self._internal.index_scols)\
+            .rowsBetween(-periods, -periods)
+        scol = self._scol - F.lag(self._scol, periods).over(window)
+        return Series(self._kdf._internal.copy(scol=scol), anchor=self._kdf).rename(self.name)
 
     def idxmax(self, skipna=True):
         """

databricks/koalas/tests/test_groupby.py

@@ -170,6 +170,18 @@ def test_size(self):
         self.assert_eq(kdf.groupby(['A', 'B']).size().sort_index(),
                        pdf.groupby(['A', 'B']).size().sort_index())
 
+    def test_diff(self):
+        pdf = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6],
+                            'b': [1, 1, 2, 3, 5, 8],
+                            'c': [1, 4, 9, 16, 25, 36]}, columns=['a', 'b', 'c'])
+        kdf = koalas.DataFrame(pdf)
+        self.assert_eq(kdf.groupby("b").diff().sort_index(),
+                       pdf.groupby("b").diff().sort_index())
+        self.assert_eq(kdf.groupby(['a', 'b']).diff().sort_index(),
+                       pdf.groupby(['a', 'b']).diff().sort_index())
+        self.assert_eq(repr(kdf.groupby(['b'])['a'].diff().sort_index()),
+                       repr(pdf.groupby(['b'])['a'].diff().sort_index()))
+
     def test_missing(self):
         kdf = koalas.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7, 8, 9]})
 

docs/source/reference/groupby.rst

@@ -43,3 +43,4 @@ Computations / Descriptive Stats
    GroupBy.sum
    GroupBy.var
    GroupBy.size
+   GroupBy.diff