PySpark替代Pandas apply实现DataFrame整数列转列表类型的方法

错误原因分析 #

你遇到的类型转换报错是因为PySpark的cast方法不支持直接将基础数值类型（如bigint、int）转换为数组类型，仅支持兼容的基础类型、或有对应转换规则的复合类型之间的转换，所以直接cast(ArrayType(IntegerType()))是行不通的。

最优实现方案 #

使用PySpark内置的array函数即可直接构造数组，该方法是原生实现，性能远高于自定义UDF：

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import col, array

# 构造示例DataFrame
df_sp = spark.createDataFrame([(9,2),(9,3)], schema="a int, b int")

# 将b列转换为单元素数组
df_sp = df_sp.withColumn("b", array(col("b")))

# 如果需要严格指定数组元素为int类型，可先对元素做类型转换再构造数组
# df_sp = df_sp.withColumn("b", array(col("b").cast(IntegerType())))

df_sp.show()

运行后输出结果如下，符合预期：

+---+---+
|  a|  b|
+---+---+
|  9|[2]|
|  9|[3]|
+---+---+

可选方案：自定义UDF实现（不推荐） #

如果有特殊的自定义逻辑需要实现，也可以通过UDF完成，但性能比内置函数差：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType

# 定义UDF
to_array_udf = udf(lambda x: [x], ArrayType(IntegerType()))

df_sp = df_sp.withColumn("b", to_array_udf(col("b")))

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