Repository: databricks/koalas

Issue:
UDF using DataFrame.apply
These are the patterns that I have seen with the apply() function:
```
df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7], 'b': [7, 6, 5, 4, 3, 2, 1]})

# Option 1: Use a series and apply on it to run a UDF (Technically a Series.apply)
df['c'] = df['a'].apply(lambda x: x*2)

# Option 2: Use the entire row from a dataframe
df['d'] = df.apply(lambda x: x.a*x.b, axis=1)

# Option 3: Use some columns in a row
df['e'] = df[['a', 'b']].apply(lambda x: x.a*x.b, axis=1)

# Option 4: Apply aggregates on every column to create a new dataframe
new_df = df.apply(lambda x: x.sum())

# Option 5: Apply aggregates but return multiple values - here it adjusts the shorted columns with NaNs
new_df = df.apply(lambda x: x[:x[0]])
```



Discussion:
There are some problems to use it directly.

1. Current Scalar UDF doesn't support aggregation across whole data. If we do the same thing, it will aggregate each batch unit that's internally used (arrow batch).

2. If we restrict it's 1 to 1 then we can consider `transform` instead. But we shouldn't forget we should check if the input & output length are same.

3. When `axis` is different, it will need a different execution style to row by row, which will probably be like https://github.com/databricks/koalas/commit/d397c1a0067fb41336de7447b60735e1888d91d0 way.
Actually, I think we could hack this via `df.groupby(F.spark_partition_id()).apply(pands_udf)`. Let me take a look soon.
Would groupby-apply introduce exchange/shuffle on every `apply`? If so, seems like that would be an expensive operation.
Yeah, we should avoid. With Spark 3.0, we can avoid a shuffle with `mapInPandas` API.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -41,7 +41,8 @@
 from pyspark.sql import functions as F, Column
 from pyspark.sql.functions import pandas_udf
 from pyspark.sql.types import (BooleanType, ByteType, DecimalType, DoubleType, FloatType,
-                               IntegerType, LongType, NumericType, ShortType)
+                               IntegerType, LongType, NumericType, ShortType, StructType,
+                               StructField)
 from pyspark.sql.window import Window
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
@@ -1800,6 +1801,229 @@ def transpose(self):
 
     T = property(transpose)
 
