Repository: databricks/koalas

Issue:
GroupBy.head
Implement `GroupBy.head`.

https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.core.groupby.GroupBy.head.html


Discussion:
Hi,
I looked into how one could implement this.
I think that combining the existing standalone head logic with the groupby logic should be enough. That is, after doing group by as usual, we use sdf_limit(n) on the result.

If this makes sense, should I go ahead with this implementation?

Thanks!
@arrayslayer Thanks for the comment!
I'm afraid I don't think the logic works at the first glance, but I might miss something.
Please go ahead and feel free to submit the PR with your logic.
Thanks!
@arrayslayer hi mate, 😸 Are you planning to work on this??
@itholic please go ahead to work on this. Seems inactive.
@HyukjinKwon Thanks for the comment :)
okay then i'm going to work on this from now on.

Resolving patch:
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -30,7 +30,7 @@
 
 from pyspark.sql import Window, functions as F
 from pyspark.sql.types import FloatType, DoubleType, NumericType, StructField, StructType
-from pyspark.sql.functions import PandasUDFType, pandas_udf, Column
+from pyspark.sql.functions import PandasUDFType, pandas_udf, Column, monotonically_increasing_id
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
 from databricks.koalas.typedef import _infer_return_type
@@ -1401,6 +1401,69 @@ def ffill(self, limit=None):
 
     pad = ffill
 
+    def head(self, n=5):
+        """
+        Return first n rows of each group.
+
+        Returns
+        -------
+        DataFrame or Series
+
+        Examples
+        --------
+        >>> df = ks.DataFrame({'a': [1, 1, 1, 1, 2, 2, 2, 3, 3, 3],
+        ...                    'b': [2, 3, 1, 4, 6, 9, 8, 10, 7, 5],
+        ...                    'c': [3, 5, 2, 5, 1, 2, 6, 4, 3, 6]},
+        ...                   columns=['a', 'b', 'c'],
+        ...                   index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6])
+        >>> df
+            a   b  c
+        7   1   2  3
+        2   1   3  5
+        4   1   1  2
+        1   1   4  5
+        3   2   6  1
+        4   2   9  2
+        9   2   8  6
+        10  3  10  4
+        5   3   7  3
+        6   3   5  6
+
+        >>> df.groupby('a').head(2)
+            a   b  c
+        7   1   2  3
+        2   1   3  5
+        10  3  10  4
+        5   3   7  3
+        3   2   6  1
+        4   2   9  2
+
+        >>> df.groupby('a')['b'].head(2)
+        7      2
+        2      3
+        10    10
+        5      7
+        3      6
+        4      9
+        Name: b, dtype: int64
+        """
+        groupkeys = self._groupkeys
+        tmp_col = '__row_number__'
+        sdf = self._kdf._sdf
+        window = Window.partitionBy([s._scol for s in groupkeys]) \
+                       .orderBy(F.monotonically_increasing_id())
+        sdf = sdf.withColumn(tmp_col, F.row_number().over(window)).filter(F.col(tmp_col) <= n)
+        sdf = sdf.select(self._kdf._internal.scols)
+
+        if isinstance(self, DataFrameGroupBy):
+            internal = self._kdf._internal.copy(
+                sdf=sdf,
+                column_scols=[scol_for(sdf, col) for col in self._kdf._internal.data_columns])
+        else:
+            internal = self._kser._internal.copy(sdf, scol=None)
+
+        return DataFrame(internal)
+
     def shift(self, periods=1, fill_value=None):
         """
         Shift each group by periods observations.
@@ -1952,6 +2015,9 @@ def idxmax(self, skipna=True):
 
     idxmax.__doc__ = GroupBy.idxmax.__doc__
 
+    def head(self, n=5):
+        return _col(super(SeriesGroupBy, self).head(n))
+
     # TODO: add keep parameter
     def nsmallest(self, n=5):
         """
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
@@ -52,7 +52,6 @@ class _MissingPandasLikeDataFrameGroupBy(object):
     cumcount = unsupported_function('cumcount')
     describe = unsupported_function('describe')
     get_group = unsupported_function('get_group')
