Repository: databricks/koalas

Issue:
Understanding Groupbyapply
Hello there, firstly thank you for such an amazing package that bridges the gap between Pandas and PySpark. 
I started using koalas approximately 1 week back and everything was intuitive till the time i stumbled upon `koalas.Groupby.Apply`. 

Code:
```
if __name__ == '__main__':

        ks_df = ks.DataFrame(features_data)
        ks_df_info_abt_train = ks_df.groupby(['div_nbr', 'store_nbr']).apply(_koalas_train)
        
        def _koalas_train(frame):
                  out_frame = frame.copy()
                  out_frame = frame['trans_type_value'].sum()
                  return out_frame
```
Here features_data is a  pd.Dataframe.

Output from Koalas.Groupby.Apply:

![Screen Shot 2019-09-27 at 2 27 26 PM](https://user-images.githubusercontent.com/33138515/65793013-f2d26580-e132-11e9-9a39-9e0fecb5bd3f.png)

Output from Pandas.Groupby.Apply:
![Screen Shot 2019-09-27 at 2 30 13 PM](https://user-images.githubusercontent.com/33138515/65793255-65dbdc00-e133-11e9-8cde-d47aed06ce3b.png)


As you can see, the output from pandas Groupby apply is as expected, but the output from Koalas Groupby apply is not right. Could you guid me towards the right direction by pointing out any logical mistake that i might have made or anything else. 
Thank you once again.

Koalas version - 0.18.0 
Pandas version - 0.23.4
PySpark -  2.4.3


Discussion:
I think the root cause is the highlighted line below in `Groupby.apply` ([link](https://github.com/databricks/koalas/blob/1c11cc71f1710325f8ee2911f4f78a57940205c4/databricks/koalas/groupby.py#L819-L828)).

```
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._kdf.head(limit + 1)._to_internal_pandas()
    
    ##### This line is causing the issue #####
    pdf = pdf.groupby(input_groupnames).apply(func)

    kdf = DataFrame(pdf)
    return_schema = kdf._sdf.schema
    if len(pdf) <= limit:
        return kdf
```
After 
`pdf = pdf.groupby(input_groupnames).apply(func)`, `len(pdf)` becomes much smaller than `limit` (which is 1000 by default) and `kdf` created from the first 1001 rows is returned.
Notebook to reproduce the issue.
https://databricks-prod-cloudfront.cloud.databricks.com/public/4027ec902e239c93eaaa8714f173bcfc/4475210719991087/2980686694056275/4764411236566651/latest.html
Thank you @harupy for reproducing the issue and explaining it, any work arounds?
@Devarsh-UTD 
You could use `agg` to take the sum of `trans_type_value` in each group like below.

```
ks_df.groupby(['div_nbr', 'store_nbr']).agg({'trans_type_value': 'sum'})
```
`DataFrameGroupBy.apply` doesn't support a function which returns a series or a scalar yet.

https://koalas.readthedocs.io/en/latest/reference/api/databricks.koalas.groupby.GroupBy.apply.html
> A callable that takes a DataFrame as its first argument, and returns a dataframe.

Thanks for the prompt reply. Since i wanted to post this on git hub i took a simple example of 'sum()', in reality what i wanted to achieve through this function was to groupby ID and then for each group train a separate scikit learn model using pandas_wraps. Any suggestions to achive this ade welcome. Also, this means that this is an actual issue right, nothing wrong on my side in terms of coding. 
Let me test it out. :)
Just to note that i usually return from the function like this - 

Return pd.Dataframe(return_value)
> what i wanted to achieve through this function was to groupby ID and then for each group train a separate scikit learn model using pandas_wraps.

Like this?

```
def train_model(each_group_df):
   model = ModelConstructor(...)
   model.fit(each_group_df)
   # do something with the trained model.
   ...
```
```
def train_model(each_group_df):
   feature_name =['feature_1', 'feature_2']
   df = each_group_df.copy()
   model = ModelConstructor(...)
   model.fit(df[features_name])
   df['model_str'] = pickle.dumps(model)
    df = df[['model_str', 'ID']]
  Return df

ks_df=ks_df.groupby('ID').apply(trainmodel)
```
@Devarsh-UTD

Not sure if this is a proper way to achieve what you want to do, but this works.
```python
import pickle
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression


def train_model(df) -> ks.DataFrame[int, str]:
  model = LinearRegression()
  model.fit(df[['feature_1', 'feature_2']], df['label'])
  return pd.DataFrame({
    'id': df['id'].unique(),
    'model_str': [str(pickle.dumps(model))]
  })

N = 10000
data = {
  'id': np.random.randint(10, size=N),
  'feature_1': np.random.rand(N),
  'feature_2': np.random.rand(N),
  'label': np.random.randint(2, size=N)
}

kdf = ks.DataFrame(data)
kdf.groupby('id').apply(train_model)
```
Well what you are doing here is basically what i did. But the output is not proper as discussed above, which is the main problem.

If somehow we could make the output right , your or my function will work perfectly fine. 
> Well what you are doing here is basically what I did.

Actually, it is a bit different because I put type hint (`ks.DataFrame[int, str]`) on `train_model` to skip schema inference.
Hmmm .. https://github.com/databricks/koalas/issues/834#issuecomment-536175397 this looks actually a potential issue ..  @harupy do you mind if I ask to document this limitation for now? Maybe we should think about a way to fix it properly (e.g., do not allow to return the pandas DataFrame directly as a shortcut but only uses inferred schema).
@HyukjinKwon 
No, not at all. I'll create a PR for this.
Thanks!
@HyukjinKwon 
I just started thinking the current implementation is **HIGHLY** dangerous.

```python
from databricks import koalas as ks
import pandas as pd
import numpy as np

np.random.seed(seed=0)

def train(df):
  return df['c'].mean()

# data larger than shortcut limit
data = {
  'a': np.random.randint(0, 2, 10000),
  'b': np.random.randint(3, 5, 10000),
  'c': np.random.randint(6, 8, 10000)
}

# pandas
pdf = pd.DataFrame(data)
pdf_agg = pdf.groupby(['a', 'b']).apply(train)
>>> pdf_agg

   a  b
0  3  6.491192
   4  6.500404
1  3  6.490500
   4  6.494812

# koalas
kdf = ks.DataFrame(pdf)
kdf_agg = kdf.groupby(['a', 'b']).apply(train)
>>> kdf_agg

      0
a  b	
0  3  6.502092
   4  6.459144
1  3  6.458167
   4  6.496063
```

Here pandas returns correct values, but koalas doesn't because it only uses the first 1001 records. The problem is this code can run without raising any error or warning to users. @Devarsh-UTD 's case was lucky because the difference is obvious and noticeable.
Let me get back to you by tomorrow evening. 
@harupy Ack. we should fix it. We can fix it first to remove the shortcut but only infer its schema (see https://github.com/databricks/koalas/pull/845). Let's fix this one before the next release.

There's no perfect way to infer the schema. Ideally we should specify the type hint properly.
We can increase the size of configuration to infer, or disable by None to throw an error saying type hint should be specified. One alternative is sample and use it to infer its schema.
ok! I'll try!
After removing the if-clause which returns pdf as a shortcut like below,  `test_apply_with_new_dataframe`([code](https://github.com/databricks/koalas/blob/1462291490c4a3e084ece28edaef92721bce14fd/databricks/koalas/tests/test_groupby.py#L598)) starts failing.

```
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._kdf.head(limit + 1)._to_internal_pandas()
    pdf = pdf.groupby(input_groupnames).apply(func)
    kdf = DataFrame(pdf)
    return_schema = kdf._sdf.schema
   
    # remove these two lines to disable shortcut
    # if len(pdf) <= limit:
    #    return kdf
```
Seems like there's inconsistency between `spark_df.groupby(...).apply(...)` and `pandas_df.groupby(...).apply(...)` about handling indexes.

Can you try to fix below?

```diff
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
index ea5ee45..9368e4d 100644
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -824,11 +824,13 @@ class GroupBy(object):
             pdf = pdf.groupby(input_groupnames).apply(func)
             kdf = DataFrame(pdf)
             return_schema = kdf._sdf.schema
-            if len(pdf) <= limit:
-                return kdf
+            # if len(pdf) <= limit:
+            #     return kdf

         sdf = self._spark_group_map_apply(
-            func, return_schema, retain_index=should_infer_schema)
+            lambda pdf: pdf.groupby(input_groupnames).apply(func),
+            return_schema,
+            retain_index=should_infer_schema)

         if should_infer_schema:
             # If schema is inferred, we can restore indexes too.

```

In this way, we will handle `groupby().apply()` identically as pandas'

It can be debugged as below:

```diff
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
index ea5ee45..d8fffce 100644
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -916,6 +916,8 @@ class GroupBy(object):

             pdf = func(pdf)

+            print(pdf)
+
             if retain_index:
                 # If schema should be inferred, we don't restore index. Pandas seems restoring
                 # the index in some cases.
```

**Before:**

```
    column
0     0.0
```

**After:**

```
          column
car_id          
A      0     0.0
```
BTW, it should be safe to do because it's guaranteed to have the same whole grouped data in `func` from `spark_df.groupby(...).apply(func)`.

Say, if we have a Spark DataFrame as below:

```python
>>> df.show()
```
```
+---+---+
|  a|  b|
+---+---+
|  1|  1|
|  1|  2|
|  2|  2|
|  2|  3|
+---+---+
``` 

```python
from pyspark.sql.functions import PandasUDFType, pandas_udf

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def func(pdf):
    print(pdf)
    return pdf

df.groupby("a").apply(func).show()
```

The `func`'s `pdf` becomes, each grouped by `a`.

```
   a  b
0  1  1
1  1  2
```
```
   a  b
0  2  2
1  2  3
```

So, applying `pdf.groupby(...).apply(...)` won't change its output (although it wastes a bit of computation by additional groupby in pandas) but only correct the index as pandas' style:

```python
from pyspark.sql.functions import PandasUDFType, pandas_udf

@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
def func(pdf):
    pdf = pdf.groupby("a").apply(lambda pdf: pdf)
    print(pdf)
    return pdf

df.groupby("a").apply(func).show()
```

```
   a  b
0  1  1
1  1  2
```

```
   a  b
0  2  2
1  2  3
```
@HyukjinKwon Thanks!

Resolving patch:
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -799,7 +799,7 @@ def apply(self, func):
 
         >>> def plus_min(x):
         ...    return x + x.min()
-        >>> g.apply(plus_min)  # doctest: +NORMALIZE_WHITESPACE
+        >>> g.apply(plus_min).sort_index()  # doctest: +NORMALIZE_WHITESPACE
             A  B   C
         0  aa  2   8
         1  aa  3  10
@@ -838,17 +838,16 @@ def apply(self, func):
 
         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._kdf.head(limit + 1)._to_internal_pandas()
+            pdf = self._kdf.head(limit)._to_internal_pandas()
             pdf = pdf.groupby(input_groupnames).apply(func)
             kdf = DataFrame(pdf)
             return_schema = kdf._sdf.schema
-            if len(pdf) <= limit:
-                return kdf
 
         sdf = self._spark_group_map_apply(
-            func, return_schema, retain_index=should_infer_schema)
+            lambda pdf: pdf.groupby(input_groupnames).apply(func),
+            return_schema,
+            retain_index=should_infer_schema)
 
         if should_infer_schema:
             # If schema is inferred, we can restore indexes too.


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
@@ -553,6 +553,9 @@ def test_apply(self):
         self.assert_eq(kdf.groupby(['b'])['a'].apply(lambda x: x).sort_index(),
                        pdf.groupby(['b'])['a'].apply(lambda x: x).sort_index())
 
+        with self.assertRaisesRegex(TypeError, "<class 'int'> object is not callable"):
+            kdf.groupby("b").apply(1)
+
         # multi-index columns
         columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'c')])
         pdf.columns = columns
@@ -563,38 +566,6 @@ def test_apply(self):
         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())
 
-        # Less than 'compute.shortcut_limit' will execute a shortcut
-        # by using collected pandas dataframe directly.
-        # now we set the 'compute.shortcut_limit' as 1000 explicitly
-        set_option('compute.shortcut_limit', 1000)
-        try:
-            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 = koalas.DataFrame(pdf)
-            self.assert_eq(kdf.groupby("b").apply(lambda x: x + 1).sort_index(),
-                           pdf.groupby("b").apply(lambda x: x + 1).sort_index())
-            self.assert_eq(kdf.groupby(['a', 'b']).apply(lambda x: x * x).sort_index(),
-                           pdf.groupby(['a', 'b']).apply(lambda x: x * x).sort_index())
-            self.assert_eq(kdf.groupby(['b'])['a'].apply(lambda x: x).sort_index(),
-                           pdf.groupby(['b'])['a'].apply(lambda x: x).sort_index())
-            with self.assertRaisesRegex(TypeError, "<class 'int'> object is not callable"):
-                kdf.groupby("b").apply(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")).apply(lambda x: x + 1).sort_index(),
-                           pdf.groupby(("x", "b")).apply(lambda x: 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())
-        finally:
-            reset_option('compute.shortcut_limit')
-
     def test_apply_with_new_dataframe(self):
         pdf = pd.DataFrame({
             "timestamp": [0.0, 0.5, 1.0, 0.0, 0.5],
@@ -606,20 +577,28 @@ def test_apply_with_new_dataframe(self):
             kdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
             pdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index())
 
-        set_option('compute.shortcut_limit', 1000)
-        try:
-            # 1000+ records will only infer the schema.
-            pdf = pd.DataFrame({
-                "timestamp": [0.0, 0.5, 1.0, 0.0, 0.5] * 300,
-                "car_id": ['A', 'A', 'A', 'B', 'B'] * 300
-            })
-            kdf = koalas.DataFrame(pdf)
+        self.assert_eq(
+            kdf.groupby('car_id')
+            .apply(lambda df: pd.DataFrame({'mean': [df['timestamp'].mean()]})).sort_index(),
+            pdf.groupby('car_id')
+            .apply(lambda df: pd.DataFrame({"mean": [df['timestamp'].mean()]})).sort_index())
 
-            self.assert_eq(
-                kdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
-                pdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index())
-        finally:
-            reset_option('compute.shortcut_limit')
+        # dataframe with 1000+ records
+        pdf = pd.DataFrame({
+            "timestamp": [0.0, 0.5, 1.0, 0.0, 0.5] * 300,
+            "car_id": ['A', 'A', 'A', 'B', 'B'] * 300
+        })
+        kdf = koalas.DataFrame(pdf)
+
+        self.assert_eq(
+            kdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
+            pdf.groupby('car_id').apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index())
+
+        self.assert_eq(
+            kdf.groupby('car_id')
+            .apply(lambda df: pd.DataFrame({"mean": [df['timestamp'].mean()]})).sort_index(),
+            pdf.groupby('car_id')
+            .apply(lambda df: pd.DataFrame({"mean": [df['timestamp'].mean()]})).sort_index())
 
     def test_transform(self):
         pdf = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6],