spark sql 自定义函数实例（udf、udaf、udtf）

Spark SQL中的自定义函数（UDF、UDAF、UDTF）是用户自己定义的函数，可以用于对数据进行处理和转换。下面是一些自定义函数的实例：

1. UDF（User-Defined Function）：用户自定义函数，可以将一个或多个输入参数转换为输出值。例如，我们可以定义一个UDF来计算两个数的和：

import org.apache.spark.sql.functions.udf

val sumUDF = udf((a: Int, b: Int) => a + b)

val df = Seq((1, 2), (3, 4)).toDF("a", "b")

df.select(sumUDF($"a", $"b")).show()

2. UDAF（User-Defined Aggregate Function）：用户自定义聚合函数，可以对一组数据进行聚合操作，例如求和、平均值等。例如，我们可以定义一个UDAF来计算一组数的平均值：

import org.apache.spark.sql.expressions.MutableAggregationBuffer
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction
import org.apache.spark.sql.types._

class AvgUDAF extends UserDefinedAggregateFunction {
  // 输入数据类型
  def inputSchema: StructType = StructType(StructField("value", DoubleType) :: Nil)

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

  // 输出数据类型
  def dataType: DataType = DoubleType

  // 是否是确定性的
  def deterministic: Boolean = true

  // 初始化聚合缓冲区
  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0.0
    buffer(1) = 0L
  }

  // 更新聚合缓冲区
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getDouble(0) + input.getDouble(0)
    buffer(1) = buffer.getLong(1) + 1L
  }

  // 合并聚合缓冲区
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  // 计算最终结果
  def evaluate(buffer: Row): Any = {
    buffer.getDouble(0) / buffer.getLong(1)
  }
}

val avgUDAF = new AvgUDAF()

val df = Seq(1.0, 2.0, 3.0, 4.0).toDF("value")

df.agg(avgUDAF($"value")).show()

3. UDTF（User-Defined Table-Generating Function）：用户自定义表生成函数，可以将一个或多个输入参数转换为一个表。例如，我们可以定义一个UDTF来将一个字符串拆分成多个单词：

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.{explode, udf}
import org.apache.spark.sql.types._

class SplitUDTF extends UserDefinedFunction {
  // 输入数据类型
  def inputSchema: StructType = StructType(StructField("value", StringType) :: Nil)

  // 输出数据类型
  def dataType: DataType = ArrayType(StringType)

  // 是否是确定性的
  def deterministic: Boolean = true

  // 计算结果
  def apply(value: Row): Any = {
    value.getString(0).split(" ")
  }
}

val splitUDTF = udf(new SplitUDTF(), ArrayType(StringType))

val df = Seq("hello world", "spark sql").toDF("value")

df.select(explode(splitUDTF($"value"))).show()