pyspark Pandas UDF多种写法python

一、什么是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前需要填充空值