Python函数转PySpark UDF如何实现嵌套列表及多结果输出

• 单返回值场景报错有两个诱因：
  1. UDF返回值schema定义错误：你的函数输出是整数列表构成的嵌套列表，对应PySpark类型应为ArrayType(ArrayType(IntegerType()))，你之前写的ArrayType(IntegerType())仅匹配单层整数列表，类型不匹配
  2. 调用UDF时传参错误：PySpark UDF的入参必须是列对象（Column），如果要传入字面量Python列表，需要用F.array()将列表转换为列类型，不能直接传入原生Python列表
• 多返回值场景报错核心原因：同样是嵌套列表类型定义错误，你定义的每个StructField是ArrayType(IntegerType())，仅匹配单层列表，嵌套列表需要嵌套两层ArrayType

1. 单返回值（嵌套整数列表）适配方案 #

首先定义符合业务逻辑的原生Python函数，再按正确类型定义UDF：

# 示例原生函数，输入单层整数列表，输出嵌套整数列表
def my_function(input_list: list[int]) -> list[list[int]]:
    # 此处替换为你的业务逻辑，示例为按2个元素分组
    return [input_list[i:i+2] for i in range(0, len(input_list), 2)]

# 正确定义UDF
from pyspark.sql.types import ArrayType, IntegerType
from pyspark.sql.functions import udf

pyspark_my_function = udf(
    my_function, 
    returnType=ArrayType(ArrayType(IntegerType()))
)

调用示例：

from pyspark.sql import functions as F
# 场景1：对DataFrame的数组类型列调用
df = spark.createDataFrame([([4,5,7,8,10,11],)], schema=["input_arr"])
df.withColumn("output_arr", pyspark_my_function("input_arr")).show(truncate=False)

# 场景2：传入字面量列表测试，需先转成列对象
spark.range(1).select(
    pyspark_my_function(F.array(F.lit(4),F.lit(5),F.lit(7),F.lit(8),F.lit(10),F.lit(11))).alias("test_output")
).show(truncate=False)

2. 多返回值（多个嵌套整数列表）适配方案 #

原生函数返回和schema字段顺序对应的元组即可，schema按嵌套列表类型定义：

# 示例多返回值原生函数
def my_multi_output_function(input_list: list[int]) -> tuple[list[list[int]], list[list[int]]]:
    output1 = [input_list[i:i+2] for i in range(0, len(input_list), 2)]
    output2 = [input_list[i:i+3] for i in range(0, len(input_list), 3)]
    return output1, output2

# 正确定义多返回值schema和UDF
from pyspark.sql.types import StructType, StructField, ArrayType, IntegerType

multi_schema = StructType([
    StructField("output1", ArrayType(ArrayType(IntegerType())), nullable=False),
    StructField("output2", ArrayType(ArrayType(IntegerType())), nullable=False)
])

pyspark_multi_func = udf(my_multi_output_function, returnType=multi_schema)

调用示例：

df = spark.createDataFrame([([4,5,7,8,10,11],)], schema=["input_arr"])
df.withColumn("multi_output", pyspark_multi_func("input_arr"))\
  .select("input_arr", 
          "multi_output.output1", 
          "multi_output.output2")\
  .show(truncate=False)

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