SparkSql自定义强类型、弱类型聚合函数

自定义弱类型

package com.chen.sparksql.func

import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object UDAFDemo {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession.builder().appName("UDAFDemo").master("local[2]").getOrCreate()

    val df: DataFrame = spark.read.json("d:/user.json")

    df.createOrReplaceTempView("user")


    // 注册聚合函数
    spark.udf.register("mySum",new MySum)
    spark.udf.register("myAvg",new MyAvg)

    spark.sql("select mySum(salary),myAvg(salary) from user").show

    spark.close()
  }

}

/**
 * 自定义聚合函数 实现sum求和功能
 */
class MySum extends UserDefinedAggregateFunction {
  // 输入数据类型
  override def inputSchema: StructType = StructType(StructField("ele", DoubleType) :: Nil)

  // 缓冲区类型
  override def bufferSchema: StructType = StructType(StructField("sum", DoubleType) :: Nil)

  // 最终聚合结果的数据类型
  override def dataType: DataType = DoubleType

  // 相同的输入是否返回相同的输出
  override def deterministic: Boolean = true

  // 对缓冲区初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    // 在缓冲集合中初始化
    buffer(0) = 0D // 等同     buffer.update(0,0D)
  }

  // 分区内聚合
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    // input 是指使用聚合函数时，缓存过来的参数封装到了Row中
    if (!input.isNullAt(0)) { // 考虑到每行数据传入时 可能有nill的情况 会报错
      //      val v: Double = input.getDouble(0)
      val v: Double = input.getAs[Double](0)
      buffer(0) = buffer.getDouble(0) + v
    }

  }

  // 分区间聚合
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    // buffer1 和buffer2聚合在一起 之后写回buffer1
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
  }

  // 最终返回值
  override def evaluate(buffer: Row): Any = buffer(0)
}


/**
 * 自定义聚合函数 实现avg 求平均值功能
 */
class MyAvg extends UserDefinedAggregateFunction {
  // 输入数据类型
  override def inputSchema: StructType = StructType(StructField("ele", DoubleType) :: Nil)

  // 缓冲区类型
  override def bufferSchema: StructType = StructType(StructField("sum", DoubleType) :: StructField("count", LongType)::Nil)

  // 最终聚合结果的数据类型
  override def dataType: DataType = DoubleType

  // 相同的输入是否返回相同的输出
  override def deterministic: Boolean = true

  // 对缓冲区初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    // 在缓冲集合中初始化
    buffer(0) = 0D // 等同     buffer.update(0,0D)
    buffer(1) = 0L // count
  }

  // 分区内聚合
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    // input 是指使用聚合函数时，缓存过来的参数封装到了Row中
    if (!input.isNullAt(0)) { // 考虑到每行数据传入时 可能有nill的情况 会报错
      //      val v: Double = input.getDouble(0)
      val v: Double = input.getAs[Double](0)
      buffer(0) = buffer.getDouble(0) + v
      buffer(1) = buffer.getLong(1) + 1L
    }

  }

  // 分区间聚合
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    // buffer1 和buffer2聚合在一起 之后写回buffer1
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  // 最终返回值
  override def evaluate(buffer: Row): Any = buffer.getDouble(0) / buffer.getLong(1)
}

自定义强类型

package com.chen.sparksql.func

import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.{DataFrame, Dataset, Encoder, Encoders, Row, SparkSession, TypedColumn}


case class Dog(name: String, age: Int)

case class AgeAvg(sum: Int, count: Int) {
  def avg = sum.toDouble / count
}

object UDAFDemo2{
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession.builder().appName("UDAFDemo").master("local[2]").getOrCreate()

    import spark.implicits._
    val ds: Dataset[Dog] = List(Dog("dahuang", 8), Dog("xiaohuang", 4), Dog("zhonghuang", 6)).toDS()

    // 强类型使用方式 需要转成TypedColumn类型
    val avg: TypedColumn[Dog, Double] = new MyAvg2().toColumn.name("abg")
    ds.select(avg).show
    
    spark.close()
  }

}


/**
 * 自定义强类型 聚合函数 实现avg 求平均值功能
 * 可在ds中使用
 */
class MyAvg2 extends Aggregator[Dog, AgeAvg, Double] {
  // 对缓冲区进行初始化
  override def zero: AgeAvg = AgeAvg(0, 0)

  // 分区内聚合
  override def reduce(b: AgeAvg, a: Dog): AgeAvg = a match {
    case Dog(name, age) => AgeAvg(b.sum + age, b.count + 1)
    // 如果为null 直接返回
    case _ => b
  }

  // 分区间聚合
  override def merge(b1: AgeAvg, b2: AgeAvg): AgeAvg = AgeAvg(b1.sum + b2.sum, b1.count + b2.count)

  // 返回最终的值
  override def finish(reduction: AgeAvg): Double = reduction.avg

  // 对缓冲区进行编码
  override def bufferEncoder: Encoder[AgeAvg] = Encoders.product // 若为样例类，直接返回product编码器

  // 对返回值进行编码
  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}