Databricks中如何在PySpark DataFrame上使用predict_proba？

问题原因 #

你通过mlflow.pyfunc.load_model加载的是本地单进程模型，仅支持Pandas DataFrame。PySpark DataFrame是分布式数据集，无法直接调用这类模型的方法，需要把模型包装成Spark UDF（用户自定义函数），才能实现分布式推理。

实现步骤 #

1. 加载模型并创建Spark UDF #

不用直接把模型加载为pyfunc对象，而是用mlflow.pyfunc.spark_udf直接从模型仓库或Run URI创建支持Spark的UDF，同时指定返回类型（因为predict_proba返回概率数组，要对应Spark的数据类型）。

示例代码：

from pyspark.sql.types import ArrayType, DoubleType
import mlflow

# 方式1：用注册的模型名称和版本创建UDF
predict_proba_udf = mlflow.pyfunc.spark_udf(
    spark=spark,
    model_uri=f"models:/{model_name}/{model_version}",
    result_type=ArrayType(DoubleType())  # 匹配predict_proba返回的概率数组格式
)

# 方式2：如果还是想用训练Run的URI
predict_proba_udf = mlflow.pyfunc.spark_udf(
    spark=spark,
    model_uri=model_uri,
    result_type=ArrayType(DoubleType())
)

2. 在PySpark DataFrame上调用UDF #

创建好UDF后，用withColumn就能在PySpark DataFrame上生成概率列：

# 假设你的PySpark数据集是spark_df，包含模型需要的所有特征列
spark_df_with_proba = spark_df.withColumn(
    "probabilities",
    predict_proba_udf(*spark_df.columns)  # 传入所有特征列，也可指定具体列名列表
)

# 查看结果
spark_df_with_proba.show()

3. 注意事项 #

• 要保证模型的输入特征列和训练时完全一致，UDF会自动把PySpark列映射成模型需要的格式。
• 如果是二分类任务，predict_proba可能返回单个概率值，这时候把result_type改成DoubleType()即可；多分类任务用ArrayType(DoubleType())对应多类别概率数组。
• 用Databricks Runtime的话，AutoML生成的模型已经自带必要依赖，不需要额外安装。

小数据量临时方案（不推荐大数据场景） #

如果你的数据集很小，也可以先转成Pandas DataFrame推理，再转回Spark，但这种方式在大数据量下会性能不足：

pandas_df = spark_df.toPandas()
pandas_df["probabilities"] = model.predict_proba(pandas_df)
spark_df_with_proba = spark.createDataFrame(pandas_df)

内容的提问来源于stack exchange，提问作者cgu