如何判断PySpark DataFrame列是否为Vector类型及相关报错处理

错误根因梳理 #

1. 类型对比返回False的原因：predictions.schema["probability"].dataType返回的是VectorUDT类的实例对象，和字符串'VectorUDT'数据类型不一致，因此等值判断返回False。
2. TypeError: unhashable type: 'Column'报错的原因：Spark DataFrame的列类型是全局固定的元数据，dict(predictions.dtypes)['probability'] == 'vector'是Driver端执行的Python布尔值，而when函数要求传入行级的列表达式，你把全局判断逻辑写在行级运算中，类型不匹配触发报错。

正确实现方式 #

场景1：全局判断列类型（99%的适用场景） #

Spark DataFrame的列是强类型，整列数据类型统一，不需要逐行判断，直接在Driver端完成类型校验后再生成列即可：

from pyspark.ml.linalg import VectorUDT
from pyspark.sql import functions as F

# 方法1：用类型实例判断，准确性更高
is_vector_col = isinstance(predictions.schema["probability"].dataType, VectorUDT)
# 方法2：也可以用dtypes判断，写法更简单
# is_vector_col = dict(predictions.dtypes)['probability'] == 'vector'

predictions = predictions.withColumn('test', F.lit(1) if is_vector_col else F.lit(0))

场景2：行级判断数据类型（极少使用） #

如果存在列类型为混合类型的特殊场景，需要逐行判断数据类型，可以通过UDF实现：

from pyspark.ml.linalg import Vector
from pyspark.sql import functions as F
from pyspark.sql.types import BooleanType

# 定义判断向量类型的UDF
is_vector_udf = F.udf(lambda x: isinstance(x, Vector), BooleanType())

predictions = predictions.withColumn('test', F.when(is_vector_udf('probability'), 1).otherwise(0))

补充说明 #

如果要通过类型名做字符串对比，可以取类型的类名属性判断：

type(predictions.schema["probability"].dataType).__name__ == 'VectorUDT'

该判断会返回True。

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