Repository: databricks/koalas

Issue:
Unexpected Koalas groupby.apply groupby.agg behavior
I found some unexpected groupby.agg or groupby.apply behavior.
Here's a reproducible example.

```
test = pd.DataFrame({'secid': {0: 125047, 1: 125047, 2: 125047, 3: 125047, 4: 125047, 5: 125047},
 'date': {0: pd.Timestamp('2018-03-05 00:00:00'),
  1: pd.Timestamp('2018-03-05 00:00:00'),
  2: pd.Timestamp('2018-03-05 00:00:00'),
  3: pd.Timestamp('2018-03-05 00:00:00'),
  4: pd.Timestamp('2018-03-05 00:00:00'),
  5: pd.Timestamp('2018-03-05 00:00:00')},
 'days': {0: 182, 1: 182, 2: 182, 3: 182, 4: 182, 5: 182},
 'delta': {0: -20, 1: -15, 2: -10, 3: 10, 4: 15, 5: 20},
 'impl_volatility': {0: 0.251335,
  1: 0.321205,
  2: 0.433905,
  3: 0.327283,
  4: 0.307678,
  5: 0.28428},
'abs_delta': {0: 20, 1: 15, 2: 10, 3: 10, 4: 15, 5: 20}})


def my_divider(x) -> float:
    try:
        return x.iloc[0]/x.iloc[1]
    except IndexError: 
        pass


test
secid | date | days | delta | impl_volatility | abs_delta
-- | -- | -- | -- | -- | --
125047 | 2018-03-05 | 182 | -20 | 0.251335 | 20
125047 | 2018-03-05 | 182 | -15 | 0.321205 | 15
125047 | 2018-03-05 | 182 | -10 | 0.433905 | 10
125047 | 2018-03-05 | 182 | 10 | 0.327283 | 10
125047 | 2018-03-05 | 182 | 15 | 0.307678 | 15
125047 | 2018-03-05 | 182 | 20 | 0.284280 | 20


```
I was trying to group by multiple columns and calculate the ratio using ```my_divider``` function. 

When I run below codes in pandas, it works as expected.
```
test.groupby(['secid','date','days','abs_delta'])['impl_volatility'].apply(sum)
>>>
secid   date        days  abs_delta
125047  2018-03-05  182   10           0.761188
                          15           0.628883
                          20           0.535615
Name: impl_volatility, dtype: float64

test.groupby(['secid','date','days','abs_delta'])['impl_volatility'].apply(my_divider)
>>>secid   date        days  abs_delta
125047  2018-03-05  182   10           1.325779
                          15           1.043965
                          20           0.884111
Name: impl_volatility, dtype: float64

test.groupby(['secid','date','days','abs_delta']).agg(vol_ratio = ('impl_volatility',my_divider))
>>>
secid	date	days	abs_delta	vol_ratio
125047	2018-03-05	182	10	1.325779
                                                15	1.043965
                                                20	0.884111
```


However, in Koalas dataframe

```
test2 = ks.from_pandas(test)
test2.groupby(['secid','date','days','abs_delta'])['impl_volatility'].apply(sum)
>>>
0    0.535615
1    0.628883
2    0.761188
3    0.761188
4    0.628883
5    0.535615
# six outputs (three * 2 duplicates)


test2.groupby(['secid','date','days','abs_delta']).agg(mycolumn = ('impl_volatility','sum'))
>>>
secid	date	days	abs_delta	mycolumn
125047	2018-03-05	182	15	0.628883
                                                20	0.535615
                                                10	0.761188

```
.agg works with 'sum' function as expected but .agg does not understand function defined above.
```
test2.groupby(['secid','date','days','abs_delta']).agg(mycolumn = ('impl_volatility','my_divider'))
>>>
Py4JJavaError: An error occurred while calling o3520.agg.
: org.apache.spark.sql.AnalysisException: Undefined function: 'my_divider'. This function is neither a registered temporary function nor a permanent function registered in the database 'default'
AnalysisException: "Undefined function: 'my_divider'. This function is neither a registered temporary function nor a permanent function registered in the database 'default'.; line 1 pos 0"
```
So I tried registering function, but
```
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('double') 
def my_divider_udf(x):
    try:
        return x.iloc[0]/x.iloc[1]
    except:
        pass

spark.udf.register('my_udf',my_divider_udf)

```
I get 
```

test2.groupby(['secid','date','days','abs_delta']).agg(pc_ratio=('impl_volatility','my_udf'))

Py4JJavaError: An error occurred while calling o2588.agg.
: org.apache.spark.sql.AnalysisException: expression '`impl_volatility`' is neither present in the group by, nor is it an aggregate function. Add to group by or wrap in first() (or first_value) if you don't care which value you get.;;
```




