pandas - Pyspark Pandas_UDF 错误，参数无效，而不是字符串或列

Question

我创建了一个 Pandas UDF，它将输入一个数据帧，在 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.

Answer 1

是 df在您的最后一个代码片段中可能是 Pandas DataFrame？该代码应该是一个 PySpark DataFrame 才能工作； Pandas UDF 将应用于 Spark。