+    def apply(self, func, axis=0):
+        """
+        Apply a function along an axis of the DataFrame.
+
+        Objects passed to the function are Series objects whose index is
+        either the DataFrame's index (``axis=0``) or the DataFrame's columns
+        (``axis=1``).
+
+        .. note:: when `axis` is 0 or 'index', the `func` is unable to access
+            to the whole input series. Koalas internally splits the input series into multiple
+            batches and calls `func` with each batch multiple times. Therefore, operations
+            such as global aggregations are impossible. See the example below.
+
+            >>> # This case does not return the length of whole series but of the batch internally
+            ... # used.
+            ... def length(s) -> int:
+            ...    return len(s)
+            ...
+            >>> df = ks.DataFrame({'A': range(1000)})
+            >>> df.apply(length, axis=0)  # doctest: +SKIP
+            0     83
+            1     83
+            2     83
+            ...
+            10    83
+            11    83
+            Name: 0, dtype: int32
+
+        .. note:: this API executes the function once to infer the type which is
+            potentially expensive, for instance, when the dataset is created after
+            aggregations or sorting.
+
+            To avoid this, specify return type in ``func``, for instance, as below:
+
+            >>> def square(s) -> ks.Series[np.int32]:
+            ...     return s ** 2
+
+            Koalas uses return type hint and does not try to infer the type.
+
+            In case when axis is 1, it requires to specify `DataFrame` with type hints
+            as below:
+
+            >>> def plus_one(x) -> ks.DataFrame[float, float]:
+            ...    return x + 1
+
+            If the return type is specified, the output column names become
+            `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
+            DataFrame in ``func``. See examples below.
+
+
+        Parameters
+        ----------
+        func : function
+            Function to apply to each column or row.
+        axis : {0 or 'index', 1 or 'columns'}, default 0
+            Axis along which the function is applied:
+
+            * 0 or 'index': apply function to each column.
+            * 1 or 'columns': apply function to each row.
+
+        Returns
+        -------
+        Series or DataFrame
+            Result of applying ``func`` along the given axis of the
+            DataFrame.
+
+        See Also
+        --------
+        DataFrame.applymap: For elementwise operations.
+        DataFrame.aggregate: Only perform aggregating type operations.
+        DataFrame.transform: Only perform transforming type operations.
+
+        Examples
+        --------
+        >>> df = ks.DataFrame([[4, 9]] * 3, columns=['A', 'B'])
+        >>> df
+           A  B
+        0  4  9
+        1  4  9
+        2  4  9
+
+        Using a numpy universal function (in this case the same as
+        ``np.sqrt(df)``):
+
+        >>> def sqrt(x) -> ks.Series[float]:
+        ...    return np.sqrt(x)
+        ...
+        >>> df.apply(sqrt, axis=0)
+             A    B
+        0  2.0  3.0
+        1  2.0  3.0
+        2  2.0  3.0
+
+        You can omit the type hint and let Koalas infer its type.
+
+        >>> df.apply(np.sqrt, axis=0)
+             A    B
+        0  2.0  3.0
+        1  2.0  3.0
+        2  2.0  3.0
+
+        When `axis` is 1 or 'columns', it applies the function for each row.
+
+        >>> def summation(x) -> np.int64:
+        ...    return np.sum(x)
+        ...
+        >>> df.apply(summation, axis=1)
+        0    13
+        1    13
+        2    13
+        Name: 0, dtype: int64
+
+        Likewise, you can omit the type hint and let Koalas infer its type.
+
+        >>> df.apply(np.sum, axis=1)
+        0    13
+        1    13
+        2    13
+        Name: 0, dtype: int64
+
+        Returning a list-like will result in a Series
+
+        >>> df.apply(lambda x: [1, 2], axis=1)
+        0    [1, 2]
+        1    [1, 2]
+        2    [1, 2]
+        Name: 0, dtype: object
+
+        In order to specify the types when `axis` is '1', it should use DataFrame[...]
+        annotation. In this case, the column names are automatically generated.
+
+        >>> def identify(x) -> ks.DataFrame[np.int64, np.int64]:
+        ...     return x
+        ...
+        >>> df.apply(identify, axis=1)
+           c0  c1
+        0   4   9
+        1   4   9
+        2   4   9
+        """
+        from databricks.koalas.groupby import GroupBy
+        from databricks.koalas.series import _col
+
+        if isinstance(func, np.ufunc):
+            f = func
+            func = lambda *args, **kwargs: f(*args, **kwargs)
+
+        assert callable(func), "the first argument should be a callable function."
+
+        axis = validate_axis(axis)
+        should_return_series = False
+        spec = inspect.getfullargspec(func)
+        return_sig = spec.annotations.get("return", None)
+        should_infer_schema = return_sig is None
+
+        def apply_func(pdf):
+            pdf_or_pser = pdf.apply(func, axis=axis)
+            if isinstance(pdf_or_pser, pd.Series):
+                return pdf_or_pser.to_frame()
+            else:
+                return pdf_or_pser
+
+        if should_infer_schema:
+            # Here we execute with the first 1000 to get the return type.
+            # If the records were less than 1000, it uses pandas API directly for a shortcut.
+            limit = get_option("compute.shortcut_limit")
+            pdf = self.head(limit + 1)._to_internal_pandas()
+            applied = pdf.apply(func, axis=axis)
+            kser_or_kdf = ks.from_pandas(applied)
+            if len(pdf) <= limit:
+                return kser_or_kdf
+
+            kdf = kser_or_kdf
+            if isinstance(kser_or_kdf, ks.Series):
+                should_return_series = True
+                kdf = kser_or_kdf.to_frame()
+
+            return_schema = kdf._internal._sdf.drop(*HIDDEN_COLUMNS).schema
+
+            sdf = GroupBy._spark_group_map_apply(
+                self, apply_func, (F.spark_partition_id(),),
+                return_schema, retain_index=True)
+
+            # If schema is inferred, we can restore indexes too.
+            internal = kdf._internal.copy(sdf=sdf,
+                                          column_scols=[scol_for(sdf, col)
+                                                        for col in kdf._internal.data_columns])
+        else:
+            return_schema = _infer_return_type(func).tpe
+            require_index_axis = getattr(return_sig, "__origin__", None) == ks.Series
+            require_column_axis = getattr(return_sig, "__origin__", None) == ks.DataFrame
+            if require_index_axis:
+                if axis != 0:
+                    raise TypeError(
+                        "The given function should specify a scalar or a series as its type "
+                        "hints when axis is 0 or 'index'; however, the return type "
+                        "was %s" % return_sig)
+                fields_types = zip(self.columns, [return_schema] * len(self.columns))
+                return_schema = StructType([StructField(c, t) for c, t in fields_types])
+            elif require_column_axis:
+                if axis != 1:
+                    raise TypeError(
+                        "The given function should specify a scalar or a frame as its type "
+                        "hints when axis is 1 or 'column'; however, the return type "
+                        "was %s" % return_sig)
+            else:
+                # any axis is fine.
+                should_return_series = True
+                return_schema = StructType([StructField("0", return_schema)])
+
+            sdf = GroupBy._spark_group_map_apply(
+                self, apply_func, (F.spark_partition_id(),),
+                return_schema, retain_index=False)
+
+            # Otherwise, it loses index.
+            internal = _InternalFrame(sdf=sdf, index_map=None)
+
+        result = DataFrame(internal)
+        if should_return_series:
+            return _col(result)
+        else:
+            return result
+
     def transform(self, func):
         """
         Call ``func`` on self producing a Series with transformed values
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -879,8 +879,10 @@ def apply(self, func):
             kdf = DataFrame(pdf)
             return_schema = kdf._sdf.drop(*HIDDEN_COLUMNS).schema
 
-        sdf = self._spark_group_map_apply(
+        sdf = GroupBy._spark_group_map_apply(
+            self._kdf,
             lambda pdf: pdf.groupby(input_groupnames).apply(func),
+            self._groupkeys_scols,
             return_schema,
             retain_index=should_infer_schema)
 
@@ -938,17 +940,18 @@ def filter(self, func):
         def pandas_filter(pdf):
             return pdf.groupby(groupby_names).filter(func)
 
-        sdf = self._spark_group_map_apply(
-            pandas_filter, data_schema, retain_index=True)
+        sdf = GroupBy._spark_group_map_apply(
+            self._kdf, pandas_filter, self._groupkeys_scols, data_schema, retain_index=True)
         return DataFrame(self._kdf._internal.copy(
             sdf=sdf,
             column_scols=[scol_for(sdf, col) for col in self._kdf._internal.data_columns]))
 
-    def _spark_group_map_apply(self, func, return_schema, retain_index):
-        index_columns = self._kdf._internal.index_columns
-        index_names = self._kdf._internal.index_names
-        data_columns = self._kdf._internal.data_columns
-        column_index = self._kdf._internal.column_index
+    @staticmethod
+    def _spark_group_map_apply(kdf, func, groupkeys_scols, return_schema, retain_index):
+        index_columns = kdf._internal.index_columns
+        index_names = kdf._internal.index_names
+        data_columns = kdf._internal.data_columns
+        column_index = kdf._internal.column_index
 
         def rename_output(pdf):
             # TODO: This logic below was borrowed from `DataFrame.pandas_df` to set the index
@@ -1019,8 +1022,8 @@ def rename_output(pdf):
 
         grouped_map_func = pandas_udf(return_schema, PandasUDFType.GROUPED_MAP)(rename_output)
 
-        sdf = self._kdf._sdf.drop(*HIDDEN_COLUMNS)
-        input_groupkeys = [s for s in self._groupkeys_scols]
+        sdf = kdf._sdf.drop(*HIDDEN_COLUMNS)
+        input_groupkeys = [s for s in groupkeys_scols]
         sdf = sdf.groupby(*input_groupkeys).apply(grouped_map_func)
 
         return sdf
@@ -1637,8 +1640,9 @@ def pandas_transform(pdf):
             if len(pdf) <= limit:
                 return kdf
 
-            sdf = self._spark_group_map_apply(
-                pandas_transform, return_schema, retain_index=True)
+            sdf = GroupBy._spark_group_map_apply(
+                self._kdf, pandas_transform, self._groupkeys_scols,
+                return_schema, retain_index=True)
             # If schema is inferred, we can restore indexes too.
             internal = kdf._internal.copy(sdf=sdf,
                                           column_scols=[scol_for(sdf, col)
@@ -1649,8 +1653,9 @@ def pandas_transform(pdf):
             return_schema = StructType([
                 StructField(c, return_type) for c in data_columns if c not in input_groupnames])
 
-            sdf = self._spark_group_map_apply(
-                pandas_transform, return_schema, retain_index=False)
+            sdf = GroupBy._spark_group_map_apply(
+                self._kdf, pandas_transform, self._groupkeys_scols,
+                return_schema, retain_index=False)
             # Otherwise, it loses index.
             internal = _InternalFrame(sdf=sdf, index_map=None)
 
diff --git a/databricks/koalas/missing/frame.py b/databricks/koalas/missing/frame.py
--- a/databricks/koalas/missing/frame.py
+++ b/databricks/koalas/missing/frame.py
@@ -40,7 +40,6 @@ class _MissingPandasLikeDataFrame(object):
 
     # Functions
     align = unsupported_function('align')
-    apply = unsupported_function('apply')
     asfreq = unsupported_function('asfreq')
     asof = unsupported_function('asof')
     at_time = unsupported_function('at_time')


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -2267,24 +2267,77 @@ def test_transform(self):
                            columns=['a', 'b', 'c'],
                            index=np.random.rand(600))
         kdf = ks.DataFrame(pdf)
+        self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(),
+                       pdf.transform(lambda x: x + 1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(),
+                           pdf.transform(lambda x: x + 1).sort_index())
 
-        def _test(kdf, expected):
-            self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(), expected)
+        with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"):
+            kdf.transform(1)
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')])
+        pdf.columns = columns
+        kdf.columns = columns
+
+        self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(),
+                       pdf.transform(lambda x: x + 1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(),
+                           pdf.transform(lambda x: x + 1).sort_index())
+
+    def test_apply(self):
+        pdf = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6] * 100,
+                            'b': [1., 1., 2., 3., 5., 8.] * 100,
+                            'c': [1, 4, 9, 16, 25, 36] * 100},
+                           columns=['a', 'b', 'c'],
+                           index=np.random.rand(600))
+        kdf = ks.DataFrame(pdf)
 
-            with option_context("compute.shortcut_limit", 500):
-                self.assert_eq(kdf.transform(lambda x: x + 1).sort_index(), expected)
+        self.assert_eq(kdf.apply(lambda x: x + 1).sort_index(),
+                       pdf.apply(lambda x: x + 1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.apply(lambda x: x + 1).sort_index(),
+                           pdf.apply(lambda x: x + 1).sort_index())
 
-        _test(kdf, pdf.transform(lambda x: x + 1).sort_index())
+        # returning a Series
+        self.assert_eq(kdf.apply(lambda x: len(x), axis=1).sort_index(),
+                       pdf.apply(lambda x: len(x), axis=1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.apply(lambda x: len(x), axis=1).sort_index(),
+                           pdf.apply(lambda x: len(x), axis=1).sort_index())
 
         with self.assertRaisesRegex(AssertionError, "the first argument should be a callable"):
-            kdf.transform(1)
+            kdf.apply(1)
+
+        with self.assertRaisesRegex(TypeError, "The given function.*1 or 'column'; however"):
+            def f1(_) -> ks.DataFrame[int]:
+                pass
+            kdf.apply(f1, axis=0)
+
+        with self.assertRaisesRegex(TypeError, "The given function.*0 or 'index'; however"):
+            def f2(_) -> ks.Series[int]:
+                pass
+            kdf.apply(f2, axis=1)
 
         # multi-index columns
         columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')])
         pdf.columns = columns
         kdf.columns = columns
 
-        _test(kdf, pdf.transform(lambda x: x + 1).sort_index())
+        self.assert_eq(kdf.apply(lambda x: x + 1).sort_index(),
+                       pdf.apply(lambda x: x + 1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.apply(lambda x: x + 1).sort_index(),
+                           pdf.apply(lambda x: x + 1).sort_index())
+
+        # returning a Series
+        self.assert_eq(kdf.apply(lambda x: len(x), axis=1).sort_index(),
+                       pdf.apply(lambda x: len(x), axis=1).sort_index())
+        with option_context("compute.shortcut_limit", 500):
+            self.assert_eq(kdf.apply(lambda x: len(x), axis=1).sort_index(),
+                           pdf.apply(lambda x: len(x), axis=1).sort_index())
 
     def test_empty_timestamp(self):
         pdf = pd.DataFrame({'t': [datetime(2019, 1, 1, 0, 0, 0),