PySpark之SparkSQL 使用Pandas UDF

一、介绍

• Pandas UDF是用户定义的函数，由Spark执行，使用Arrow传输数据，Pandas执行数据，允许向量化操作。Pandas UDF是使用pandas_udf()作为装饰器或包装函数来定义的，不需要额外的配置。Pandas UDF通常表现为常规的PySpark函数API
• pyspark.sql.functions.pandas_udf(f=None, returnType=None, functionType=None)
  • f: 用户定义的函数；
  • returnType: 用户自定义函数的返回值类型，该值可以是pyspark.sql.types.DataType对象或DDL格式的类型字符串
  • functionType: pyspark.sql.functions.PandasUDFType中的枚举值。 默认值：SCALAR.存在此参数是为了兼容性。 鼓励使用Python类型。

二、类型1：seriesToSeries

• 类型提示可以表示为pandas.Series,…->pandas.Series

• 通过pandas_udf()与具有上述类型提示的函数一起使用，它会创建一个Pandas UDF，其中给定的函数采用一个或多个pandas.Series并输出一个pandas.Series。函数的输出应始终与输入的长度相同。在内部，PySpark将通过拆分为批次并作为数据的子集调用每个批次的函数，然后将结果连接在一起来执行Pandas UDF

import os
from pyspark.sql.types import LongType
from pyspark.sql import SparkSession
import pandas as pd
from pyspark.sql.functions import pandas_udf,col

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession\
        .builder\
        .appName("test")\
        .getOrCreate()
    sc = spark.sparkContext
    # 方式1：普通方式创建pandas_func
    def multiply_func(a:pd.Series,b:pd.Series)->pd.Series:
        return a*b
    multiply = pandas_udf(multiply_func,returnType=LongType())
    x = pd.Series([1,2,3])
    print(multiply_func(x,x))
    df = spark.createDataFrame(pd.DataFrame(x,columns=['x']))
    df.select(multiply(col("x"),col("x"))).show()
    print("="*100)
    # 方式2：装饰器方法
    @pandas_udf(LongType())
    def multiply_func1(a:pd.Series,b:pd.Series) -> pd.Series:
        return a*b
    df.select(multiply_func(col("x"),col("x")))\
        .withColumnRenamed("multiply_func1(x,x)","xxx").show()
    spark.stop()

三、类型2: series迭代器到series迭代器[与类型1非常相似]

• 类型提示可以表示为Iterator[pandas.Series]-> Iterator[pandas.Series]。
• 通过pandas_udf()与具有上述类型提示的函数一起使用，它会创建一个 Pandas UDF，其中给定的函数采用迭代器pandas.Series并输出的迭代器pandas.Series。函数的整个输出的长度应该与整个输入的长度相同；因此，只要长度相同，它就可以从输入迭代器中预取数据。在这种情况下，在调用 Pandas UDF 时，创建的 Pandas UDF 需要一个输入列。
• 要使用多个输入列，需要不同的类型提示。请查看后面多个系列的迭代器到系列的迭代器。

import os
from typing import Iterator
from pyspark.sql import SparkSession
import pandas as pd
from pyspark.sql.functions import pandas_udf

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession\
        .builder\
        .appName("test")\
        .getOrCreate()
    sc = spark.sparkContext
    pdf = pd.DataFrame([1,2,3],columns=["x"])
    df = spark.createDataFrame(pdf)
    @pandas_udf("long")
    def plus_one(iterator:Iterator[pd.Series])->Iterator[pd.Series]:
        for x in iterator:
            yield x + 1
    df.select(plus_one("x")).show()
    spark.stop()

四、多个series的迭代器到series的迭代器

• 类型提示可以表示为 Iterator[Tuple[pandas.Series, …]]->Iterator[pandas.Series]
• 通过pandas_udf()与具有上述类型提示的函数一起使用，它创建了一个 Pandas UDF，其中给定的函数采用多个元组pandas.Series的迭代器并输出迭代器pandas.Series。在这种情况下，当调用 Pandas UDF 时，创建的 Pandas UDF需要与元组中的Series一样多的多个输入列。否则，它具有与 Iterator of Series 到 Iterator of Series 情况相同的特性和限制。

