Repository: databricks/koalas

Issue:
DataFrame.GroupBy.apply()
I'm continuously missing this behavior when converting my pandas scripts to koalas:

```python
df = pd.DataFrame({"timestamp":[0.0, 0.5, 1.0, 0.0, 0.5],
              "car_id": ['A','A','A','B','B'], 
              "battery_charge": [100, 90, 80,100,90]
              })

print(df)

def calc_battery_usage(x):
    start_bat = x.sort_values('timestamp').iloc[0]['battery_charge']
    stop_bat = x.sort_values('timestamp').iloc[-1]['battery_charge']
    return start_bat - stop_bat

print(df.groupby('car_id').apply(calc_battery_usage))
```
That prints out:
```python
   timestamp car_id  battery_charge
0        0.0      A             100
1        0.5      A              90
2        1.0      A              80
3        0.0      B             100
4        0.5      B              90

car_id
A    20
B    10
```

On koalas I just get:
```
PandasNotImplementedError: The method `pd.groupby.GroupBy.apply()` is not implemented yet.
```

I can take over and start implementing it myself, but would probably need some general guidelines on how to do it (probably using pysparks pandas_udf).


Discussion:
Is this similar to https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#grouped-map and https://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#grouped-aggregate ? Perhaps we can just map to that.

Yes, this is similar. In my current scripts these are the methods I'm using to "manually" make the apply on groupby's. 
Hi @patryk-oleniuk! I think it'd great to have groupby apply in koalas. 

Spark's groupby apply and Pandas's groupby apply are pretty similar. However, two main difference that I can think of are:

- Spark's groupby apply need full return type specification on the udf definition
- Spark's groupby apply doesn't keep index, while pandas groupby apply use the grouping key as index.

As a start, I'd probably try this something like:
```
@pandas_wraps(full_return_schema):
def calc_battery_usage(x):
    start_bat = x.sort_values('timestamp').iloc[0]['battery_charge']
    stop_bat = x.sort_values('timestamp').iloc[-1]['battery_charge']
    return start_bat - stop_bat

kdf.groupby('car_id').apply(calc_battery_usage)
```
And just wrap it with Spark's groupby apply + pandas_udf

Does that make sense?


@icexelloss that's nice details. Let me give a shot!

Resolving patch:
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
--- a/databricks/koalas/generic.py
+++ b/databricks/koalas/generic.py
@@ -929,7 +929,7 @@ def skew(self, numeric_only=True):
 
         On a DataFrame:
 
