Ошибка Pyspark Pandas_UDF с неверным аргументом, а не строкой или столбцом

Question

Я создал UDF Pandas, который будет вводить фрейм данных, прогнозировать и выводить фрейм данных для Primary_Key и Predictions.

schema = StructType([StructField('primary_id', IntegerType()),
                     StructField('prediction', FloatType())])

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def apply_model(sample_df):  
    # run the model on the partitioned data set   
    ids = sample_df['primary_id']
    x_train = sample_df.drop(['primary_id', 'partition_id'], axis = 1)
    pred = model_broadcast.value.predict_proba(x_train) 

    return pd.DataFrame({'primary_id': ids, 'prediction': pred[:,1]})

sample_df - это входной фрейм данных

Код Работает хорошо, когда я тестирую его, как показано ниже:

a = apply_model.func(df)

Вывод a.dtypes дает

прогноз float64 primary_id int64

При запуске приведенный ниже код:

results = df.groupby('partition_id').apply(apply_model)

приведенное выше утверждение завершилось ошибкой:

TypeError: Invalid argument, not a string or column:
[26 rows x 32 columns] of type <class 'pandas.core.frame.DataFrame'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.