Pyspark Pandas_UDF erroring with Invalid argument, not a string or column

Question

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.

Timo V · Accepted Answer · 2019-12-07 19:55:19Z

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.

answered Dec 7, 2019 at 19:55

Timo V

133 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Pyspark Pandas_UDF erroring with Invalid argument, not a string or column

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related