如何将带位置参数的UDF转换为单个PySpark Pandas矢量化UDF？

需求：把接收邮箱字符串和位置整数参数的普通UDF改造为单个PySpark Pandas矢量化UDF，避免分别创建处理位置0和1的两个UDF，核心疑问是位置参数position在矢量化UDF中的类型定义方式。

原普通UDF代码 #

@udf("string")
def splitEmailUDF(email: str, position: int) -> str:
    return email.split("@")[position]

当前拆分实现（待优化） #

目前只能通过创建两个独立的矢量化UDF实现需求，但希望合并为一个：

from pyspark.sql.functions import pandas_udf

@pandas_udf("string")
def splitFirstNameUDFVec(email: pd.Series) -> pd.Series:
    return email.str.split("@").str[0]

@pandas_udf("string")
def splitDomainUDFVec(email: pd.Series) -> pd.Series:
    return email.str.split("@").str[1]

解决方案 #

在PySpark Pandas矢量化UDF中，标量参数（如这里的position）直接标注为普通Python类型（int），但调用时需确保该参数是常量或通过lit()传递的标量值（不能是DataFrame的列，列参数需标注为pd.Series）。

最终合并后的矢量化UDF代码 #

from pyspark.sql.functions import pandas_udf, lit, col
import pandas as pd

@pandas_udf("string")
def splitEmailUDFVec(email: pd.Series, position: int) -> pd.Series:
    # 对Series中每个邮箱字符串拆分后，取指定位置的元素
    return email.str.split("@").str[position]

调用示例 #

# 提取用户名（位置0）
df = df.withColumn("username", splitEmailUDFVec(col("email"), lit(0)))

# 提取域名（位置1）
df = df.withColumn("domain", splitEmailUDFVec(col("email"), lit(1)))

原理说明 #

当向矢量化UDF传递标量参数（如lit(0)）时，Spark会将该标量值直接传入函数的position参数；而email参数对应DataFrame的列，会被转为pd.Series。利用Pandas的str方法可以高效地对整个Series批量处理，实现和普通UDF一致的逻辑但性能更优。

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