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() } }