使用Spark pandas_udf返回字符串数组报错，如何解决？

这个报错的核心原因是：构造返回结果时，错误地将整个Series对象作为数组元素返回，而非每个元素对应一个字符串数组，导致Arrow类型校验失败。

正确实现示例 #

以下是可正常运行的代码，演示如何用pandas_udf返回字符串数组：

1. 导入依赖并初始化Spark会话

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import ArrayType, StringType
import pandas as pd

spark = SparkSession.builder.appName("StringArrayUDFDemo").getOrCreate()

2. 定义正确的pandas_udf

@pandas_udf(ArrayType(StringType()))
def generate_str_array(s: pd.Series) -> pd.Series:
    # 核心：对Series中的每个元素，返回一个字符串列表（数组）
    # 这里以字符串拆分为例，可替换为你的业务逻辑
    return s.apply(lambda x: x.split(","))

3. 创建测试数据并调用UDF

# 构造含字符串列的DataFrame
df = spark.createDataFrame([("apple,banana,orange",), ("cat,dog,bird",)], ["input_str"])
# 调用UDF生成字符串数组列
result_df = df.withColumn("str_array", generate_str_array("input_str"))
result_df.show(truncate=False)

关键注意事项 #

• 必须确保UDF返回的是pd.Series，其中每个元素是字符串列表（即对应Spark的array<string>类型）
• 不要直接返回整个Series作为单个数组元素，需通过apply逐元素处理生成数组
• UDF的返回类型必须严格指定为ArrayType(StringType())，与实际返回的元素类型完全匹配

固定字符串数组场景示例 #

如果需要为每个输入元素返回固定的字符串数组，可按如下方式实现：

@pandas_udf(ArrayType(StringType()))
def fixed_str_array_udf(s: pd.Series) -> pd.Series:
    # 为每个输入元素返回相同的字符串数组
    return pd.Series([["foo", "bar", "baz"]] * len(s))

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