PySpark使用Pandas UDF拆分字符串数组提取国家代码报错如何解决

报错原因 #

你定义Pandas UDF时声明的返回类型为"string"，但UDF实际返回的是每个元素为list的Pandas Series，Arrow序列化过程中需要匹配Spark声明的字符串类型、预期接收字节数据，但是拿到了列表对象，因此触发类型错误。

修复方案 #

只需要将Pandas UDF的返回类型声明修改为和实际输出匹配的字符串数组类型即可，原有逻辑无需调整，完整代码如下：

# 首先导入需要的类型定义
from pyspark.sql.types import ArrayType, StringType
import pyspark.sql.functions as F

# 修改Pandas UDF的返回类型声明
@F.pandas_udf(returnType=ArrayType(StringType()))
def func(s):
    return s.apply(lambda x: [y.split("/")[0] for y in x])

# 调用逻辑保持不变
df.withColumn('source_countries', 
              func(F.map_keys(F.col("country_mapping")))).show()

执行后即可得到你预期的["GBR", "USA"]格式的数组结果。

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