如何将Pyspark标量UDF转换为向量化Pandas UDF

以下是和原标量UDF逻辑完全对齐的向量化实现：

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

@pandas_udf("string")
def redact(colVal: pd.Series, offset: int = 0) -> pd.Series:
    # offset为0时直接批量返回固定结果
    if offset == 0:
        return pd.Series(['X' * 8] * len(colVal))
    
    def _process_single(s):
        # 处理单条字符串的逻辑，和原UDF完全一致
        if pd.isna(s) or not s:
            return 'X' * 8
        return 'X' * (len(s) - offset) + s[-offset:]
    
    # 批量处理整个Series
    return colVal.apply(_process_single)

使用注意事项 #

• 该UDF的调用方式和你原来的标量UDF完全相同，不需要调整上层业务代码，传入列和offset参数即可
• 批量处理的特性相比标量UDF能减少大量JVM和Python进程的通信开销，性能提升幅度通常在数倍到数十倍区间
• 如果你的PySpark版本在3.0以上，该实现完全兼容最新的pandas UDF规范，不需要额外配置

内容的提问来源于stack exchange，提问作者ASHISH M.G