Spark/Databricks中Pandas UDF生成新列报错问题排查

核心原因 #

@pandas_udf是Spark为分布式Spark DataFrame设计的矢量化用户定义函数，仅能适配Spark的Column对象或在Spark API调度下接收Pandas Series，无法直接作用于本地Pandas DataFrame的列。你直接将本地Pandas Series传入该装饰器修饰的函数，Spark的UDF处理逻辑无法识别本地Pandas对象，因此抛出类型错误。

两种解决路径 #

路径1：针对本地Pandas DataFrame，移除@pandas_udf装饰器 #

直接编写普通Pandas矢量化函数处理列即可：

import pandas as pd

# 去掉@pandas_udf装饰器，改为普通Pandas函数
def clean_text(s: pd.Series) -> pd.Series:
    # 示例数据清洗逻辑，可替换为你的实际操作
    return s.str.strip().str.lower().replace(r'\s+', ' ', regex=True)

# 直接给本地Pandas DataFrame赋值新列
pandasDF['newColumn'] = clean_text(pandasDF['oldColumn'])

路径2：针对Spark DataFrame，用Spark API调用@pandas_udf #

如果源数据实际是Spark分布式DataFrame，需通过Spark的withColumn方法调用该UDF：

from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf('string')
def clean_text(s: pd.Series) -> pd.Series:
    # 你的数据清洗逻辑
    return s.str.strip().str.lower().replace(r'\s+', ' ', regex=True)

# 对Spark DataFrame执行列操作
sparkDF = sparkDF.withColumn('newColumn', clean_text(sparkDF['oldColumn']))

# 若需转成本地Pandas DataFrame查看
pandasDF = sparkDF.toPandas()

额外注意 #

在Databricks环境中务必明确区分：spark.sql.DataFrame是分布式数据集，pandas.DataFrame是本地单机数据集，两者的API规则、适配的UDF类型完全不同。@pandas_udf的Pandas Series参数只能由Spark分布式计算时自动传入，不能手动传递本地Pandas对象。

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