如何将标量 Pyspark UDF 转换为 Pandas UDF？

Question

我有一个如下的 UDF，它是一个普通的标量 Pyspark UDF：

@udf()
def redact(colVal: column, offset: int = 0):
    if not colVal or not offset:
        return 'X'*8
    else:
        charList=list(colVal)
        charList[:-offset]='X'*(len(colVal)-offset)
        return "".join(charList)

当我阅读时尝试将其转换为 pandas_udf 时，使用矢量化 UDF 代替标量 UDF 有显着的性能改进，但我遇到了很多与 pandas 相关的问题，但我在这方面经验不足。

请帮助我将此 UDF 转换为矢量化 Pandas UDF

Answer 1

redact 函数可以包装在一个函数中，该函数将redact 应用于pd.Series 的每个项目。

需要应用柯里化，因为要传递标量 offset 值。

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

def pandas_wrapper(values: pd.Series, offset: int) -> pd.Series:
    def redact(colVal: str, offset: int = 0):
        if not colVal or not offset:
            return 'X'*8
        else:
            charList=list(colVal)
            charList[:-offset]='X'*(len(colVal)-offset)
            return "".join(charList)
    return values.apply(lambda value: redact(value, offset))

def curried_wrapper(offset: int):
    return F.pandas_udf(lambda x: pandas_wrapper(x, offset), "string")

df = spark.createDataFrame([("abcdef", ), ("12yz", ), (None,)], ("data_col", ))

df.withColumn("redacted", curried_wrapper(2)(F.col("data_col"))).show()

输出

+--------+--------+
|data_col|redacted|
+--------+--------+
|  abcdef|  XXXXef|
|    12yz|    XXyz|
|    null|XXXXXXXX|
+--------+--------+