Where .apply throws different type of error.
```
test2.groupby(['secid','date','days','abs_delta'])['impl_volatility'].apply(my_divider)
>>>
Py4JJavaError: An error occurred while calling o1084.collectToPython.
: org.apache.spark.SparkException: Job aborted due to stage failure: Task 94 in stage 22.0 failed 1 times, most recent failure: Lost task 94.0 in stage 22.0 (TID 167, localhost, executor driver): org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 377, in main
    process()
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 372, in process
    serializer.dump_stream(func(split_index, iterator), outfile)
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pyspark/python/lib/pyspark.zip/pyspark/serializers.py", line 290, in dump_stream
    for series in iterator:
  File "<string>", line 1, in <lambda>
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 113, in wrapped
    result = f(pd.concat(value_series, axis=1))
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pyspark/python/lib/pyspark.zip/pyspark/util.py", line 99, in wrapper
    return f(*args, **kwargs)
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/databricks/koalas/groupby.py", line 978, in rename_output
    pdf = func(pdf)
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/databricks/koalas/groupby.py", line 1628, in pandas_transform
    return pdf.transform(func)
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pandas/core/frame.py", line 6735, in transform
    return super().transform(func, *args, **kwargs)
  File "/Users/matthewson/miniconda3/lib/python3.7/site-packages/pandas/core/generic.py", line 10814, in transform
    raise ValueError("transforms cannot produce " "aggregated results")
ValueError: transforms cannot produce aggregated results

```

I'm not sure why I'm seeing last two errors in Koalas with my limited experience. I feel I might be doing something wrong here.


Discussion:
@matthewgson, thanks for reporting an issue. It was a horrible mistake .. The fix is on the way

Resolving patch:
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -26,6 +26,7 @@
 from typing import Any, List, Tuple, Union
 
 import numpy as np
+import pandas as pd
 from pandas._libs.parsers import is_datetime64_dtype
 from pandas.core.dtypes.common import is_datetime64tz_dtype
 
@@ -814,8 +815,8 @@ def apply(self, func):
 
              To avoid this, specify return type in ``func``, for instance, as below:
 
