PySpark UDF返回多列表：如何定义正确的返回数据类型？

1. 先定义单个字典对应的结构类型 #

两个列表里的字典格式一致，先把单个字典的结构用StructType和StructField定义好，字段类型根据实际数据调整（比如字符串用StringType，整数用IntegerType）：

from pyspark.sql.types import StructType, StructField, StringType

# 单个字典对应的Schema，根据实际字段修改类型和名称
item_schema = StructType([
    StructField("aa", StringType(), nullable=True),
    StructField("cc", StringType(), nullable=True)
])

2. 定义UDF的返回类型 #

你的函数返回两个列表，对应PySpark里包含两个数组字段的结构体类型，每个数组的元素类型是上面定义的item_schema：

from pyspark.sql.types import ArrayType

# UDF返回的整体Schema：两个数组字段分别对应list1和list2
udf_return_schema = StructType([
    StructField("list1", ArrayType(item_schema), nullable=True),
    StructField("list2", ArrayType(item_schema), nullable=True)
])

3. 注册并使用UDF #

把你的Python函数用@udf装饰器注册，指定返回类型为上面定义的udf_return_schema：

from pyspark.sql.functions import udf

@udf(returnType=udf_return_schema)
def function(parameter):
    list1 = [{"aa":"bb", "cc":"dd"}, {"aa":"ee", "cc":"ff"}]
    list2 = [{"aa":"gg", "cc":"hh"}, {"aa":"ii", "cc":"jj"}]
    # 这里写你的业务逻辑
    return list1, list2

4. 提取并使用返回的两个列表 #

调用UDF后，结果会是一个结构体列，你可以通过点语法提取其中的list1和list2，也可以展开数组做后续处理：

# 假设原数据框是df，调用UDF生成结果列
result_df = df.withColumn("udf_output", function(df.your_parameter_column))

# 提取两个列表到单独列
result_df = result_df.withColumn("list1", result_df.udf_output.list1) \
                    .withColumn("list2", result_df.udf_output.list2)

# 示例：展开list1中的元素，查看每个字典的字段
from pyspark.sql.functions import explode
result_df.select(explode("list1").alias("dict_item")) \
         .select("dict_item.aa", "dict_item.cc") \
         .show()

内容的提问来源于stack exchange，提问作者dpv