pyspark dataframe之udf

PySpark UDF概念引出

在pandas中自定义函数，通过遍历行的方式，便捷实现工程师的需求。但是对于数据量较大的数据处理，会出现速度过慢甚至超内存的问题。Spark作为替代pandas处理海量数据的工具，参照 pandas udf 定义了名为PandasUDFType的类，通过自定义函数的方式spark处理数据的灵活度和高效率有很大亮点。
从spark 1.3到2.3udf函数有row-at-a-time UDF与Pandas UDFs(spark 2.3)两种定义方式，下面就这两种方式详细介绍。

row-at-a-time UDF

传统的UDF定义方式，从1.3版本开始，意如其名与pandas一样，一次执行一行遍历实现自定义函数的功能。spark dataframe是基于行存储的一种结构，最常见的使用场景是用withColumn新建列。
先定义好spark dataframe测试数据，下例是分别使用注册UDF和修饰函数的方式实现自定函数。

from pyspark.sql import HiveContext
from pyspark import SparkContext
import numpy as np
import pandas as pd

sc = SparkContext(appName='spark_udf')
sqlContext = HiveContext(sc)
pdf = pd.DataFrame(np.random.rand(1000, 3))
pdf = pdf.rename(columns={0: "one", 1: "two", 2: "three"})
pdf["id"] = np.random.randint(0, 50, size=len(pdf))
sdf = sqlContext.createDataFrame(pdf)

注册自定义函数

不同于pandas即定义即用的方式，使用前需要调用pyspark.sql.functions包中的udf函数声明，定义函数的返回数据类型，注册自定义函数，如下例所示。

from pyspark.sql.types import DoubleType

def plus_one(a):
    return a + 1
    
plus_one_udf = udf(plus_one, returnType=DoubleType())
sdf = sdf.withColumn("one_processed", plus_one_udf(sdf["one"]))

修饰自定义函数

指定数据返回类型后使用udf注释定义的函数，调用自定义的函数也可以实现spark调用自定义函数。

@udf(returnType=DoubleType())
def plus_one(a):
    return a + 1

sdf = sdf.withColumn("one_processed", plus_one(sdf["one"]))

Pandas UDFs

传统的自定义函数一次只能操作一行，因此会遭遇高序列化和调用开销。在spark 2.3以上版本，基于 Apache Arrow 构建Pandas UDF 完全用 Python 定义低开销，高性能 UDF的能力，此外自定义函数不仅针对数据列，而且可以针对Group和Window。
在 Spark 2.3 中，有两种类型的 Pandas UDF: 标量(scalar)和分组映射(grouped map)。

Scalar Pandas UDFs

使用pandas_udf修饰定义的函数，函数接收double 类型的参数，定义函数接收 pandas.Series 类型的参数 “a”，并将 “a + 1” 的结果作为pandas.Series 返回。因为 “a + 1” 是在 pandas.Series 上进行矢量化的，所以 Pandas 版本比 row-at-a-time 的版本快得多。
自定义Pandas UDF如下：

from pyspark.sql.functions import pandas_udf, PandasUDFType

#使用 pandas_udf 定义一个 Pandas UDF
@pandas_udf('double', PandasUDFType.SCALAR)
#输入/输出都是 double 类型的 pandas.Series
def pandas_plus_one(a):
    return a + 1

sdf.withColumn("one_processed", pandas_plus_one(sdf["one"]))

scala pandas UDF定义函数的特点 ：

• 输入、输出数据类型 -> pandas.Series
• 输出数据大小 -> 和输入大小一样
• 函数装饰器中的返回类型 -> 一个 DataType，用于指定返回的 pandas.Series 的类型

Grouped Map Pandas UDFs

Grouped Map Pandas UDF 是针对某些组的所有数据进行操作。Grouped Map Pandas UDF 首先根据 groupby 运算符中指定的条件将 Spark DataFrame 分组，然后将用户定义的函数（pandas.DataFrame -> pandas.DataFrame）应用于每个组，并将结果组合并作为新的 Spark DataFrame 返回。

与Scalar自定义函数对应，Grouped map自定义函数特点如下：

• 输入、输出数据类型 -> pandas.DataFrame
• 输出数据大小 -> 任何尺寸
• 函数装饰器中的返回类型 -> 一个 StructType，用于指定返回的 pandas.DataFrame 中每列的名称和类型

下面用两个例子阐述grouped map Pandas UDF 的使用场景

Subtract Mean

下例显示了使用 grouped map Pandas UDFs从组中的three列每个值减去平均值。每个输入到自定义函数的 pandas.DataFrame 具有相同的 “id” 值。这个用户定义函数的输入和输出模式是相同的，所以我们将“df.schema” 传递给装饰器 pandas_udf 来指定模式。

@pandas_udf(sdf.schema, PandasUDFType.GROUPED_MAP)
#Input/output are both a pandas.DataFrame
def subtract_mean(df):
    return df.assign(new_col=df.three - df.three.mean())

sdf.groupby("id").apply(subtract_mean)

性能比较

参考性能图

从 plus_one 和 substract_mean两个函数的性能比较可知，Pandas UDF 的表现比 row-at-a-time UDF 好得多。

参考链接

1. https://www.jianshu.com/p/87d70918e16e