PySpark中VectorAssembler的输入数据类型要求及流水线适配咨询

报错根因 #

你复现代码中列a的类型为array<bigint>，未通过VectorAssembler的输入类型校验，因此抛出IllegalArgumentException。

VectorAssembler 支持的合法输入列类型 #

• 单值数值类型：包含ByteType、ShortType、IntegerType、LongType、FloatType、DoubleType，这类列的单个值会直接作为独立特征按顺序拼接进输出向量
• Spark ML标准向量类型（VectorUDT）：即DenseVector、SparseVector对应的列类型，这类列本身承载一组特征，会被整体拼接进输出向量
• 同元素数值数组（仅PySpark 3.0+版本支持）：仅当数组元素全为DoubleType时可直接通过校验，整型（Long/Int/Short等）数组需要提前将元素转为DoubleType才能被识别，你代码中的array<bigint>就是因为元素为整型未转Double，触发了类型校验异常

注：你之前推测输入必须为Vector类型存在认知偏差，单值数值列、DoubleType元素的数值数组均属于合法输入，无需全部转为Vector类型

流水线串联适配规则 #

你设计的流水线各环节输出需要对齐VectorAssembler的输入要求，具体适配要点：

• 前序标准组件输出天然兼容：StringIndexer输出为DoubleType编码列，OneHotEncoder输出为VectorUDT独热向量列，均属于合法输入，无需额外调整
• 自定义UnaryTransformer的返回类型要求：
  • 如果转换器的上游输入是单值数值列，输出需保持为数值类型（推荐统一返回DoubleType，规避不同数值类型的兼容问题），禁止返回Python原生列表或非数值类型
  • 如果转换器的上游输入是VectorUDT向量列，输出必须保持VectorUDT类型，即完成向量元素的数值修改（你需求中“将取值为1的元素乘以10”的逻辑）后，返回同类型的向量对象，不能返回Python列表或原生数组

复现代码修正方案 #

你之前的测试代码可通过两种方式解决类型兼容问题：

1. 将原数组列转为VectorUDT类型

   from pyspark.ml.feature import VectorAssembler
   from pyspark.ml.linalg import Vectors
   # 将原Python列表值转为稠密向量作为列值
   df = spark.createDataFrame([(Vectors.dense([1, 2, 3]), 0, 3)], ["a", "b", "c"])
   vecAssembler = VectorAssembler(outputCol="features", inputCols=["a", "b", "c"])
   vecAssembler.transform(df).show()
   

2. 将原数组列的元素统一转为DoubleType

   from pyspark.ml.feature import VectorAssembler
   from pyspark.sql.functions import transform, col
   from pyspark.sql.types import DoubleType
   df = spark.createDataFrame([([1, 2, 3], 0, 3)], ["a", "b", "c"])
   # 将a列的bigint元素数组逐元素转为DoubleType
   df = df.withColumn("a", transform(col("a"), lambda item: item.cast(DoubleType())))
   vecAssembler = VectorAssembler(outputCol="features", inputCols=["a", "b", "c"])
   vecAssembler.transform(df).show()
   

Vector类型列的常用构造方式 #

• 手动构造：使用pyspark.ml.linalg.Vectors工具类，稠密向量用Vectors.dense(特征值列表)创建，稀疏向量用Vectors.sparse(向量总长度, {索引:对应值})创建
• 数组转向量：PySpark 3.0+可使用内置函数array_to_vector，直接将DoubleType数组列转为VectorUDT类型列
• 多列拼向量：对拆分的多列单值特征，直接通过VectorAssembler拼接为统一向量列

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