如何从PySpark DataFrame的probability向量列提取首个元素

我来帮你解决这个问题！你这里的probability列是PySpark ML模型输出的Vector类型（比如逻辑回归这类分类模型的概率结果），要提取每个向量的第一个元素并添加为新列，有两种常用方法，我给你详细说明：

这是最通用的方法，不管你用哪个Spark版本都能生效：

首先导入需要的模块：

from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

然后定义一个UDF，用来从Vector中提取第一个元素：

# 定义UDF：接收Vector类型，返回第一个元素的浮点值
extract_first_prob = udf(lambda vec: vec[0], DoubleType())

最后把这个UDF应用到你的DataFrame上，生成新列：

# 添加名为first_probability的新列，值为probability列的第一个元素
result_df = df.withColumn("first_probability", extract_first_prob(df.probability))

如果你的Spark版本在3.0及以上，更推荐用原生函数，性能会比UDF更好：

先导入所需的函数：

from pyspark.sql.functions import vector_to_array, element_at

先把Vector类型的列转成数组，再提取第一个元素（注意element_at的索引是从1开始的）：

result_df = df.withColumn("prob_array", vector_to_array(df.probability)) \
              .withColumn("first_probability", element_at("prob_array", 1)) \
              .drop("prob_array")  # 可以删掉中间生成的数组列

验证结果 #

运行完上面的代码后，你可以用result_df.show()查看效果，新列first_probability就会显示每个probability向量的第一个元素，比如你示例里的0.72498853530094...。

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