背景
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'))
相关文章:pyspark 并行调用udf函数