基于Pyspark的Pandas_udf使用方法

运行环境：

• Spark2.4.2
• python 3.6

有一些关于spark df和pandas df的介绍，可以详细看基于Pyspark的Pandas_udf使用方法
这里介绍我在使用过程中遇到的问题

Pandas UDFs

pandas udf是用户定义的函数，是由spark用arrow传输数据，pandas去处理数据。我们可以使用pandas_udf作为decorator或者registor来定义一个pandas udf函数，不需要额外的配置。目前，pandas udf有三种类型：标量映射（Scalar）和分组映射（Grouped Map）和分组聚合（Grouped Aggregate）。

Grouped Map（分组映射）

Grouped Map pandas_udf与groupBy().apply()一起使用，后者实现了split-apply-combine模式。拆分应用组合包括三个步骤：

1. df.groupBy()对数据分组
2. apply()对每个组进行操作，输入和输出都是dataframe格式
3. 汇总所有结果到一个dataframe中

使用groupBy().apply()，用户需要定义以下内容：

1. 一个函数，放在apply()里
2. 一个输出的schema（参考的作者说输入输出的schema必须相同，这一点不能苟同）

请注意，在应用函数之前，组的所有数据都将加载到内存中。这可能导致内存不足异常，尤其是当组的大小skwed的时候。maxRecordsPerBatch不适用于这里。所以，用户需要来确保分组的数据适合可用内存。也就是说groupby操作有可能导致数据倾斜的问题，如果groupby后某个key的数据超出了executor memory的大小，则会报出内存错误。

from pyspark.sql.functions import pandas_udf, PandasUDFType, lit

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))
df = df.withColumn('v2', lit(0))
schema = df.select('id', 'v2').schema
# or df.schema
@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def subtract_mean(pdf):
    # pdf is a pandas.DataFrame
    pdf['v2'] = pdf.v
    return pdf[['id', 'v2']]

df.groupby("id").apply(subtract_mean).show()

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)中的schema可以也定义为"id long, v2 double"。也就是说只要定义好输出的schema，就输出任意的dataframe。

如何自定义schema

上面的代码可以修改为：

from pyspark.sql.types import *
from pyspark.sql.functions import pandas_udf, PandasUDFType, lit

defined_schema = StructType([
    StructField("id", IntegerType(), True),
    StructField("v2", IntegerType(), True), 
])

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))
# or df.schema
@pandas_udf(defined_schema, PandasUDFType.GROUPED_MAP)
def subtract_mean(pdf):
    # pdf is a pandas.DataFrame
    pdf['v2'] = pdf.v
    return pdf[['id', 'v2']]

df.groupby("id").apply(subtract_mean).show()

schema里可以使用任意spark支持的类型，参考pyspark入门系列 - 08 pyspark.sql.types数据类型汇总

注：对于pyspark不支持的type怎么办？参考More Efficient UD(A)Fs with PySpark

标量映射（Scalar）

其用于向量化标量操作。它们可以与select和withColumn等函数一起使用。python函数应该以pandas.series作为输入，并返回一个长度相同的pandas.series。在内部，spark将通过将列拆分为batch，并将每个batch的函数作为数据的子集调用，然后将结果连接在一起，来执行padas UDF。也就是对dataframe的某一列进行操作，输入和输出都是Series!
第一种写法：

import pandas as pd
from pyspark.sql.functions import col, pandas_udf
from pyspark.sql.types import LongType

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))
# Declare the function and create the UDF
@pandas_udf(returnType=LongType())
def multiply_func(a, b):
    return a * b
# or multiply = pandas_udf(multiply_func, returnType=LongType())
# Execute function as a Spark vectorized UDF
df.select(multiply_func(col("v"), col("v"))).show()
df.withColumn("v^2", multiply_func(col("v"), col("v"))).show()

第二种写法：

import pandas as pd

from pyspark.sql.functions import col, pandas_udf
from pyspark.sql.types import LongType
import pyspark.sql.functions as F

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))
# Declare the function and create the UDF
# @pandas_udf(returnType=LongType())
def multiply_func(a, b):
    return a * b

multiply = pandas_udf(multiply_func, returnType=LongType())
# Execute function as a Spark vectorized UDF
df.select(multiply(col("v"), col("v"))).show()
df.withColumn("v^2", multiply(col("v"), col("v"))).show()

第三种写法：

import pyspark.sql.functions as F

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v"))

def multiply_func(a, b):
    return a * b

df.withColumn('v^2', F.UserDefinedFunction(multiply_func)(df.v, df.v)).show()
df.withColumn('v^2', F.UserDefinedFunction(multiply_func,returnType=LongType())(df.v, df.v)).schema

但是这里v^2的类型是StringType，不是LongType。第一种写法

参考：

基于Pyspark的Pandas_udf使用方法

相似问题：
1、pyspark如何在groupby后使用自定义的聚合函数
2、pyspark在groupby后如何操作