Spark用户自定义函数UDAF

用户自定义函数

–1.UDF

package com.qf.sql.day02

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

object _06UserDefineFunction {
    def main(args: Array[String]): Unit = {
        val spark = SparkSession.builder().master("local").appName("test").getOrCreate()
        val df: DataFrame = spark.read.json("file:///D:\\Users\\Michael\\Documents\\IdeaProjects\\sz2003\\sz2003_sparksql01/data/emp.json")

        //维护一张临时表
        df.createTempView("emp")
        //注册函数
        spark.udf.register("getLevel",getLevel _)


       /* val sql =
            """
              |select ename,
              |job,
              |sal,
              |case when sal>3000 then 'level1'
              |when sal>1500 then 'level2'
              |else 'level3' end as level
              |from
              |emp
              |""".stripMargin*/
       val sql =
           """
             |select ename,
             |job,
             |sal,
             |getLevel(sal) as level
             |from
             |emp
             |""".stripMargin
        spark.sql(sql).show()
        spark.stop()
    }
    def getLevel(sal:Int)={
        if(sal>3000){
            "level1"
        }else if(sal>1500){
            "level2"
        }else{
            "level3"
        }
    }
}

–2.UDAF

package day02

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DataTypes, StructField, StructType}
import org.apache.spark.sql.{Dataset, Row, SparkSession}

/**
 * 用户自定义分析函数UDAF案例演示:
 *    自定义一个求平均值的函数：   多对一。
 * @param id
 * @param name
 * @param score
 */

case class Score(id:Int,name:String,score:Double)
object _03UDAF {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("udaf").master("local").getOrCreate()
    val list = List(
      Score(1,"张三",99.9),
      Score(2,"李四",88.9),
      Score(3,"小明",77.9),
      Score(1,"张三",91.9),
      Score(2,"李四",81.9),
      Score(3,"小明",71.9)
    )

    import spark.implicits._
    //scala的对象可以直接转DS或者是DF， 当然需要导入隐式转换
    val ds: Dataset[Score] = list.toDS()
    /*维护一张表*/
    ds.createTempView("tmp")
    //注册函数：
    spark.udf.register("myavg",new MyAvgUDAF)

    val sql = "select avg(score), myavg(score),name from tmp group by name"
    spark.sql(sql).show()
  }
}

/**
 * 1:需要继承 UserDefinedAggregateFunction
 * 2:重写方法
 */
class  MyAvgUDAF extends  UserDefinedAggregateFunction{
  /**
   *  指定用户自定义udaf输入参数的元数据
   * @return
   */
  override def inputSchema: StructType = {
    StructType(Array(StructField("score",DataTypes.DoubleType)))
  }

  /**
   * udaf自定义函数求解过程中的临时变量的数据类型
   *
   * 因为要求平均值，正常逻辑是求出总和以及个数，然后做除法运算，因此要有两个临时变量
   * @return
   */
  override def bufferSchema: StructType = {
    StructType(Array(
      StructField("sum",DataTypes.DoubleType),
      StructField("count",DataTypes.IntegerType)
    ))
  }

  /**
   * udaf返回值的数据类型
   * @return
   */
  override def dataType: DataType = DataTypes.DoubleType

  override def deterministic: Boolean = true

  /**
   * 临时变量是存储在buffer中，我们定义的buffer有两个元素，第一元素是sum，第二个元素是count
   * @param buffer
   */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer.update(0,0d) //设置第一个元素的的初始值为0d
    buffer.update(1,0) // 设置第二个元素的初始值为0
  }

  /**
   * 分区内的局部累加操作
   * @param buffer
   * @param input
   */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val score: Double = input.getAs[Double](0)
    buffer.update(0,buffer.getDouble(0)+score) //将新进来的行中的分数累加到第一个元素sum上
    buffer.update(1,buffer.getInt(1)+1)        //将第二个元素count累加一个1
  }

  /**
   * 分区间的累加操作
   * @param buffer1
   * @param buffer2
   */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1.update(0,buffer1.getDouble(0)+buffer2.getDouble(0))
    buffer1.update(1,buffer1.getInt(1)+buffer2.getInt(1))
  }
  /**
   * 用于计算输出结果
   * @param buffer
   * @return
   */
  override def evaluate(buffer: Row): Any = buffer.getDouble(0)/buffer.getInt(1)
}