SparkSQL UDF函数

前言

• 无论是Hive还是SparkSQL分析处理数据时，往往需要使用函数，SparkSQL模块本身自带很多实现公共功能的函数，在pyspark.sql.functions中。Hive中常见的自定义函数有三种UDF（一对一）、UDAF（多对一）、UDTF（一对多）。在SparkSQL中，目前仅仅支持UDAF与UDF，而python仅支持UDF。

SparkSQL 定义UDF函数

1. sparksession.udf.register()：注册的UDF可以用DSL和SQL，返回值用于DSL风格，传参内给的名字用于SQL风格。

udf = sparksession.udf.register(param1, param2, param3)
# param1：UDF名称，可用于SQL风格。
# param2：被注册成UDF的方法名。
# param3：声明UDF的返回值类型。

2. pyspark.sql.functions.udf，仅能用于DSL风格。

udf = F.udf(param1, param2)
# param1：被注册成UDF的方法名
# param2：声明UDF的返回值类型
# F是from pyspark.sql import functions as F

SparkSQL UDF相关代码

UDF 返回Float类型

• 代码示例

# coding : utf8
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StringType, IntegerType, FloatType
from pyspark.sql import functions as F
import numpy as np

if __name__ == '__main__':
    ss = SparkSession.builder \
        .appName("test") \
        .master("local[*]") \
        .getOrCreate()
    sc = ss.sparkContext

    # 构建一个RDD
    rdd = sc.parallelize([1.52, 2.78, 3.62, 4.14, 5.26, 6.83, 7.91, 8.06, 9.31]).map(lambda x: [x])
    # print(rdd.collect())
    df = rdd.toDF(["emb"])

    # TODO 1：sparksession方法创建UDF
    # 自定义一个sigmoid函数
    def sigmoid(emb):
        return (1 / (1 + np.exp(- emb))).__float__()
    #    return 2 * emb
    udf2 = ss.udf.register("udf1", sigmoid, FloatType())
    # 参数1: 注册的UDF的名称, 这个udf名称, 仅可以用于 SQL风格
    # 参数2: UDF的处理逻辑, 是一个单独的方法
    # 参数3: 声明UDF的返回值类型, 注意: UDF注册时候, 必须声明返回值类型, 并且UDF的真实返回值一定要和声明的返回值一致
    # 返回值对象: 这是一个UDF对象, 仅可以用于 DSL 语法
    # 当前这种方式定义的UDF, 可以通过参数1的名称用于SQL风格, 通过返回值对象用户DSL风格

    # SQL风格中使用
    # selectExpr 以SELECT的表达式执行, 表达式 SQL风格的表达式(字符串)
    # select方法, 接受普通的字符串字段名, 或者返回值是Column对象的计算
    df.selectExpr("udf1(emb)").withColumnRenamed("udf1(emb)", "emb_rs").show()

    # DSL 风格中使用
    # 返回值UDF对象 如果作为方法使用, 传入的参数 一定是Column对象
    df.select(udf2(df['emb'])).withColumnRenamed("udf1(emb)", "emb_rs").show()

    # TODO 2: functions方法创建udf
    udf3 = F.udf(sigmoid, FloatType())
    df.select(udf3(df['emb'])).withColumnRenamed("sigmoid(emb)", "emb_rs").show()

• 演示结果
  

UDF 返回Array类型

• 代码示例

# coding : utf8
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StringType, IntegerType, FloatType, ArrayType
from pyspark.sql import functions as F

if __name__ == '__main__':
    ss = SparkSession.builder \
        .appName("test") \
        .master("local[*]") \
        .getOrCreate()
    sc = ss.sparkContext

    # 构建一个RDD
    rdd = sc.parallelize([["hadoop spark flink python"], ["hadoop flink java hive"], ["hive spark java es"]])
    df = rdd.toDF(["big_data"])

    # 注册UDF, UDF的执行函数定义
    def split_line(data):
        # 返回值是一个Array对象
        return data.split(" ")

    # TODO 1：sparksession方法创建UDF
    udf2 = ss.udf.register("udf1", split_line, ArrayType(StringType()))

    # DLS风格
    df.select(udf2(df['big_data'])).withColumnRenamed('udf1(big_data)', 'tech').show()
    # SQL风格
    df.createTempView("big_datas")
    # 显示全量 truncate=False
    ss.sql("SELECT udf1(big_data) AS tech FROM big_datas").show(truncate=False)

    # TODO 2：functions方法创建udf
    udf3 = F.udf(split_line, ArrayType(StringType()))
    df.select(udf3(df['big_data'])).show(truncate=False)

• 演示结果
  

UDF 返回Dict类型

• 代码示例

# coding : utf8
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StringType, IntegerType, FloatType, ArrayType
from pyspark.sql import functions as F

if __name__ == '__main__':
    ss = SparkSession.builder \
        .appName("test") \
        .master("local[*]") \
        .getOrCreate()
    sc = ss.sparkContext

    # 构建一个RDD
    rdd = sc.parallelize([1,2,3,4]).map(lambda x: [x])
    df = rdd.toDF(["num"])

    # 需求判断奇偶
    # 输出返回 {"num":1, "type": "odd"}
    def odd_even(data):
        if data %2 == 1:
            return {"num": data, "type": "odd"}
        else:
            return {"num": data, "type": "even"}

    """
    UDF的返回值是字典的话, 需要用StructType来接收
    """
    udf1 = ss.udf.register("udf1", odd_even, StructType().add("num", IntegerType(), nullable=True).add("type", StringType(), nullable=True))

    df.selectExpr("udf1(num)").withColumnRenamed("udf1(num)", "num_dict").show(truncate=False)
    df.select(udf1(df['num'])).withColumnRenamed("udf1(num)", "num_dict").show(truncate=False)

• 演示结果