如何在PySpark中创建一个返回字符串数组的自定义函数？

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

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

# 示例数据
data = [("John Doe", 25), ("Jane Smith", 30), ("Bob Johnson", 35)]
df = spark.createDataFrame(data, ["name", "age"])

# 自定义函数
def split_name(name):
    return name.split(" ")

# 注册自定义函数
split_name_udf = udf(split_name, ArrayType(StringType()))

# 使用自定义函数
df = df.withColumn("name_array", split_name_udf(df["name"]))

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

+-----------+---+----------------+
|name       |age|name_array      |
+-----------+---+----------------+
|John Doe   |25 |[John, Doe]     |
|Jane Smith |30 |[Jane, Smith]   |
|Bob Johnson|35 |[Bob, Johnson]  |
+-----------+---+----------------+