Spark SQL UDAF示例

UDAF：用户自定义聚合函数

Scala 2.10.7，spark 2.0.0

package UDF_UDAF

import java.util

import org.apache.spark.SparkConf
import org.apache.spark.sql.{Row, RowFactory, SparkSession}
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DataTypes, StructField, StructType}

class UDAF extends UserDefinedAggregateFunction {
  /**
    * 指定输入字段的字段及类型
    */
  override def inputSchema: StructType =
    DataTypes.createStructType(Array(DataTypes.createStructField("namexxx",DataTypes.StringType,true)))

  /**
    * 在进行聚合操作的时候所要处理的数据的结果的类型
    * */
  override def bufferSchema: StructType =
    DataTypes.createStructType(Array(DataTypes.createStructField("buffer",DataTypes.IntegerType,true)))

  /**
    * 指定UDAF计算后返回的结果类型
    * @return
    */
  override def dataType: DataType = DataTypes.IntegerType

  /**
    * 确保一致性 一般用true,用以标记针对给定的一组输入，UDAF是否总是生成相同的结果。
    */
  override def deterministic: Boolean = true

  /**
    * 初始化一个内部的自己定义的值,在Aggregate之前每组数据的初始化结果
    */
  override def initialize(buffer: MutableAggregationBuffer): Unit = buffer.update(0,0)

  /**
    * 更新 可以认为一个一个地将组内的字段值传递进来 实现拼接的逻辑
    * buffer.getInt(0)获取的是上一次聚合后的值
    * 相当于map端的combiner，combiner就是对每一个map task的处理结果进行一次小聚合
    * 大聚和发生在reduce端.
    * 这里即是:在进行聚合的时候，每当有新的值进来，对分组后的聚合如何进行计算
    */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = buffer.update(0, buffer.getInt(0)+1)

  /**
    * 合并 update操作，可能是针对一个分组内的部分数据，在某个节点上发生的 但是可能一个分组内的数据，会分布在多个节点上处理
    * 此时就要用merge操作，将各个节点上分布式拼接好的串，合并起来
    * buffer1.getInt(0) : 大聚合的时候 上一次聚合后的值
    * buffer2.getInt(0) : 这次计算传入进来的update的结果
    * 这里即是：最后在分布式节点完成后需要进行全局级别的Merge操作
    */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = buffer1.update(0, buffer1.getInt(0)+buffer2.getInt(0))

  /**
    * 最后返回一个和dataType方法的类型要一致的类型，返回UDAF最后的计算结果
    */
  override def evaluate(buffer: Row): Any = buffer.getInt(0)
}

object UDAF{
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setMaster("local").setAppName("udaf")
    val sparkSession = SparkSession.builder().config(conf).config("spark.sql.warehouse.dir","/test/warehouse").getOrCreate()
    val sc = sparkSession.sparkContext

    val parallelize = sc.parallelize(Array("zhangsan","lisi","wanger","zhaosi","zhangsan","lisi"))
    val rowRDD = parallelize.map(s=>RowFactory.create(s))

    val fields = new util.ArrayList[StructField]()
    fields.add(DataTypes.createStructField("name",DataTypes.StringType,true))
    val schema = DataTypes.createStructType(fields)

    val df = sparkSession.createDataFrame(rowRDD, schema)
    df.createOrReplaceTempView("user")

    sparkSession.udf.register("StringCount",new UDAF())

    sparkSession.sql("select name, StringCount(name) as StrCount from user group by name").show()

    sparkSession.stop()

  }
}