本节来学习pyspark.sql.functions中的pandas_udf函数。博客中代码基于spark 2.4.4版本。不同版本函数会有不同,详细请参考官方文档。博客案例中用到的数据可以点击此处下载(提取码:2bd5)
pyspark.sql.functions.pandas_udf(f=None, returnType=None, functionType=None)
pandas_udf是用户定义的函数,由Spark使用Arrow来传输数据,并使用Pandas来处理数据,从而实现矢量化操作。使用pandas_udf,可以方便的在PySpark和Pandas之间进行互操作,并且保证性能;其核心思想是将Apache Arrow作为序列化格式。
Pandas UDF通常表现为常规的PySpark函数API。
- f: 用户定义的函数;
- returnType: 用户自定义函数的返回值类型,该值可以是pyspark.sql.types.DataType对象或DDL格式的类型字符串
- functionType: pyspark.sql.functions.PandasUDFType中的枚举值。 默认值:SCALAR.存在此参数是为了兼容性。 鼓励使用Python类型。
Pandas_UDF类型
目前,有三种类型的Pandas_UDF,分别是Scalar(标量映射),GROUPED_MAP(分组映射),GROUPED_AGG(分组聚合),对应的输入输出格式如下表
pandas_udf类型 | 输入 | 输出 |
---|---|---|
SCALAR | 一个或多个pd.Series | 一个pd.Series |
GROUPED_MAP | pd.DataFrame | pd.DataFrame |
GROUPED_AGG | 一个或多个pd.Series | SCALAR |
# 在学习之前先导入必要的包和数据
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType, FloatType
import pandas as pd
from pyspark.sql.functions import pandas_udf, udf
spark = SparkSession.Builder().master('local').appName('GroupedData').getOrCreate()
df = spark.read.csv('../data/data.csv', header=True)
df = df.withColumn('投