4

I created a Pandas UDF, which will input a dataframe, predict and output a dataframe on Primary_Key and 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 - is Input dataframe

Code Runs good when I test it as below:

a = apply_model.func(df)

Output a.dtypes is giving

prediction float64 primary_id int64

When Running the below code:

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

above statement failing with the error:

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.
Pawan Kalyan
  • 51
  • 5

1 Answers1

1

Is df in your last code snippet perhaps a Pandas DataFrame? It should be a PySpark DataFrame for that code to work; Pandas UDFs are to be applied in Spark.

Timo V
  • 13
  • 3