一、什么是Pandas UDF
尽管Apache Spark具有大量内置功能,但Spark的灵活性允许数据工程师和数据科学家定义自己的功能。这些被称为用户自定义函数(UDF)。pandas_udf是PySpark 2.3版本引入的新功能,可以将Pandas的函数应用于Spark的DataFrame。简而言之就是说,spark运用时有些函数并不能满足用户的需求,使用自定义pandas_udf来引入,函数内部使用pandas操作。
二、如何使用
pandasUDF共有三种类型,分别是SCALER,GROUPED_AGG和GROUPED_MAP
1.SCALER
- 可以进行一些简单的操作
- 输入:pandas.Series(一个或多个)
- 输出:pandas.Series(只能有一个)
写法1:一个输入
import math
from pyspark.sql.functions import pandas_udf, PandasUDFType
LANE_WIDTH_M = 3
cos_udf=pandas_udf(lambda x:(LANE_WIDTH_M / 2) *x.map(math.cos),'float',PandasUDFType.SCALAR)
df = df.withColumn('cos_angle_l',cos_udf('angle_l'))
写法2:一个输入
import math
from pyspark.sql.functions import pandas_udf, PandasUDFType
LANE_WIDTH_M = 3
@pandas_udf('float', PandasUDFType.SCALAR)
def cos_udf(col:pd.Series)->pd.Series:
return col.map(lambda x:(LANE_WIDTH_M / 2) *x.map(math.cos))
df = df.withColumn('cos_angle_l',cos_udf('angle_l'))
写法3:多个输入
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('int', PandasUDFType.SCALAR)
def flag_udf(col1,col2,col3):
df=pd.DataFrame({'obj_id': col1,'dis': col2,'class': col3})
'''
中间一系列pandas操作,最后返回df的一列
'''
return df['flag']
df = df.withColumn('flag', flag_udf('obj_id','dis','class'))
2.GROUPED_AGG
- 和SCALER用法差不多,配合groupby.agg使用
- 输入:pandas.Series(一个或多个)
- 输出:pandas.Series(只能有一个)
下面是多个输入的两种写法:
@pandas_udf("double", PandasUDFType.GROUPED_AGG)
def average_column(col1: pd.Series, col2: pd.Series) -> float:
return (col1 + col2).mean()
df = df.groupby('obj_id').agg(average_column(F.col('obj_x'), F.col('obj_y')).alias('mean_col'))
df.show()
from pyspark.sql.types import StructType, StructField,DoubleType
def mean_udf(col1: pd.Series, col2: pd.Series) -> float:
return (col1 + col2).mean()
average_column = pandas_udf(mean_udf, DoubleType(), PandasUDFType.GROUPED_AGG)
df = df.groupby('obj_id').agg(average_column(F.col('obj_x'), F.col('obj_y')).alias('mean_col'))
df.show()
3.GROUPED_MAP
- 当你想返回整个df时使用,返回的结果df可以是和之前完全不同的结构,也可以是新增几列后的df
- 输入:pandas.DataFrame
- 输出:pandas.DataFrame
假如你是单纯的想要增加一列或几列,我建议先创建此列(给他一个默认值,让他识别到此列的类型)
df = df.withColumn('new_index', F.lit(1))
df_schema = df.schema
@pandas_udf(df_schema, PandasUDFType.GROUPED_MAP)
def fix_dataframe_index(df):
df = df.sort_values('obj_timestamp')
df['new_index'] = df['obj_frame_index'].rank(method='dense')
return df
df = df.groupBy('obj_id').apply(fix_dataframe_index)
df.show()
'''
df_schema = StructType([
StructField("header timestamp", StringType()),
StructField("object.id", IntegerType()),StructField("class_label_pred", StringType()),
StructField("pose.position.x", DoubleType()),StructField("pose.position.y", DoubleType()),
StructField("pose.position.z", DoubleType()), StructField("yaw", DoubleType()),
StructField("dimension_x", DoubleType()), StructField("dimension_y", DoubleType()),
StructField("dimension_z", DoubleType()), StructField("abs_velocity", DoubleType()), StructField("direction", StringType()),
StructField("data_name", StringType())
])
'''
我们可以看到,GROUPED_MAP返回的类型是一个df,在spark中pandasudf返回的df必须指定其每个字段的数据类型,如果你的列很多写的就会很麻烦。
所以我这里选用了先增加一列的办法,告诉其new_index类型为int。
否则你需要用上面类似的''' '''中的写法,当然,如果你返回的是一个与之前输入完全不同的,就必须要这样写了
三、GroupedData.applyInPandas
还有一个很GROUPED_MAP类似的groupby方法,区别在于:
-
运行环境:
groupby.apply
: Spark操作,运行在Spark集群上。groupby.inPandas
: Pandas操作,运行在本地Python环境中。
-
使用场景:
groupby.apply
: 在大数据集上可能更快,因为它充分利用了Spark的分布式计算能力。groupby.inPandas
: 在小数据集或中等大小的数据集上可能更快,因为它直接在本地Python环境中运行。
- 使用方法与GROUPED_MAP类似:
df = df.withColumn('new_index', F.lit(1))
def fix_dataframe_index(df):
df = df.sort_values('obj_timestamp')
df['new_index'] = df['obj_frame_index'].rank(method='dense')
return df
df = df.groupBy('obj_id').applyInPandas(fix_dataframe_index, schema=df.schema)
df.show()
schema问题和上面的同样
注意:此方法pandas操作后,返回的df不能含有空值,否则报错,在return前需要填充空值