SparkSQL中的自定义函数UDF

在Spark中，也支持Hive中的自定义函数。自定义函数大致可以分为三种：

• UDF(User-Defined-Function)，即最基本的自定义函数，类似to_char,to_date等
• UDAF（User- Defined Aggregation Funcation），用户自定义聚合函数，类似在group by之后使用的sum,avg等
• UDTF(User-Defined Table-Generating Functions),用户自定义生成函数，有点像stream里面的flatMap

Notes: 自定义一个UDF函数需要继承UserDefinedAggregateFunction类，并实现其中的8个方法

自定义函数UDF的编写

GetDistinctCityUDF.scala

package com.UDF.TestUDF

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

/**
  * 获取不同城市的UDF函数
  */
object GetDistinctCityUDF extends UserDefinedAggregateFunction{

  /**
    * 输入的数据类型
    */
  override def inputSchema: StructType = StructType(
    StructField("status",StringType,true) :: Nil
  )

  /**
    * 缓存字段类型
    */
  override def bufferSchema: StructType = {
    StructType(
      Array(
        StructField("buffer_city_info",StringType,true)
      )
    )
  }

  /**
    * 输出结果类型
    */
  override def dataType: DataType = StringType

  /**
    * 输入类型和输出类型是否一种
    */
  override def deterministic: Boolean = true

  /**
    * 对辅助字段进行初始化
    */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer.update(0,"")
  }

  /**
    * 修改辅助字段的值
    */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {

    // 获取最后一次的值
    var last_str = buffer.getString(0)

    // 获取当前的值
    val current_str = input.getString(0)

    if (!last_str.contains(current_str)) {
      if (last_str.equals("")) {
        last_str = current_str
      } else {
        last_str += "," + current_str
      }
    }
    buffer.update(0,last_str)
  }

  /**
    * 对分区结果进行合并
    * buffer1是机器Slave1上的结果
    * buffer2是机器Slave2上的结果
    */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    var buf1 = buffer1.getString(0)
    val buf2 = buffer2.getString(0)

    // 将buf2里面存在的数据，而在buf1里面没有的数据，追加到buf1
    // buf2的数据按照 逗号 进行切分
    for (s <- buf2.split(",")) {
      if (!buf1.contains(s)) {
        if (buf1.equals("")) {
          buf1 = s
        } else {
          buf1 += s
        }
      }
    }
    buffer1.update(0,buf1)
  }

  /**
    * 最终的计算结果
    */
  override def evaluate(buffer: Row): Any = {
    buffer.getString(0)
  }

}

注册自定义的UDF函数为临时函数

TestUDFApp.scala

package com.UDF.TestUDF

import org.apache.spark.sql.SparkSession

/**
  * 注册自定义的UDF函数为临时函数
  */
object TestUDFApp extends App {

  /**
    * 第一步： 创建程序入口
    */
  val spark = SparkSession
    .builder()
    .appName("TestUDFApp")
    .master("local[2]")
    .getOrCreate()

  /**
    * 注册成为临时函数
    */
  spark.udf.register("get_distinct_city",GetDistinctCityUDF)

  /**
    * 注册成为临时函数
    */
  spark.udf.register("get_product_status", (str:String) => {
    var status = 0
    for (s <- str.split(",")) {
      if (s.contains("product_status")) {
        status = s.split(":")(1).toInt
      }
    }
  })
}