如何将向量类型的VectorAssembler输出转换为数组

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.sql.functions import expr

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

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

# 使用VectorAssembler将多个特征列合并为向量类型的特征列
assembler = VectorAssembler(inputCols=["col1", "col2", "col3"], outputCol="features")
output = assembler.transform(df)

# 使用vector_to_array函数将向量类型的列转换为数组类型的列
output = output.withColumn("features_array", expr("vector_to_array(features)"))

# 显示结果
output.show(truncate=False)