import os
from typing import Iterator,Tuple
from pyspark.sql import SparkSession
import pandas as pd
from pyspark.sql.functions import pandas_udf

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession.builder\
        .appName("test")\
        .getOrCreate()
    sc = spark.sparkContext
    pdf = pd.DataFrame([1,2,3],columns=["x"])
    df = spark.createDataFrame(pdf)
    @pandas_udf("long")
    def multiply_two_cols(iterator:Iterator[Tuple[pd.Series,pd.Series]])->Iterator[pd.Series]:
        for a, b in iterator:
            yield a*b
    df.select(multiply_two_cols("x","x")).show()
    spark.stop()

五、类型4: series到标量sclaer

• 类型提示可以表示为pandas.Series, … -> Any。
• 通过pandas_udf()与具有上述类型提示的函数一起使用，它创建了一个类似于 PySpark 的聚合函数的 Pandas UDF。给定的函数采用pandas.Series并返回一个标量值。返回类型应该是原始数据类型，返回的标量可以是python原始类型，例如，int或float或numpy数据类型，例如，numpy.int64或numpy.float64。 Any理想情况下应该是一个特定的标量类型。
• 此 UDF 还可与GroupedData.agg()和Window 一起使用。它定义了从一个或多个pandas.Series到标量值的聚合，其中每个值pandas.Series 代表组或窗口中的一列。
• 注意，这种类型的 UDF 不支持部分聚合，组或窗口的所有数据都将加载到内存中。此外，目前 Grouped 聚合 Pandas UDF 仅支持无界窗口。

import os
from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf
from pyspark.sql import Window
import pandas as pd

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession.builder\
        .appName('test')\
        .getOrCreate()
    sc = spark.sparkContext
    df = spark.createDataFrame([(1,1.0),(1,2.0),(2,3.0),(2,5.0),(2,10.0)],("id","v"))

    @pandas_udf("double")
    def mean_udf(v:pd.Series) -> float:
        return v.mean()
    df.select(mean_udf(df['v'])).show()
    df.groupby("id").agg(mean_udf(df['v'])).show()
    w = Window\
        .partitionBy("id")\
        .rowsBetween(Window.unboundedPreceding,Window.unboundedFollowing)
    df.withColumn('mean_v',mean_udf(df['v']).over(w)).show()
    spark.stop()

六、grouped Map

• 支持Pandas实例的分组映射操作，DataFrame.groupby().applyInPandas() 它需要一个 Python 函数，该函数接受一个pandas.DataFrame并返回另一个pandas.DataFrame。它将每个组映射到pandas.DataFrame的Python函数中每个组。

• 这个 API 实现了“split-apply-combine”模式，它包括三个步骤：

  • 使用将数据分组DataFrame.groupBy()。
  • 对每个组应用一个函数。函数的输入和输出都是pandas.DataFrame。输入数据包含每个组的所有行和列。
  • 将结果合并到一个新的PySpark中DataFrame。

• 要使用DataFrame.groupBy().applyInPandas()，用户需要定义以下内容：

  • 定义每个组的计算的 Python 函数。
  • StructType对象或定义输出PySpark的DataFrame的Scheme。
  • 返回的列标签pandas.DataFrame必须与定义的输出模式中的字段名称（如果指定为字符串）匹配，或者如果不是字符串，则按位置匹配字段数据类型，例如整数索引。

• 请注意，在应用该函数之前，组的所有数据都将加载到内存中。这可能会导致内存不足异常，尤其是在组大小有偏差的情况下。maxRecordsPerBatch的配置不适用于组，由用户来确保分组数据适合可用内存。

import os
from pyspark.sql import SparkSession

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession\
        .builder\
        .appName('test')\
        .getOrCreate()
    sc = spark.sparkContext
    df = spark.createDataFrame(
        [(1,1.0),(1,2.0),(2,3.0),(2,5.0),(2,10.0)],
        ("id","v")
    )
    def subtract_mean(pdf):
        v = pdf.v
        return pdf.assign(v=v-v.mean())
    df.groupby("id").applyInPandas(subtract_mean,schema="id long,v double").show()
    spark.stop()

七、pandas_UDF 与 topandas的区别

• @pandas_udf 创建一个向量化的用户定义函数(UDF)，利用了panda的矢量化特性，是udf的一种更快的替代方案，因此适用于分布式数据集。
• toPandas将分布式spark数据集转换为pandas数据集，对pandas数据集进行本地化，并且所有数据都驻留在驱动程序内存中，因此此方法仅在预期生成的pandas DataFrame较小的情况下使用。
• 换句话说，@pandas_udf使用panda API来处理分布式数据集，而toPandas()将分布式数据集转换为本地数据，然后使用pandas进行处理。