Python pyspark pandas_udf用法及代碼示例

創建一個 pandas 用戶定義函數(又名矢量化用戶定義函數)。

Pandas UDF 是用戶定義的函數，由 Spark 使用 Arrow 執行來傳輸數據，並使用 Pandas 來處理數據，從而允許矢量化操作。 Pandas UDF 是使用 pandas_udf 作為裝飾器或包裝函數來定義的，不需要額外的配置。一般來說，Pandas UDF 的行為與常規 PySpark 函數 API 相同。

參數：

f：函數，可選

用戶自定義函數。如果用作獨立函數，則為 python 函數

returnType： pyspark.sql.types.DataType 或 str，可選

用戶定義函數的返回類型。該值可以是 pyspark.sql.types.DataType 對象或 DDL 格式的類型字符串。

functionType：整數，可選

pyspark.sql.functions.PandasUDFType 中的枚舉值。默認值：標量。存在此參數是為了兼容性。鼓勵使用 Python 類型提示。

注意：

用戶定義的函數不支持條件表達式或布爾表達式中的短路，最終會全部在內部執行。如果函數在特殊行上失敗，解決方法是將條件合並到函數中。

用戶定義的函數在調用方不采用關鍵字參數。

用戶定義函數返回的 pandas.Series 的數據類型應與定義的 returnType 匹配(參見 types.to_arrow_type() 和 types.from_arrow_type() )。當它們之間不匹配時，Spark 可能會對返回的數據進行轉換。不能保證轉換是正確的，用戶應檢查結果的準確性。

目前， pyspark.sql.types.TimestampType 的 pyspark.sql.types.ArrayType 和嵌套的 pyspark.sql.types.StructType 目前不支持作為輸出類型。

例子：

為了使用此 API，通常需要導入以下內容：

>>> import pandas as pd
>>> from pyspark.sql.functions import pandas_udf

從 Spark 3.0 和 Python 3.6+，Python type hints 檢測函數類型如下：

>>> @pandas_udf(IntegerType())
... def slen(s: pd.Series) -> pd.Series:
...     return s.str.len()

在 Spark 3.0 之前，pandas UDF 使用 functionType 來決定執行類型，如下所示：

>>> from pyspark.sql.functions import PandasUDFType
>>> from pyspark.sql.types import IntegerType
>>> @pandas_udf(IntegerType(), PandasUDFType.SCALAR)
... def slen(s):
...     return s.str.len()

最好為 pandas UDF 指定類型提示，而不是通過 functionType 指定 pandas UDF 類型，這將在未來的版本中棄用。

請注意，類型提示在所有情況下都應使用 pandas.Series，但有一種變體，當輸入或輸出列為 pyspark.sql.types.StructType 時，應將 pandas.DataFrame 用於其輸入或輸出類型提示。以下示例顯示了一個 Pandas UDF，它采用 long 列、字符串列和 struct 列，並輸出一個 struct 列。它需要函數指定pandas.Series和pandas.DataFrame的類型提示，如下所示：

>>> @pandas_udf("col1 string, col2 long")
>>> def func(s1: pd.Series, s2: pd.Series, s3: pd.DataFrame) -> pd.DataFrame:
...     s3['col2'] = s1 + s2.str.len()
...     return s3
...
>>> # Create a Spark DataFrame that has three columns including a struct column.
... df = spark.createDataFrame(
...     [[1, "a string", ("a nested string",)]],
...     "long_col long, string_col string, struct_col struct<col1:string>")
>>> df.printSchema()
root
|-- long_column: long (nullable = true)
|-- string_column: string (nullable = true)
|-- struct_column: struct (nullable = true)
|    |-- col1: string (nullable = true)
>>> df.select(func("long_col", "string_col", "struct_col")).printSchema()
|-- func(long_col, string_col, struct_col): struct (nullable = true)
|    |-- col1: string (nullable = true)
|    |-- col2: long (nullable = true)

在以下部分中，它說明了支持的類型提示的組合。為簡單起見，pandas.DataFrame 變體被省略。

• 係列到係列

  pandas.Series , ... -> pandas.Series

  該函數接受一個或多個 pandas.Series 並輸出一個 pandas.Series 。函數的輸出應始終與輸入具有相同的長度。

  >>> @pandas_udf("string")
  ... def to_upper(s: pd.Series) -> pd.Series:
  ...     return s.str.upper()
  ...
  >>> df = spark.createDataFrame([("John Doe",)], ("name",))
  >>> df.select(to_upper("name")).show()
  +--------------+
  |to_upper(name)|
  +--------------+
  |      JOHN DOE|
  +--------------+

  >>> @pandas_udf("first string, last string")
  ... def split_expand(s: pd.Series) -> pd.DataFrame:
  ...     return s.str.split(expand=True)
  ...
  >>> df = spark.createDataFrame([("John Doe",)], ("name",))
  >>> df.select(split_expand("name")).show()
  +------------------+
  |split_expand(name)|
  +------------------+
  |       [John, Doe]|
  +------------------+

  注意

  輸入的長度不是整個輸入列的長度，而是用於每次調用函數的內部批處理的長度。

