在Databricks中用PySpark提取VectorUDT类型概率列的首个值

你的probability列是Spark MLlib特有的VectorUDT向量类型，不属于常规的Struct/Array/MAP类型，所以之前的扩展方式不适用，以下是两种可行的解决方法：

方法1：使用element_at函数（Spark 3.0及以上版本） #

Spark 3.0开始支持用element_at直接操作VectorUDT类型，通过1-based索引提取第一个元素：

from pyspark.sql.functions import element_at

# 添加新列存储预测正确的概率
df = df.withColumn("correct_pred_prob", element_at("probability", 1))

方法2：自定义UDF（兼容低版本Spark） #

如果你的Spark版本低于3.0，可通过自定义UDF提取向量的第一个元素：

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

# 定义提取向量第一个元素的函数
def extract_first_vec_element(vector):
    return vector[0] if vector is not None else None

# 注册UDF
extract_first_udf = udf(extract_first_vec_element, DoubleType())

# 生成新列
df = df.withColumn("correct_pred_prob", extract_first_udf("probability"))

报错原因说明 #

• col("probability.*")仅适用于Struct类型列，而VectorUDT是MLlib专用的向量类型，无法用Struct的扩展语法处理。
• 直接尝试提取vectorType字段失败，是因为VectorUDT不属于常规的复杂类型（STRUCT/ARRAY/MAP）范畴，需要用MLlib相关的操作方法或上述两种方式处理。

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