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 .. #834 (comment) 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.