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|
  +---+----+------+

  注意

  出于性能原因，不会复制窗口函数的输入序列。因此，不允许对输入序列进行变异，这将导致不正确的结果。出于同样的原因，用户也不应该依赖输入序列的索引。