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

Question

我正在使用 pyspark 2.4.2，因此根据此版本的文档，可以执行此操作来创建 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称为普通的python function，它通过了pandas ZBA834BA01217A3788E459。 但如果你这样做，它将起作用：

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均值。 如何使用注释方法做到这一点？ 从文档中不清楚如何做到这一点。 注释了哪些 function 以及为f参数给出了哪些 function？

Answer 1

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

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