如何向PySpark中的Pandas UDF传递多个自定义参数？

可行解决方案 #

以下3种方案均可以避免使用全局变量，实现参数灵活传入：

方案1：闭包（工厂函数）（最推荐通用场景使用） #

通过外层函数接收自定义参数，内层定义Pandas UDF并返回，实现参数封装隔离：

from cape_privacy.pandas.transformations import Tokenizer
import pandas as pd
from pyspark.sql.functions import pandas_udf

def create_tokenize_udf(max_token_len: int):
    @pandas_udf("string")
    def tokenize(column: pd.Series) -> pd.Series:
        tokenizer = Tokenizer(max_token_len)
        return tokenizer(column)
    return tokenize

# 调用时按需传入参数即可，支持创建多个不同参数的UDF实例
tokenize_udf = create_tokenize_udf(max_token_len=5)
spark_df = spark_df.withColumn("name", tokenize_udf("name"))

方案2：使用functools.partial绑定参数 #

适合逻辑简单的场景，无需额外编写工厂函数：

from functools import partial
from cape_privacy.pandas.transformations import Tokenizer
import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("string")
def tokenize(column: pd.Series, max_token_len: int) -> pd.Series:
    tokenizer = Tokenizer(max_token_len)
    return tokenizer(column)

# 调用前绑定固定参数
tokenize_udf = partial(tokenize, max_token_len=5)
spark_df = spark_df.withColumn("name", tokenize_udf("name"))

方案3：参数作为字面量列传入（Spark 3.0+兼容） #

如果需要支持动态参数（比如参数来自表字段、不同行使用不同参数），可以把参数转换为Spark字面量列传入UDF：

from cape_privacy.pandas.transformations import Tokenizer
import pandas as pd
from pyspark.sql.functions import pandas_udf, lit

@pandas_udf("string")
def tokenize(column: pd.Series, max_token_len_col: pd.Series) -> pd.Series:
    # 常量列所有值一致，取第一个值即可
    max_token_len = max_token_len_col.iloc[0]
    tokenizer = Tokenizer(max_token_len)
    return tokenizer(column)

# 调用时传入字面量参数，也可替换为表中的参数字段
spark_df = spark_df.withColumn("name", tokenize("name", lit(5)))

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