Koalas: Understanding Groupbyapply

[devarshml]

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:

Output from Pandas.Groupby.Apply:

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

[harupy]

I think the root cause is the highlighted line below in Groupby.apply (link).

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.

[harupy]

Notebook to reproduce the issue.
https://databricks-prod-cloudfront.cloud.databricks.com/public/4027ec902e239c93eaaa8714f173bcfc/4475210719991087/2980686694056275/4764411236566651/latest.html

[devarshml]

Thank you @harupy for reproducing the issue and explaining it, any work arounds?

[harupy]

@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'})

[harupy]

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.

[devarshml]

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.

[devarshml]

Let me test it out. :)

[devarshml]

Just to note that i usually return from the function like this -

Return pd.Dataframe(return_value)

[harupy]

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.
   ...

[devarshml]

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)

[harupy]

@Devarsh-UTD

Not sure if this is a proper way to achieve what you want to do, but this works.

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)

[devarshml]

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.

[harupy]

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.

[HyukjinKwon]

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).

[harupy]

@HyukjinKwon
No, not at all. I'll create a PR for this.

[HyukjinKwon]

Thanks!

[harupy]

@HyukjinKwon
I just started thinking the current implementation is HIGHLY dangerous.

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.

[devarshml]

Let me get back to you by tomorrow evening.

[HyukjinKwon]

@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.

[harupy]

ok! I'll try!

[harupy]

After removing the if-clause which returns pdf as a shortcut like below, test_apply_with_new_dataframe(code) 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

[HyukjinKwon]

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 --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 --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:

0     0.0

After:

car_id          
A      0     0.0

[HyukjinKwon]

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:

>>> df.show()

|  a|  b|
+---+---+
|  1|  1|
|  1|  2|
|  2|  2|
|  2|  3|
+---+---+



md5-2c5ea977f5806aace5a5195eaeb7ebb9




The `func`'s `pdf` becomes, each grouped by `a`.



md5-3d6a12b57f8e62573e01ea9e873223ff

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

[harupy]

@HyukjinKwon Thanks!

[devarshml]

Thank you guys.

[HyukjinKwon]

Thank YOU guys :-)

[harupy]

Thank you guys :)