-             >>> def pandas_div_sum(x) -> ks.DataFrame[float, float]:
-             ...    return x[['B', 'C']] / x[['B', 'C']].sum()
+             >>> def pandas_div(x) -> ks.DataFrame[float, float]:
+             ...    return x[['B', 'C']] / x[['B', 'C']]
 
              If the return type is specified, the output column names become
              `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
@@ -832,11 +833,12 @@ def apply(self, func):
 
         Returns
         -------
-        applied : DataFrame
+        applied : DataFrame or Series
 
         See Also
         --------
         aggregate : Apply aggregate function to the GroupBy object.
+        DataFrame.apply : Apply a function to a DataFrame.
         Series.apply : Apply a function to a Series.
 
         Examples
@@ -853,24 +855,6 @@ def apply(self, func):
         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
-
-        You can omit the type hint and let Koalas infer its type.
-
         >>> def plus_min(x):
         ...    return x + x.min()
         >>> g.apply(plus_min).sort_index()  # doctest: +NORMALIZE_WHITESPACE
@@ -879,11 +863,28 @@ def apply(self, func):
         1  aa  3  10
         2  bb  6  10
 
+        You can specify the type hint and prevent schema inference for better performance.
+
+        >>> def pandas_div(x) -> ks.DataFrame[float, float]:
+        ...    return x[['B', 'C']] / x[['B', 'C']]
+        >>> g.apply(pandas_div).sort_index()  # doctest: +NORMALIZE_WHITESPACE
+            c0   c1
+        0  1.0  1.0
+        1  1.0  1.0
+        2  1.0  1.0
+
+        >>> def pandas_length(x) -> int:
+        ...    return len(x)
+        >>> g.apply(pandas_length).sort_index()  # doctest: +NORMALIZE_WHITESPACE
+        0    1
+        1    2
+        Name: 0, dtype: int32
+
         In case of Series, it works as below.
 
         >>> def plus_max(x) -> ks.Series[np.int]:
         ...    return x + x.max()
-        >>> df.B.groupby(df.A).apply(plus_max)
+        >>> df.B.groupby(df.A).apply(plus_max).sort_index()
         0    6
         1    3
         2    4
@@ -891,36 +892,81 @@ def apply(self, func):
 
         >>> def plus_min(x):
         ...    return x + x.min()
-        >>> df.B.groupby(df.A).apply(plus_min)
+        >>> df.B.groupby(df.A).apply(plus_min).sort_index()
         0    2
         1    3
         2    6
         Name: B, dtype: int64
+
+        You can also return a scalar value as a aggregated value of the group:
+
+        >>> def plus_max(x) -> np.int:
+        ...    return len(x)
+        >>> df.B.groupby(df.A).apply(plus_max).sort_index()
+        0    1
+        1    2
+        Name: B, dtype: int32
         """
         if not isinstance(func, Callable):
             raise TypeError("%s object is not callable" % type(func))
 
         spec = inspect.getfullargspec(func)
         return_sig = spec.annotations.get("return", None)
-        if return_sig is None:
-            return_schema = None  # schema will inferred.
-        else:
-            return_schema = _infer_return_type(func).tpe
-
-        should_infer_schema = return_schema is None
+        should_infer_schema = return_sig is None
         input_groupnames = [s.name for s in self._groupkeys]
 
+        should_return_series = False
+        is_series_groupby = isinstance(self, SeriesGroupBy)
+        if is_series_groupby:
+            name = self._kser.name
+
         if should_infer_schema:
             # Here we execute with the first 1000 to get the return type.
             limit = get_option("compute.shortcut_limit")
-            pdf = self._kdf.head(limit)._to_internal_pandas()
-            pdf = pdf.groupby(input_groupnames).apply(func)
-            kdf = DataFrame(pdf)
+            pdf = self._kdf.head(limit + 1)._to_internal_pandas()
+            if is_series_groupby:
+                pser_or_pdf = pdf.groupby(input_groupnames)[name].apply(func)
+            else:
+                pser_or_pdf = pdf.groupby(input_groupnames).apply(func)
+            kser_or_kdf = ks.from_pandas(pser_or_pdf)
+            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._sdf.drop(*HIDDEN_COLUMNS).schema
+        else:
+            if not is_series_groupby and getattr(return_sig, "__origin__", None) == ks.Series:
+                raise TypeError(
+                    "Series as a return type hint at frame groupby is not supported "
+                    "currently; however got [%s]. Use DataFrame type hint instead." % return_sig
+                )
+
+            return_schema = _infer_return_type(func).tpe
+            if not isinstance(return_schema, StructType):
+                should_return_series = True
+                if is_series_groupby:
+                    return_schema = StructType([StructField(name, return_schema)])
+                else:
+                    return_schema = StructType([StructField("0", return_schema)])
+
+        def pandas_groupby_apply(pdf):
+            if is_series_groupby:
+                pdf_or_ser = pdf.groupby(input_groupnames)[name].apply(func)
+            else:
+                pdf_or_ser = pdf.groupby(input_groupnames).apply(func)
+
+            if not isinstance(pdf_or_ser, pd.DataFrame):
+                return pd.DataFrame(pdf_or_ser)
+            else:
+                return pdf_or_ser
 
         sdf = GroupBy._spark_group_map_apply(
             self._kdf,
-            lambda pdf: pdf.groupby(input_groupnames).apply(func),
+            pandas_groupby_apply,
             self._groupkeys_scols,
             return_schema,
             retain_index=should_infer_schema,
@@ -932,7 +978,11 @@ def apply(self, func):
         else:
             # Otherwise, it loses index.
             internal = _InternalFrame(spark_frame=sdf, index_map=None)
-        return DataFrame(internal)
+
+        if should_return_series:
+            return _col(DataFrame(internal))
+        else:
+            return DataFrame(internal)
 
     # TODO: implement 'dropna' parameter
     def filter(self, func):
@@ -2096,11 +2146,6 @@ def agg(self, *args, **kwargs):
     def aggregate(self, *args, **kwargs):
         return _MissingPandasLikeSeriesGroupBy.aggregate(self, *args, **kwargs)
 
-    def apply(self, func):
-        return _col(super(SeriesGroupBy, self).transform(func))
-
-    apply.__doc__ = GroupBy.transform.__doc__
-
     def transform(self, func):
         return _col(super(SeriesGroupBy, self).transform(func))
 


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
@@ -1082,8 +1082,8 @@ def test_apply(self):
             pdf.groupby(["a", "b"]).apply(lambda x: x * x).sort_index(),
         )
         self.assert_eq(
-            kdf.groupby(["b"])["c"].apply(lambda x: x).sort_index(),
-            pdf.groupby(["b"])["c"].apply(lambda x: x).sort_index(),
+            kdf.groupby(["b"])["c"].apply(lambda x: 1).sort_index(),
+            pdf.groupby(["b"])["c"].apply(lambda x: 1).sort_index(),
         )
 
         with self.assertRaisesRegex(TypeError, "<class 'int'> object is not callable"):
@@ -1095,14 +1095,25 @@ def test_apply(self):
         kdf.columns = columns
 
         self.assert_eq(
-            kdf.groupby(("x", "b")).apply(lambda x: x + 1).sort_index(),
-            pdf.groupby(("x", "b")).apply(lambda x: x + 1).sort_index(),
+            kdf.groupby(("x", "b")).apply(lambda x: 1).sort_index(),
+            pdf.groupby(("x", "b")).apply(lambda x: 1).sort_index(),
         )
         self.assert_eq(
             kdf.groupby([("x", "a"), ("x", "b")]).apply(lambda x: x * x).sort_index(),
             pdf.groupby([("x", "a"), ("x", "b")]).apply(lambda x: x * x).sort_index(),
         )
 
+    def test_apply_without_shortcut(self):
+        with option_context("compute.shortcut_limit", 0):
+            self.test_apply()
+
+    def test_apply_negative(self):
+        def func(_) -> ks.Series[int]:
+            return pd.Series([1])
+
+        with self.assertRaisesRegex(TypeError, "Series as a return type hint at frame groupby"):
+            ks.range(10).groupby("id").apply(func)
+
     def test_apply_with_new_dataframe(self):
         pdf = pd.DataFrame(
             {"timestamp": [0.0, 0.5, 1.0, 0.0, 0.5], "car_id": ["A", "A", "A", "B", "B"]}
@@ -1146,6 +1157,10 @@ def test_apply_with_new_dataframe(self):
             .sort_index(),
         )
 
+    def test_apply_with_new_dataframe_without_shortcut(self):
+        with option_context("compute.shortcut_limit", 0):
+            self.test_apply_with_new_dataframe()
+
     def test_transform(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]},
@@ -1179,44 +1194,9 @@ def test_transform(self):
             pdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x * x).sort_index(),
         )
 
-        with option_context("compute.shortcut_limit", 1000):
-            pdf = pd.DataFrame(
-                {
-                    "a": [1, 2, 3, 4, 5, 6] * 300,
-                    "b": [1, 1, 2, 3, 5, 8] * 300,
-                    "c": [1, 4, 9, 16, 25, 36] * 300,
-                },
-                columns=["a", "b", "c"],
-            )
-            kdf = ks.from_pandas(pdf)
-            self.assert_eq(
-                kdf.groupby("b").transform(lambda x: x + 1).sort_index(),
-                pdf.groupby("b").transform(lambda x: x + 1).sort_index(),
-            )
-            self.assert_eq(
-                kdf.groupby(["a", "b"]).transform(lambda x: x * x).sort_index(),
-                pdf.groupby(["a", "b"]).transform(lambda x: x * x).sort_index(),
-            )
-            self.assert_eq(
-                kdf.groupby(["b"])["a"].transform(lambda x: x).sort_index(),
-                pdf.groupby(["b"])["a"].transform(lambda x: x).sort_index(),
-            )
-            with self.assertRaisesRegex(TypeError, "<class 'int'> object is not callable"):
-                kdf.groupby("b").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.groupby(("x", "b")).transform(lambda x: x + 1).sort_index(),
-                pdf.groupby(("x", "b")).transform(lambda x: x + 1).sort_index(),
-            )
-            self.assert_eq(
-                kdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x * x).sort_index(),
-                pdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x * x).sort_index(),
-            )
+    def test_transform_without_shortcut(self):
+        with option_context("compute.shortcut_limit", 0):
+            self.test_transform()
 
     def test_filter(self):
         pdf = pd.DataFrame(