如何将PySpark DataFrame的Vector类型列转为浮点数组以匹配指定Schema？

要将Spark DataFrame中的Vector类型列转换为算法要求的ArrayType(FloatType)，同时匹配指定的非空约束Schema，可以按以下步骤操作：

1. 导入依赖模块 #

首先导入PySpark所需的核心库和类型定义：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, FloatType, IntegerType, StructType, StructField
from pyspark.ml.linalg import Vector

2. 定义类型转换UDF #

创建一个用户自定义函数（UDF），将Spark的Vector对象转换为浮点数组：

# 把Vector转为Float类型数组的UDF
vector_to_float_array = udf(
    lambda vec: [float(val) for val in vec.toArray()],
    ArrayType(FloatType())
)

3. 转换DataFrame并匹配目标Schema #

假设你的原始DataFrame名为input_df，执行以下操作完成转换：

# 1. 转换Vector列为数组，并重命名为目标列名"values"
# 2. 调整id列类型，移除原"value"列
temp_df = input_df.withColumn("values", vector_to_float_array(input_df["value"])) \
                  .drop("value") \
                  .withColumn("id", input_df["id"].cast(IntegerType()))

# 强制应用目标Schema，确保非空约束生效
target_schema = StructType([
    StructField("id", IntegerType(), False),
    StructField("values", ArrayType(FloatType()), False)
])

final_df = spark.createDataFrame(temp_df.rdd, schema=target_schema)

关键说明 #

• Vector.toArray()会将稀疏/密集向量统一转为密集数组，适配算法对数组格式的要求；
• 通过spark.createDataFrame重新指定Schema，可以强制设置列的非空约束（nullable=False），无需修改算法的Schema定义；
• 若原始id列存在空值，转换前需先做数据清洗（如过滤或填充），否则会因违反非空约束报错。

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