pyspark中的自定义函数

由于目前的pyspark不支持dataset，因此只能用dataframe+udf或者rdd的方式来对列进行数据处理。这里介绍一下udf，udf注册有两种方法，一种是调用register方法，一种是调用udf函数。两者都能在withColumn和apply中使用。两种方法的区别是：udf注册后可以使用dataframe的api，而register注册后可以使用spark sql。

1. Spark DF 和 Pandas DF

首先要区分spark dataframe和pandas dataframe。pandas df会将所有数据存储在driver node上，一定要慎用。
spark df 与 pandas df 相互转化性能优化，需要开启配置，默认为关闭。

spark.sql.execution.arrow.enabled true
import numpy as np
import pandas as pd

//初始化 pandas DF
pdf = pd.DataFrame(np.random.rand(100000, 3))
// pdf -> sdf
%time df = spark.createDataFrame(pdf)
// sdf -> pdf
%time result_pdf = df.select("*").toPandas()

2. Spark UDF

UDF是用户自定义函数的缩写。下面展示一下几种UDF使用方法：

from pyspark.sql.functions import udf
df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],("id", "v"))

# 使用register注册的方法
def f1(v):
    return v+1
f1_udf = spark.udf.register('f1',f1)
df.withColumn('v2',f1_udf(df.v)).show() 

# 使用udf回调的方法
@udf('double') #这里也可以简单的使用@udf
def f2(v):
    return v+1
df.withColumn('v2',f2(df.v)).show()

# 使用udf函数注册
def f3(v):
    return v+1
f3_udf = udf(f3,'double')
df.withColumn('v2',f3_udf(df.v)).show()

# 使用register注册的可以使用spark sql
df.createOrReplaceTempView("df")
spark.sql('select v,f1(v) v2 from df').show()

# 使用udf注册的可以使用dataframe api
df.select(['v',f2('v').alias('v2')]).show()

3. Pandas UDF

Pandas UDF构建在 Apache Arrow 之上。Apache Arrow是一个跨平台的在内存中以列式存储的数据层，减少了大量java和python之间序列化和反序列化的工作。

from pyspark.sql.functions import pandas_udf, PandasUDFType
df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],("id", "v"))

#使用回调函数@pandas_udf(schema,return type)来注册
@pandas_udf("id long, v double", PandasUDFType.GROUPED_MAP)
def normalize(pdf):
    v = pdf.v
    return pdf.assign(v=(v - v.mean()) / v.std())
df.groupby("id").apply(normalize).show()

#使用pandas_udf函数直接注册
def plus_one(a):
    return a + 1
plus_one_pd_udf = pandas_udf(plus_one, returnType=LongType())