PySpark实现投票分类器封装UDF出现PicklingError问题咨询

错误原因 #

你遇到的PicklingError: can't pickle dict_keys objects是两个核心问题共同导致的：

• PySpark UDF在分发到Executor执行前，需要将函数及闭包内所有依赖对象序列化，你把PySpark Pipeline模型（S1、S2）放在了UDF的默认参数里，这类Spark模型底层绑定了Driver端的Java对象引用，内部包含大量无法被普通pickle序列化的内置结构（比如dict_keys）。
• 你的原代码还存在两个逻辑错误：函数内引用了未定义的变量X_list，且你定义的label_encoder是普通列表，不存在inverse_transform方法，即使解决序列化问题也无法正常运行。

最优解决方案（原生Spark算子实现，无UDF） #

不建议把Spark模型放到UDF中执行，Spark模型本身就是分布式计算的，直接逐个模型预测后合并结果即可，性能远高于UDF方案：

from pyspark.sql.functions import element_at, greatest, when, coalesce

estimators1 = [S1, S2]
w = [1, 1]
label_list = [0, 1, 2] # 原始类别标签列表
qa = featurizer.transform(test)

# 1. 每个模型单独预测，输出加权后的类别概率
for model_idx, clf in enumerate(estimators1):
    qa = clf.transform(qa)
    for label_idx, label in enumerate(label_list):
        # 提取单类别概率乘权重，Spark向量索引从1开始
        qa = qa.withColumn(f"wprob_{model_idx}_{label_idx}", 
                           element_at("probability", label_idx + 1) * w[model_idx])
    # 清理单模型输出的冗余列，避免列冲突
    qa = qa.drop("probability", "prediction", "rawPrediction")

# 2. 计算每个类别的总加权概率
for label_idx, label in enumerate(label_list):
    total_prob = sum(qa[f"wprob_{m_idx}_{label_idx}"] for m_idx in range(len(estimators1)))
    qa = qa.withColumn(f"total_{label_idx}", total_prob)

# 3. 取概率最高的类别作为投票结果，映射回原始标签
total_cols = [f"total_{i}" for i in range(len(label_list))]
qa = qa.withColumn("max_total", greatest(*total_cols))
for label_idx, label in enumerate(label_list):
    qa = qa.withColumn(f"tmp_pred_{label_idx}", 
                       when(qa[f"total_{label_idx}"] == qa.max_total, label).otherwise(None))
qa = qa.withColumn("predictedlabel", coalesce(*[f"tmp_pred_{i}" for i in range(len(label_list))]))

# 清理中间过程列（可选）
drop_cols = [c for c in qa.columns if c.startswith(("wprob_", "total_", "tmp_pred_"))] + ["max_total"]
qa = qa.drop(*drop_cols)

qa.show()

备选方案（UDF实现，仅小数据场景使用） #

如果必须使用UDF实现，需要将Spark模型的参数提取出来，用纯Python/NumPy实现预测逻辑，不要在UDF闭包中引入任何Spark模型对象：

1. 提前从训练好的S1、S2中提取模型参数（比如逻辑回归的系数、截距，朴素贝叶斯的类别先验、特征概率等）
2. 在UDF中用NumPy手动实现概率计算逻辑，所有依赖对象都是可序列化的Python原生/NumPy对象

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