Performance for pyspark dataframe is very slow after using a @pandas_udf

RRO
Databricks Partner

Hello,

I am currently working on a time series forecasting with FBProphet. Since I have data with many time series groups (~3000) I use a @pandas_udf to parallelize the training.

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def forecast_netprofit(prophtrain):
 
     ... 
 
     return results_pd
 
 
time_series_id_column_names = ['Grp1', 'Grp2', 'Grp3']
 
results = (prophtrain
           .groupby(time_series_id_column_names)
           .apply(forecast_netprofit)
          )

Now every time I want to display or do some operations on the results dataframe the performance is really low. For example: Just to display the first 1000 rows takes around 6min.

Is there a reason why the performance of the results is so slow and can I fix that somehow?