从python字典创建一个pyspark数据框架，其中一列为np.ndarray

from pyspark.sql import SparkSession
import numpy as np
from pyspark.sql.types import ArrayType, IntegerType

# 初始化SparkSession
spark = SparkSession.builder.appName("example").getOrCreate()

# 创建一个包含NumPy数组的字典
data = {
    'id': [1, 2, 3],
    'values': [np.array([1, 2, 3]), np.array([4, 5, 6]), np.array([7, 8, 9])]
}

# 定义DataFrame的schema
schema = [
    StructField('id', IntegerType(), True),
    StructField('values', ArrayType(IntegerType()), True)
]

# 创建DataFrame
df = spark.createDataFrame(data, schema=schema)

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

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

# 定义一个UDF来处理NumPy数组
def numpy_to_list(arr):
    return arr.tolist()

numpy_to_list_udf = udf(numpy_to_list, ArrayType(IntegerType()))

# 使用UDF转换NumPy数组
df = df.withColumn('values', numpy_to_list_udf(df['values']))