-    head = unsupported_function('head')
     median = unsupported_function('median')
     ngroup = unsupported_function('ngroup')
     nth = unsupported_function('nth')
@@ -94,7 +93,6 @@ class _MissingPandasLikeSeriesGroupBy(object):
     describe = unsupported_function('describe')
     filter = unsupported_function('filter')
     get_group = unsupported_function('get_group')
-    head = unsupported_function('head')
     median = unsupported_function('median')
     ngroup = unsupported_function('ngroup')
     nth = unsupported_function('nth')


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
@@ -817,6 +817,66 @@ def test_idxmin(self):
         self.assert_eq(pdf.groupby(('x', 'a')).idxmin(skipna=False),
                        kdf.groupby(('x', 'a')).idxmin(skipna=False).sort_index())
 
+    def test_head(self):
+        pdf = pd.DataFrame({'a': [1, 1, 1, 1, 2, 2, 2, 3, 3, 3],
+                            'b': [2, 3, 1, 4, 6, 9, 8, 10, 7, 5],
+                            'c': [3, 5, 2, 5, 1, 2, 6, 4, 3, 6]},
+                           columns=['a', 'b', 'c'],
+                           index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6])
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(pdf.groupby('a').head(2).sort_index(),
+                       kdf.groupby('a').head(2).sort_index())
+        self.assert_eq(pdf.groupby('a').head(-2).sort_index(),
+                       kdf.groupby('a').head(-2).sort_index())
+        self.assert_eq(pdf.groupby('a').head(100000).sort_index(),
+                       kdf.groupby('a').head(100000).sort_index())
+
+        self.assert_eq(pdf.groupby('a')['b'].head(2).sort_index(),
+                       kdf.groupby('a')['b'].head(2).sort_index())
+        self.assert_eq(pdf.groupby('a')['b'].head(-2).sort_index(),
+                       kdf.groupby('a')['b'].head(-2).sort_index())
+        self.assert_eq(pdf.groupby('a')['b'].head(100000).sort_index(),
+                       kdf.groupby('a')['b'].head(100000).sort_index())
+
+        # multi-index
+        midx = pd.MultiIndex([['x', 'y'],
+                              ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']],
+                             [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
+                              [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
+        pdf = pd.DataFrame({'a': [1, 1, 1, 1, 2, 2, 2, 3, 3, 3],
+                            'b': [2, 3, 1, 4, 6, 9, 8, 10, 7, 5],
+                            'c': [3, 5, 2, 5, 1, 2, 6, 4, 3, 6]},
+                           columns=['a', 'b', 'c'],
+                           index=midx)
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(pdf.groupby('a').head(2).sort_index(),
+                       kdf.groupby('a').head(2).sort_index())
+        self.assert_eq(pdf.groupby('a').head(-2).sort_index(),
+                       kdf.groupby('a').head(-2).sort_index())
+        self.assert_eq(pdf.groupby('a').head(100000).sort_index(),
+                       kdf.groupby('a').head(100000).sort_index())
+
+        self.assert_eq(pdf.groupby('a')['b'].head(2).sort_index(),
+                       kdf.groupby('a')['b'].head(2).sort_index())
+        self.assert_eq(pdf.groupby('a')['b'].head(-2).sort_index(),
+                       kdf.groupby('a')['b'].head(-2).sort_index())
+        self.assert_eq(pdf.groupby('a')['b'].head(100000).sort_index(),
+                       kdf.groupby('a')['b'].head(100000).sort_index())
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')])
+        pdf.columns = columns
+        kdf.columns = columns
+
+        self.assert_eq(pdf.groupby(('x', 'a')).head(2).sort_index(),
+                       kdf.groupby(('x', 'a')).head(2).sort_index())
+        self.assert_eq(pdf.groupby(('x', 'a')).head(-2).sort_index(),
+                       kdf.groupby(('x', 'a')).head(-2).sort_index())
+        self.assert_eq(pdf.groupby(('x', 'a')).head(100000).sort_index(),
+                       kdf.groupby(('x', 'a')).head(100000).sort_index())
+
     def test_missing(self):
         kdf = ks.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7, 8, 9]})