在 PySpark Pandas UDF 中指定用户定义函数的正确方法

Question

我使用的是 pyspark 2.4.2，因此根据此版本的 docs 可以创建一个 GROUPED_MAP:

from pyspark.sql.functions import pandas_udf, PandasUDFType
df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],("id", "v"))

@pandas_udf(returnType="id long, v double", functionType=PandasUDFType.GROUPED_MAP)
def subtract_mean(pdf):
    v = pdf.v
    return pdf.assign(v=v - v.mean())

df.groupby("id").apply(subtract_mean).show()

这有效，但您不能将 subtract_mean 作为传递给 pandas DataFrame 的普通 python 函数来调用。但如果你这样做，它就会起作用：

def subtract_mean(pdf):
    v = pdf.v
    return pdf.assign(v=v - v.mean())

sub_spark = pandas_udf(f=subtract_mean, returnType="id long, v double", functionType=PandasUDFType.GROUPED_MAP)

df.groupby("id").apply(sub_spark).show()

现在您可以通过 pandas DataFrame 从 python 调用 subtract_mean。如何使用注释方法做到这一点？从文档中不清楚如何执行此操作。 f参数注解什么功能，给出什么功能？

Answer 1

这两种方式对于指定UDF是等价的。装饰器方法只是一种更简洁的做事方式。装饰器后面的函数作为 f 参数传递。

如in this answer所述，您可以使用subtract_mean.__wrapped__取回原始的未修饰函数。不过，您问题中的第二种方法更具可读性。使用 __wrapped__ 会降低代码的可读性。但如果只是为了单元测试目的，应该没问题。

在 PySpark Pandas UDF 中指定用户定义函数的正确方法

Correct Way to Specify User-Defined Function in PySpark Pandas UDF

python

user-defined-functions

pandas

apache-spark

pyspark