-        >>> df.skew()
+        >>> df.skew()  # doctest: +SKIP
         a    0.000000e+00
         b   -3.319678e-16
         dtype: float64
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -18,14 +18,18 @@
 A wrapper for GroupedData to behave similar to pandas GroupBy.
 """
 
+import inspect
+from collections import Callable
 from functools import partial
 from typing import Any, List
 import numpy as np
 
 from pyspark.sql import functions as F
 from pyspark.sql.types import FloatType, DoubleType, NumericType
+from pyspark.sql.functions import PandasUDFType, pandas_udf
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
+from databricks.koalas.typedef import _infer_return_type
 from databricks.koalas.frame import DataFrame
 from databricks.koalas.internal import _InternalFrame
 from databricks.koalas.missing.groupby import _MissingPandasLikeDataFrameGroupBy, \
@@ -340,6 +344,124 @@ def any(self):
             lambda col: F.max(F.coalesce(col.cast('boolean'), F.lit(False))),
             only_numeric=False)
 
+    # TODO: Series support is not implemented yet.
+    def apply(self, func):
+        """
+        Apply function `func` group-wise and combine the results together.
+
+        The function passed to `apply` must take a DataFrame as its first
+        argument and return a DataFrame. `apply` will
+        then take care of combining the results back together into a single
+        dataframe. `apply` is therefore a highly flexible
+        grouping method.
+
+        While `apply` is a very flexible method, its downside is that
+        using it can be quite a bit slower than using more specific methods
+        like `agg` or `transform`. Koalas offers a wide range of method that will
+        be much faster than using `apply` for their specific purposes, so try to
+        use them before reaching for `apply`.
+
+        .. note:: unlike pandas, it is required for ``func`` to specify return type hint.
+
+        .. note:: the output column names are `c0, c1, c2 ... cn`. These names
+            are positionally mapped to the returned DataFrame in ``func``. See examples below.
+
+        .. note:: the dataframe within ``func`` is actually a pandas dataframe. Therefore,
+            any pandas APIs within this function is allowed.
+
+        Parameters
+        ----------
+        func : callable
+            A callable that takes a DataFrame as its first argument, and
+            returns a dataframe.
+
+        Returns
+        -------
+        applied : DataFrame
+
+        See Also
+        --------
+        aggregate : Apply aggregate function to the GroupBy object.
+        Series.apply : Apply a function to a Series.
+
+        Examples
+        --------
+        >>> df = ks.DataFrame({'A': 'a a b'.split(),
+        ...                    'B': [1, 2, 3],
+        ...                    'C': [4, 6, 5]}, columns=['A', 'B', 'C'])
+        >>> g = df.groupby('A')
+
+        Notice that ``g`` has two groups, ``a`` and ``b``.
+        Calling `apply` in various ways, we can get different grouping results:
+
+        Below the functions passed to `apply` takes a DataFrame as
+        its argument and returns a DataFrame. `apply` combines the result for
+        each group together into a new DataFrame:
+
+        >>> def pandas_div_sum(x) -> ks.DataFrame[float, float]:
+        ...    return x[['B', 'C']] / x[['B', 'C']].sum()
+        >>> g.apply(pandas_div_sum)  # doctest: +NORMALIZE_WHITESPACE
+                 c0   c1
+        0  1.000000  1.0
+        1  0.333333  0.4
+        2  0.666667  0.6
+
+        >>> def plus_max(x) -> ks.DataFrame[str, np.int, np.int]:
+        ...    return x + x.max()
+        >>> g.apply(plus_max)  # doctest: +NORMALIZE_WHITESPACE
+           c0  c1  c2
+        0  bb   6  10
+        1  aa   3  10
+        2  aa   4  12
+        """
+        if not isinstance(func, Callable):
+            raise TypeError("%s object is not callable" % type(func))
+
+        assert callable(func), "the first argument should be a callable function."
+        spec = inspect.getfullargspec(func)
+        return_sig = spec.annotations.get("return", None)
+        if return_sig is None:
+            raise ValueError("Given function must have return type hint; however, not found.")
+
+        return_schema = _infer_return_type(func).tpe
+
+        index_columns = self._kdf._internal.index_columns
+        index_names = self._kdf._internal.index_names
+        data_columns = self._kdf._internal.data_columns
+
+        def rename_output(pdf):
+            # TODO: This logic below was borrowed from `DataFrame.pandas_df` to set the index
+            #   within each pdf properly. we might have to deduplicate it.
+            import pandas as pd
+
+            if len(index_columns) > 0:
+                append = False
+                for index_field in index_columns:
+                    drop = index_field not in data_columns
+                    pdf = pdf.set_index(index_field, drop=drop, append=append)
+                    append = True
+                pdf = pdf[data_columns]
+
+            if len(index_names) > 0:
+                if isinstance(pdf.index, pd.MultiIndex):
+                    pdf.index.names = index_names
+                else:
+                    pdf.index.name = index_names[0]
+
+            pdf = func(pdf)
+            # For now, just positionally map the column names to given schema's.
+            pdf = pdf.rename(columns=dict(zip(pdf.columns, return_schema.fieldNames())))
+            return pdf
+
+        grouped_map_func = pandas_udf(return_schema, PandasUDFType.GROUPED_MAP)(rename_output)
+
+        sdf = self._kdf._sdf
+        input_groupkeys = [s._scol for s in self._groupkeys]
+        sdf = sdf.groupby(*input_groupkeys).apply(grouped_map_func)
+        internal = _InternalFrame(
+            sdf=sdf, data_columns=return_schema.fieldNames(), index_map=[])  # index is lost.
+        return DataFrame(internal)
+
     def _reduce_for_stat_function(self, sfun, only_numeric):
         groupkeys = self._groupkeys
         groupkey_cols = [s._scol.alias('__index_level_{}__'.format(i))
@@ -426,3 +548,12 @@ def _agg_columns(self):
 
     def _reduce_for_stat_function(self, sfun, only_numeric):
         return _col(super(SeriesGroupBy, self)._reduce_for_stat_function(sfun, only_numeric))
+
+    def agg(self, func_or_funcs, *args, **kwargs):
+        raise NotImplementedError()
+
+    def aggregate(self, func_or_funcs, *args, **kwargs):
+        raise NotImplementedError()
+
+    def apply(self, func):
+        raise NotImplementedError()
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):
     take = unsupported_property('take', deprecated=True)
 
     # Functions
-    apply = unsupported_function('apply')
     backfill = unsupported_function('backfill')
     bfill = unsupported_function('bfill')
     boxplot = unsupported_function('boxplot')
@@ -115,7 +114,6 @@ class _MissingPandasLikeSeriesGroupBy(object):
     take = unsupported_property('take', deprecated=True)
 
     # Functions
-    apply = unsupported_function('apply')
     backfill = unsupported_function('backfill')
     bfill = unsupported_function('bfill')
     cumcount = unsupported_function('cumcount')


Test patch: