如何在Spark DataFrame中访问VectorUDT列的元素？

from pyspark.sql import SparkSession
from pyspark.ml.linalg import Vectors, VectorUDT
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

# 创建SparkSession
spark = SparkSession.builder.appName("VectorUDTExample").getOrCreate()

# 注册VectorUDT
spark.udf.register("vector_udt", lambda v: Vectors.dense(v), VectorUDT())

# 创建一个包含VectorUDT列的DataFrame
data = [(1, Vectors.dense([1.0, 2.0, 3.0])), (2, Vectors.dense([4.0, 5.0, 6.0]))]
columns = ["id", "features"]
df = spark.createDataFrame(data, columns)

# 显示DataFrame
df.show(truncate=False)

# 访问VectorUDT列的第一个元素
first_element_udf = udf(lambda vector: float(vector[0]), DoubleType())
df = df.withColumn("first_element", first_element_udf(df["features"]))

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