在Pyspark中将任意数量的列合并为Array类型的新列

from pyspark.sql import SparkSession
from pyspark.sql.functions import array

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例数据
data = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
df = spark.createDataFrame(data, ["col1", "col2", "col3"])

# 使用array函数将列合并为Array类型的新列
df_with_array = df.withColumn("new_col", array("col1", "col2", "col3"))

# 显示结果
df_with_array.show()

+----+----+----+---------+
|col1|col2|col3| new_col |
+----+----+----+---------+
|   1|   2|   3|[1, 2, 3]|
|   4|   5|   6|[4, 5, 6]|
|   7|   8|   9|[7, 8, 9]|
+----+----+----+---------+