将PySpark中的Python UDF转换为Pandas UDF以提升性能

你写的Pandas UDF错误在于混淆了单条数据和批量Series的处理逻辑：原Python UDF接收单个字符串，而Pandas UDF的参数是Pandas Series（批量数据集合）。你直接对Series调用row.str.split('""')后用"".join()，会把Series里的列表元素强行拼接，触发expected str got list类型错误。

from pyspark.sql.functions import pandas_udf, StringType
import json
import yaml
import pandas as pd

@pandas_udf(returnType=StringType())
def pandas_udf_process(series):
    def handle_single_row(row):
        # 处理空值（Spark NULL对应Pandas NaN）
        if row is None or pd.isna(row):
            return '[]'
        # 移除所有""
        cleaned_str = "".join(row.split('""'))
        # 原逻辑：yaml加载转json
        intermediate_json = json.dumps(yaml.safe_load(cleaned_str))
        # 移除所有"后再加载转最终json
        final_result = json.dumps(yaml.safe_load(intermediate_json.replace('"', '')))
        return final_result
    
    # 对Series中每个元素批量应用处理逻辑
    return series.apply(handle_single_row)

1. 新增内部函数handle_single_row，完全复用原Python UDF的处理逻辑，同时兼容Pandas的空值判断（pd.isna(row)）
2. 用series.apply()批量处理Series中的每一条数据，确保每个元素都按原逻辑单独处理，避免类型不匹配
3. 函数名改为更清晰的pandas_udf_process，避免和原函数重名

输入示例：

"[{\"\"abc\"\": \"\"abc\"\", \"\"def\"\": \"\"18\"\", \"\"ghi\"\": 3, \"\"jkl\"\": 0, \"\"mno\"\": []}]"

处理后输出：

"[{abc: abc, def: 18, ghi: 3, jkl: 0, mno: []}]"

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