如何在PySpark DataFrame中使用.pkl格式的sklearn模型？

完全可以实现，目前生产环境常用两种方案，首推PySpark 3.0+支持的pandas UDF方案，推理效率更高。

方案1：pandas UDF批量推理（生产推荐） #

核心逻辑是将模型广播到所有worker节点，再通过pandas UDF按批次处理分布式数据，避免单条推理的序列化开销，性能接近本地sklearn推理效率。
实现代码示例：

import joblib
import pandas as pd
from pyspark.sql.functions import pandas_udf, col
from pyspark.sql.types import ArrayType, FloatType, IntegerType

# 1. 驱动节点加载本地pkl模型，广播到所有worker节点，避免重复加载浪费资源
local_model = joblib.load('model.pkl')
broadcast_model = sc.broadcast(local_model)

# 2. 定义预测概率的pandas UDF，输入为批次特征，输出分类概率数组
# 替换为你训练时的特征列名，顺序必须和训练时完全一致
feature_cols = ["特征列1", "特征列2", "..."] 
@pandas_udf(ArrayType(FloatType()))
def predict_proba_udf(*feature_batch: pd.Series) -> pd.Series:
    model = broadcast_model.value
    # 把批次特征转为和训练时结构一致的DataFrame
    batch_df = pd.concat(feature_batch, axis=1)
    proba = model.predict_proba(batch_df)
    return pd.Series(proba.tolist())

# 3. 调用UDF对PySpark DataFrame做推理
# 假设你的PySpark原始DataFrame名为spark_df
result_df = spark_df.withColumn("predict_proba", predict_proba_udf(*[col(c) for c in feature_cols]))

如果只需要输出分类标签，把predict_proba替换为predict，同时把UDF返回类型修改为对应的数据类型（如IntegerType()）即可。

方案2：普通UDF（仅适合小数据量测试） #

如果是PySpark 3.0以下的旧版本，也可以用普通UDF实现，但单条推理序列化开销极大，生产环境不推荐使用。

必看注意事项 #

• 所有worker节点的Python依赖版本必须和模型训练时的环境完全一致，包括scikit-learn、lightgbm、joblib等相关库，否则会出现模型加载失败、预测结果异常等问题
• 特征列的顺序、类型必须和模型训练时的输入完全匹配，否则会导致推理错误
• 大模型广播前需确认worker节点的内存配置，避免出现OOM问题

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