pyspark中udf写法及其使用

背景

udf是处理dataframe比较常见且便捷的功能模块，用户自定的处理功能封装在udf中，使得dateframe的处理变得通俗可控。

udf的两种书写方式

第一种——udf不需要参数

这里的不需要参数，是指udf不需要传入公共的参数（这个里的公共是对udf中多个参数模块而言的）,且该参数非dataframe中的字段。

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, ArrayType
# 定义基础模块
def proc_func(param_val1, param_val2):
	ret = '{}\t{}'.format(param_val1, param_val2)
	return ret
# 定义udf
def proc_func_udf():
	return udf(proc_func, StringType())
# 如果需要返回array，格式需为ArrayType(StringType())
# 使用udf
your_df = your_df.withColumn('new_col_name', proc_func_udf('param_col1', 'param_col2'))
# 以下方法和上面一样（其实就是带参的简化版）
your_df = your_df.withColumn('new_col_name', proc_func_udf()('param_col1', 'param_col2'))

第二种——udf需要参数

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, ArrayType
# 定义基础模块1
def proc_func1(common_val1, param_val1):
	ret = '{}\t{}'.format(common_val1, param_val1)
	return ret
# 定义基础模块2	
def proc_func2(common_val2, param_val2):
	ret = '{}\t{}'.format(common_val2, param_val2)
	return ret
# 定义udf
def proc_func_udf(common_val1, common_val2):
	def proc_func(param_val1, param_val2):
		ret1 = proc_func1(common_val1, param_val1)
		ret2 = proc_func2(common_val2, param_val2)
		ret = ret1 + ret2
		return ret
	return udf(proc_func, StringType())

# 使用udf
common_val1 = []
common_val2 = {}
your_df = your_df.withColumn('new_col_name', proc_func_udf(common_val1, common_val2)('param_col1', 'param_col2'))