在 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__會降低代碼的可讀性。 但如果它只是用於單元測試目的，那應該沒問題。