运行环境:
- Spark2.4.2
- python 3.6
有一些关于spark df和pandas df的介绍,可以详细看基于Pyspark的Pandas_udf使用方法
这里介绍我在使用过程中遇到的问题
Pandas UDFs
pandas udf是用户定义的函数,是由spark用arrow传输数据,pandas去处理数据。我们可以使用pandas_ud
f作为decorator或者registor来定义一个pandas udf函数,不需要额外的配置。目前,pandas udf有三种类型:标量映射(Scalar)和分组映射(Grouped Map)和分组聚合(Grouped Aggregate)。
Grouped Map(分组映射)
Grouped Map pandas_udf
与groupBy().apply()
一起使用,后者实现了split-apply-combine模式。拆分应用组合包括三个步骤:
- df.groupBy()对数据分组
- apply()对每个组进行操作,输入和输出都是dataframe格式
- 汇总所有结果到一个dataframe中
使用groupBy().apply()
,用户需要定义以下内容:
- 一个函数,放在apply()里
- 一个输出的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
。第一种写法
参考:
相似问题:
1、pyspark如何在groupby后使用自定义的聚合函数
2、pyspark在groupby后如何操作