Pyspark UDF (pyspark.sql.functions.pandas_udf) 经验总结本节来学习pysp

最近要把历史数据使用Pyspark做一次回归计算，遇到了OOM的问题，虽然可以通过加机器来解决，但是最好的办法是通过优化代码逻辑来实现，原因是之前的逻辑只是要实现一个简单的工作来做的，直接使用pandas的功能来实现。现在转成spark函数其中就使用了UDF 来实现

什么是UDF

UDF(User Defined Function） 是指用户定义的函数，由Spark使用Arrow来传输数据，并使用Pandas来处理数据，从而实现矢量化操作。使用pandas_udf，可以方便的在PySpark和Pandas之间进行互操作，并且保证性能；其核心思想是将Apache Arrow作为序列化格式。

UDF API 格式

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.

Pandas_UDF returnType类型

returnType类型 | input | output | | --------------- | -------------- | ------------ | | SCALAR | 一个或多个pandas.Series | 一个pandas.Series | | GROUPED_MAP | pandas.DataFrame | pandas.DataFrame | | GROUPED_AGG | 一个或多个pandas.Series | SCALAR |

@pandas_udf 使用装饰器的方法来定义UDF

首先，创建一个基于Pyspark的环境，创建DataFrame

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, PandasUDFType
import pandas as pd
from pyspark.sql.types import StructType,StructField, StringType, MapType,IntegerType
spark = SparkSession.builder.appName("grouping_example").getOrCreate()

@pandas_udf('double', PandasUDFType.GROUPED_AGG)
def my_agg_func(df):
    # 在Pandas DataFrame上应用某个函数
    return df['value'].mean()
simpleData = [("James","Sales","NY",90000,34,10000),
              ("Michael","Sales","NY",86000,56,20000),
              ("Robert","Sales","CA",81000,30,23000),
              ("Maria","Finance","CA",90000,24,23000),
              ("Raman","Finance","CA",99000,40,24000),
              ("Scott","Finance","NY",83000,36,19000),
              ("Jen","Finance","NY",79000,53,15000),
              ("Jeff","Marketing","CA",80000,25,18000),
              ("Kumar","Marketing","NY",91000,50,21000)]

schema = ["employee_name","department","state","salary","age","bonus"]
df = spark.createDataFrame(simpleData, schema)
df.show()

输出结果为

使用GROUPED_AGG例子

AGG Panda UDF类似于Spark聚合函数。常常与groupBy().agg()和pyspark.sql.window一起使用。这种类型的UDF不支持部分聚合，组或窗口的所有数据都将加载到内存中

@pandas_udf('string', functionType=PandasUDFType.GROUPED_AGG)
def my_agg_function(pdf):
#     df =  pd.concat({'age':pdf,'bonus':pdf1},axis=1)
    b = {
            'eincomsCount': 41,
            'eincomsFirstIncomeMonthRange': 56,
            'eincomsMissingMonthMax': 6,
            'eincomsLast18MonthsAmountMax': 29250.0,
            'eincomsLast18MonthsAmountMin': 24750.0,
            'eincomsLast18MonthsAmountAvg': 25031.25,
            'eincomsLast18MonthsAmountMed': -1,
            'eincomsLast18MonthsAmountStd': 1125.0,
            'eincomsLast18MonthsAmountMedian': 24750.0
    }
    return str(b).replace("'","\"")

new_df = df.groupby('department').agg(my_agg_function(df['age']).alias("eincome_factors"))
new_df.printSchema()
new_df.show(truncate=False)

PandasUDFType.SCALAR使用例子

UDF定义了一个转换，PandasUDFType.SCALAR函数可以输入一个或多个pd.Series,输出一个pd.Series，函数的输出和输入有相同的长度 Pandas UDF用于向量化标量操作。可以与select和withColumn等函数一起使用。其中调用的pandas.Series作为输入，返回一个具有相同长度的pandas.Series类型。具体执行流程是，Spark将列分成批(每批进行向量化计算的大小由 spark.sql.execution.arrow.maxRecordsPerBatch 参数控制，默认值10000)，并将每个批作为数据的子集进行函数的调用，执行panda UDF，最后将结果连接在一起。

@pandas_udf(returnType=IntegerType(), functionType=PandasUDFType.SCALAR)
def total_monery(x, y):
    return x + y
    
df.withColumn("total",total_monery("salary","bonus")).show()

输出结果为：

Grouped Map使用例子

GROUPED_MAP UDF定义了转换：pd.DataFrame -> pd.DataFrame，returnType使用StructType描述返回值的pd.DataFrame的模式。
Grouped Map（分组映射）panda_udf与groupBy().apply()一起使用，基本实现思想是“split-apply-combine”模式。
“split-apply-combine”包括三个步骤：

使用DataFrame.groupBy将数据分成多个组。
对每个分组应用一个函数。函数的输入和输出都是pandas.DataFrame。输入数据包含每个组的所有行和列。
将结果合并到一个新的DataFrame中。要使用groupBy().apply()，需要定义以下内容：
定义每个分组的Python计算函数，这里可以使用pandas包或者Python自带方法。
一个StructType对象或字符串，它定义输出DataFrame的格式，包括输出特征以及特征类型。在return 时候需要注意，如果穿进去的PDF 对象有几列，那么returnType的StructType需要做好对应的定义

# 定义的返回类型，除了传进来的"department","bonus"，还有加上新生成的V
@pandas_udf('department string,bonus int ,v double', PandasUDFType.GROUPED_MAP)
def my_map_func(pdf):
    # 在Pandas DataFrame上应用某个函数
    v = pdf['bonus']
    return pdf.assign(v=(v - v.mean()) / v.std())

new_df = df.select("department","bonus").groupby('department').apply(my_map_func)
new_df.printSchema()
new_df.show(truncate=False)

输出结果：

注意！StructType对象中的Dataframe特征顺序需要与分组中的Python计算函数返回特征顺序保持一致。在应用该函数之前，分组中的所有数据都会加载到内存，这可能导致内存不足抛出异常。

⚠️Pandas_UDF与toPandas的区别

修改为Pandas_UDF 的也是因为toPandas性能问题

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

换句话说，@pandas_udf使用panda API来处理分布式数据集，而toPandas()将分布式数据集转换为本地数据

Pyspark UDF (pyspark.sql.functions.pandas_udf) 经验总结