• 係列迭代器到係列迭代器

  Iterator[pandas.Series] -> Iterator[pandas.Series]

  該函數采用 pandas.Series 迭代器並輸出 pandas.Series 迭代器。在這種情況下，創建的 pandas UDF 實例需要一個輸入列(當該輸入列被稱為 PySpark 列時)。函數的整個輸出的長度應該與整個輸入的長度相同；因此，隻要長度相同，它就可以從輸入迭代器中預取數據。

  當 UDF 執行需要初始化某些狀態時，它也很有用，盡管在內部它的工作方式與係列到係列的情況相同。下麵的偽代碼說明了這個例子。

  @pandas_udf("long")
  def calculate(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
      # Do some expensive initialization with a state
      state = very_expensive_initialization()
      for x in iterator:
          # Use that state for whole iterator.
          yield calculate_with_state(x, state)
  
  df.select(calculate("value")).show()

  >>> from typing import Iterator
  >>> @pandas_udf("long")
  ... def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
  ...     for s in iterator:
  ...         yield s + 1
  ...
  >>> df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"]))
  >>> df.select(plus_one(df.v)).show()
  +-----------+
  |plus_one(v)|
  +-----------+
  |          2|
  |          3|
  |          4|
  +-----------+

  注意

  每個係列的長度是內部使用的批次的長度。

• 多個係列的迭代器到係列的迭代器

  Iterator[Tuple[pandas.Series, …]] -> Iterator[pandas.Series]

  該函數采用多個 pandas.Series 元組的迭代器並輸出 pandas.Series 的迭代器。在這種情況下，創建的 pandas UDF 實例需要與係列(稱為 PySpark 列)一樣多的輸入列。否則，它具有與係列迭代器到係列迭代器情況相同的特征和限製。

  >>> from typing import Iterator, Tuple
  >>> from pyspark.sql.functions import struct, col
  >>> @pandas_udf("long")
  ... def multiply(iterator: Iterator[Tuple[pd.Series, pd.DataFrame]]) -> Iterator[pd.Series]:
  ...     for s1, df in iterator:
  ...         yield s1 * df.v
  ...
  >>> df = spark.createDataFrame(pd.DataFrame([1, 2, 3], columns=["v"]))
  >>> df.withColumn('output', multiply(col("v"), struct(col("v")))).show()
  +---+------+
  |  v|output|
  +---+------+
  |  1|     1|
  |  2|     4|
  |  3|     9|
  +---+------+

  注意

  每個係列的長度是內部使用的批次的長度。

• 係列到標量

  pandas.Series , ... -> Any

  該函數采用pandas.Series 並返回標量值。 returnType 應該是原始數據類型，返回的標量可以是 python 原始類型，例如 int 或 float，也可以是 numpy 數據類型，例如 numpy.int64 或 numpy.float64。理想情況下，Any 應該是相應的特定標量類型。

  >>> @pandas_udf("double")
  ... def mean_udf(v: pd.Series) -> float:
  ...     return v.mean()
  ...
  >>> df = spark.createDataFrame(
  ...     [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
  >>> df.groupby("id").agg(mean_udf(df['v'])).show()
  +---+-----------+
  | id|mean_udf(v)|
  +---+-----------+
  |  1|        1.5|
  |  2|        6.0|
  +---+-----------+

  此 UDF 也可以用作窗口函數，如下所示：

  >>> from pyspark.sql import Window
  >>> @pandas_udf("double")
  ... def mean_udf(v: pd.Series) -> float:
  ...     return v.mean()
  ...
  >>> df = spark.createDataFrame(
  ...     [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
  >>> w = Window.partitionBy('id').orderBy('v').rowsBetween(-1, 0)
  >>> df.withColumn('mean_v', mean_udf("v").over(w)).show()
  +---+----+------+
  | id|   v|mean_v|
  +---+----+------+
  |  1| 1.0|   1.0|
  |  1| 2.0|   1.5|
  |  2| 3.0|   3.0|
  |  2| 5.0|   4.0|
  |  2|10.0|   7.5|
  +---+----+------+

  注意

  出於性能原因，不會複製窗口函數的輸入序列。因此，不允許對輸入序列進行變異，這將導致不正確的結果。出於同樣的原因，用戶也不應該依賴輸入序列的索引。