基于pandas实现UDF函数

「这是我参与11月更文挑战的第17天，活动详情查看：2021最后一次更文挑战」

pandas UDF

pandas UDF 是用户定义的函数, 由spark来执行, 使用arrow传输数据, pandas函数处理数据(写的py函数),允许向量化(充分的利用计算机的并行性)操作 , pandas UDF 是使用 pandas_udf() 作为装饰器(相当于给pandas 函数扩充功能, 支持将其转换为spark 的函数)来定义的, 不需要额外配置, 而且 pandas_udf() 通过作为 pyspark的函数API使用(F.pandas_udf())

• 基于pandas实现UDF函数: series To series

  • 类型可以表现为: pandas.series TO pandas .series

需求: 计算两列的乘积

from pyspark import SparkContext, SparkConf
from pyspark.sql import SparkSession
from pyspark.sql.types import *
import  pandas as pd
import pyspark.sql.functions as F
import os

# 目的: 锁定远端操作环境, 避免存在多个版本环境的问题
os.environ["SPARK_HOME"] = "/export/server/spark"
os.environ["PYSPARK_PYTHON"] = "/root/anaconda3/bin/python"
os.environ["PYSPARK_DRIVER_PYTHON"] = "/root/anaconda3/bin/python"

if __name__ == '__main__':
    print("基于pandas 实现 pyspark的 UDF函数定义与使用")

    # 1) 创建sparkSession对象
    spark = SparkSession.builder.master("local[*]").appName("_02_pd_udf").getOrCreate()

    # 开启 arrow方案:
    spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", True)

    # 2) 初始一份数据集
    df = spark.createDataFrame([(1, 3), (2, 4), (1, 2), (3, 3)], ['a', 'b'])

    df.createTempView("t1")

    # 3) 定义一个接收两列的数据, 返回乘积结果的pandas 函数
    def cj_func(a:pd.Series,b:pd.Series) -> pd.Series:
        return  a * b

    # 4) 将这个pandas的函数 转换为spark SQL的函数:  pandas_udf()
    
    # 方式一: 通过原生方式来注册 pandas 的函数
    cj_dsl_pdf = spark.udf.register('cj_s_pdf', cj_func, returnType=LongType())

    # DSL:
    df.select(df['a'], df['b'], cj_dsl_pdf('a', 'b')).show()
    
    # sql:
    spark.sql("select a, b, cj_s_pdf(a,b)  from t1").show()
    
    # 方式二: 通过 API形式定义  进行注册
    pd_udf = F.pandas_udf(cj_func, returnType=LongType())

    # DSL中使用
    df.select(df['a'],df['b'],pd_udf('a','b')).show()

    # SQL中使用   方式一 无法在SQL中使用  报错
    # spark.sql("select a, b, pd_udf(a,b)  from t1").show()

    # 方式三: 基于注解(装饰器方案)
    @F.pandas_udf(returnType=LongType())
    def cj_func_2(a:pd.Series,b:pd.Series) -> pd.Series:
        return  a * b

    # DSL中使用
    df.select(df['a'], df['b'], cj_func_2('a', 'b')).show()

    # SQL中使用 无法使用的 
    #spark.sql("select a, b, cj_func_2(a,b)  from t1").show()