如何在 PySpark 中将多个参数传递给 Pandas UDF？

Question

我正在使用以下代码片段：

from cape_privacy.pandas.transformations import Tokenizer

max_token_len = 5


@pandas_udf("string")

def Tokenize(column: pd.Series)-> pd.Series:
  tokenizer = Tokenizer(max_token_len)
  return tokenizer(column)


spark_df = spark_df.withColumn("name", Tokenize("name"))

由于 Pandas UDF 仅使用 Pandas 系列，我无法max_token_len在函数调用中传递参数Tokenize("name")。

因此我必须max_token_len在函数范围之外定义参数。

这个问题中提供的解决方法并没有真正的帮助。此问题是否还有其他可能的解决方法或替代方案？

请指教

Answer 1

在尝试了多种方法后，我找到了一个毫不费力的解决方案，如下图所示：

我创建了一个包装函数 ( Tokenize_wrapper) 来包装Pandas UDF ( Tokenize_udf)，包装函数返回Pandas UDF 的函数调用。

def Tokenize_wrapper(column, max_token_len=10):

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

  return Tokenize_udf(column)



df = df.withColumn("Name", Tokenize_wrapper("Name", max_token_len=5))

使用部分函数（@Vaebhav 的答案）实际上使这个问题的实现变得困难。