spark3.0教程-SparkSQL：UDAF自定义聚合函数

作者：马育民 • • 阅读：10203

介绍

UDAF，User-Defined Aggregation Functions，用户定义 聚合函数。

应用场景

spark 自带的函数并不能完全满足业务需求，这时就需要我们自定义函数了

特点

many to one，进来多个数据，出去一个结果，是 row 级别操作，如：聚合函数 avg、sum、min

执行过程

实现

spark 3.x 实现

定义类，并继承 Aggregator[-IN, BUF, OUT] 类，然后重写下面方法：

// 初始化buffer，buffer是缓存结果，如：累加结果
def zero: Buff

// 处理一条新的记录，b是缓存结果，a是输入数据
def reduce(b: Buff, a: IN): Buff

// 合并两个buff
def merge(b1: Buff, b2: Buff): Buff

// 最终计算
def finish(reduction: Buff): OUT

// 固定写法，如果缓存是自定义类型，那么赋值 Encoders.product
def bufferEncoder: Encoder[Buff] = Encoders.product

// 最终输出结果的类型
def outputEncoder: Encoder[Double] = Encoders.scalaDouble

案例

定义 UDAF 实现 avg() 求平均数功能

MyAvg 类

该类 继承 Aggregator，实现相应的方法，实现求平均数

package udaf

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator

/**
 * 计算平均值过程中使用的Buffer
 *
 * @param sum
 * @param count
 */
case class Buff(var sum: Long, var count: Long) {
}


/**
 * 计算平均值
 *
 */
class MyAvg extends Aggregator[Long, Buff, Double] {

  // 初始化buffer
  override def zero: Buff = Buff(0L, 0L)

  // 处理一条新的记录
  override def reduce(b: Buff, a: Long): Buff = {
    b.sum += a
    b.count += 1L
    b
  }

  // 合并聚合buffer
  override def merge(b1: Buff, b2: Buff): Buff = {
    b1.sum += b2.sum
    b1.count += b2.count
    b1
  }

  //
  override def finish(reduction: Buff): Double = reduction.sum.toDouble / reduction.count

  override def bufferEncoder: Encoder[Buff] = Encoders.product

  // 最终输出结果的类型
  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble

}

主类

注册刚才的 UDAF 函数，然后在 SQL 中调用

package udaf

import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, SparkSession, functions}

object Std1 {
  def main(args: Array[String]): Unit = {
    val sc = new SparkConf().setMaster("local[*]").setAppName("std")
    val spark = SparkSession.builder().config(sc).getOrCreate()

    val df: DataFrame = spark.read
      .option("header", "false") //默认为false，表示第一行不是表头
      .option("sep", ",") //列分隔符
      .csv("cdz.txt") //读取文件
      .toDF("id","ele_num","time","start_time","end_time")

    df.printSchema()
    //创建临时视图（只能查询不能修改）
    df.createOrReplaceTempView("cdz")
    // 自定义函数
    spark.udf.register("myavg",functions.udaf(new MyAvg()))

//    val sql = "select myavg(ele_num) as avg,avg(ele_num) as avg2 from cdz"
    val sql = "select id,myavg(ele_num) as avg,avg(ele_num) as avg2 from cdz group by id"
    val res: DataFrame = spark.sql(sql)
    res.show()
    spark.close()
  }
}

参考：
https://www.cnblogs.com/cc11001100/p/9471859.html
https://zhuanlan.zhihu.com/p/380127118

原文出处：/show_1IX4MdtdAhhx.html

吉ICP备15003043号