SparkSQL之UDAF使用

Posted 2021-02-15 Runner_Jack

1.创建一个类继承UserDefinedAggregateFunction类。

---------------------------------------------------------------------

package cn.piesat.test

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DataTypes, IntegerType, StructType}

class CountUDAF extends UserDefinedAggregateFunction{
  /**
    * 聚合函数的输入类型
    * @return
    */
  override def inputSchema: StructType = {
    new StructType().add("ageType",IntegerType)
  }

  /**
    * 缓存的数据类型
    * @return
    */
  override def bufferSchema: StructType = {
    new StructType().add("bufferAgeType",IntegerType)
  }

  /**
    * UDAF返回值的类型
    * @return
    */
  override def dataType: DataType = {
    DataTypes.StringType
  }

  /**
    * 如果该函数是确定性的，那么将会返回true，一般给true就行。
    * @return
    */
  override def deterministic: Boolean = true

  /**
    * 为每个分组的数据执行初始化操作
    * @param buffer
    */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0)=0
  }

  /**
    * 更新操作，指的是每个分组有新的值进来的时候，如何进行分组对应的聚合值的计算
    * @param buffer
    * @param input
    */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val num= input.getAs[Int](0)
    buffer(0)=buffer.getAs[Int](0)+num
  }

  /**
    * 分区合并时执行的操作
    * @param buffer1
    * @param buffer2
    */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0)=buffer1.getAs[Int](0)+buffer2.getAs[Int](0)
  }

  /**
    * 最后返回的结果
    * @param buffer
    * @return
    */
  override def evaluate(buffer: Row): Any = {
    buffer.getAs[Int](0).toString
  }
}
--------------------------------------------------------------


2.在main函数中使用样例
---------------------------------------------------------------

package cn.piesat.test

import org.apache.spark.sql.SparkSession

import scala.collection.mutable.ArrayBuffer


object SparkSQLTest {

  def main(args: Array[String]): Unit = {
    val spark=SparkSession.builder().appName("sparkSql").master("local[4]")
      .config("spark.serializer","org.apache.spark.serializer.KryoSerializer").getOrCreate()
    val sc=spark.sparkContext
    val sqlContext=spark.sqlContext
    val workerRDD=sc.textFile("F://Workers.txt").mapPartitions(itor=>{
      val array=new ArrayBuffer[Worker]()
      while(itor.hasNext){
        val splited=itor.next().split(",")
        array.append(new Worker(splited(0),splited(2).toInt,splited(2)))
      }
      array.toIterator
    })
    import spark.implicits._
    //注册UDAF
    spark.udf.register("countUDF",new CountUDAF())
    val workDS=workerRDD.toDS()
    workDS.createOrReplaceTempView("worker")
    val resultDF=spark.sql("select countUDF(age) from worker")
    val resultDS=resultDF.as("WO")
    resultDS.show()

    spark.stop()

  }
}
-----------------------------------------------------------------------------------------------