PySpark中如何实现类似sklearn的predict_proba概率预测？

问题背景 #

刚接触PySpark，因为数据量庞大，要把原本用Python实现的预测逻辑迁移到PySpark上。目前已经完成以下步骤：

• 将预处理后的Pandas训练DataFrame转换为PySpark DataFrame；
• 选定训练特征列，使用VectorAssembler生成特征向量；
• 训练完成RandomForestClassifier模型，并在训练集上完成评估；

现在需要把模型应用到测试DataFrame（df_to_predict）上，已经通过select匹配了训练特征列，要实现类似sklearn中model.predict_proba(df_to_predict)[::,1]的功能——也就是获取类别1的预测概率，求具体实现方法。

原参考代码片段：

# 匹配训练特征列
df_to_predict = df_to_predict.select(training_columns)
# 需转换为PySpark实现的sklearn代码
df_to_predict["Predicted_Y_Probability"] = model.predict_proba(df_to_predict)[::, 1]

实现步骤与代码 #

PySpark的机器学习API和sklearn有差异，没有直接的predict_proba方法，需要通过以下步骤实现：

1. 对测试集生成特征向量
   和训练流程一致，必须用训练时的同一个VectorAssembler把测试集的原始特征列转换成特征向量列（默认命名为features），因为PySpark模型是基于特征向量做预测的。

2. 用模型生成预测结果
   调用模型的transform方法，会返回包含预测类别、原始预测值、概率向量的DataFrame，其中probability列是存储各类别概率的Vector类型数据。

3. 提取类别1的概率
   可以用element_at或getItem函数从probability向量中取出对应类别1的概率值，注意两者的索引计数规则不同：

   • element_at从1开始计数，类别1对应第2个元素；
   • getItem从0开始计数，类别1对应索引1。

完整代码示例

# 假设训练时定义的VectorAssembler实例名为assembler
# 1. 给测试集生成特征向量列
df_to_predict = assembler.transform(df_to_predict)

# 2. 模型预测，得到包含概率向量的DataFrame
predictions_df = model.transform(df_to_predict)

# 3. 提取类别1的概率并命名为目标列
from pyspark.sql.functions import element_at

# 使用element_at的写法（从1计数）
result_df = predictions_df.withColumn(
    "Predicted_Y_Probability",
    element_at(predictions_df["probability"], 2)
)

# 或者使用getItem的写法（从0计数）
# from pyspark.sql.functions import col
# result_df = predictions_df.withColumn(
#     "Predicted_Y_Probability",
#     col("probability").getItem(1)
# )

# 按需保留需要的列
result_df = result_df.select(*training_columns, "Predicted_Y_Probability")

注意事项 #

• 必须复用训练时的VectorAssembler，保证特征顺序和处理逻辑和训练集一致，否则会导致预测错误；
• 如果不确定类别顺序，可以通过model.labels查看模型的类别列表，比如返回[0, 1]，则索引1对应类别1；
• probability向量的长度等于类别总数，多分类场景也可以用同样的方法提取对应类别的概率。

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