Spark中的UDAF简介及其Stage task

最新推荐文章于 2025-05-29 10:41:12 发布

我叫龙翔天翼

最新推荐文章于 2025-05-29 10:41:12 发布

阅读量350

点赞数

CC 4.0 BY-SA版权

分类专栏： Spark 大数据机器学习文章标签： spark 大数据分布式 UDAF stage task

本文链接：https://blog.youkuaiyun.com/weixin_40901056/article/details/98986219

大数据同时被 3 个专栏收录

6 篇文章

订阅专栏

机器学习

4 篇文章

订阅专栏

Spark

2 篇文章

订阅专栏

本文介绍了Spark中的用户定义累积聚合函数(UDAF)的运作原理，重点解析了其在执行过程中涉及的两个阶段——partial_merge和merge。partial_merge阶段在本地进行merge操作，形成窄依赖，任务数与上一步分区数对应。而merge阶段则涉及到不同分区数据的合并，需要shuffle，任务数由`spark.default.parallelism`设定决定。内容详细阐述了UDAF在大数据处理中的应用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Spark中的UDAF简介及其Stage

udaf操作会分为两个stage：

1. partial_merge: 本地进行merge，是一种窄依赖。tasks数量取决于上一步的partitions。
1. merge：不同partition的数据进行merge，是一种宽依赖，需要shuffle，因此tasks数量取决于设置的值spark.default.parallelism

class MyAvg extends UserDefinedAggregateFunction {
    // Input type
    def inputSchema: org.apache.spark.sql.types.StructType = StructType(StructField("value", DoubleType) :: Nil)
    // This is the internal fields you keep for computing your aggregate. 计算缓存
    def bufferSchema: StructType = StructType( StructField("count", LongType) :: StructField("sum", DoubleType) :: Nil
    )
    // Return type
    def dataType: DataType = DoubleType
    //幂等性
    def deterministic: Boolean = true
    //初始值
    def initialize(buffer: MutableAggregationBuffer): Unit = {   buffer(0) = 0L //计数count
      buffer(1) = 0.0 //求和sum
    }
    //根据给定输入，更新缓存buffer
    def update(buffer: MutableAggregationBuffer,input: Row): Unit = { 
        buffer(0) = buffer.getAs[Long](0) + 1
        buffer(1) = buffer.getAs[Double](1) + input.getAs[Double](0)
    }
    //合并merge两个buffer：包括计算partial和合并partial
    def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = { buffer1(0) = buffer1.getAs[Long](0) + buffer2.getAs[Long](0) buffer1(1) = buffer1.getAs[Double](1) + buffer2.getAs[Double](1)
    }
    //最终输出值
    def evaluate(buffer: Row): Any = { buffer.getDouble(1) / buffer.getLong(0)
    } 
  }

使用：

val df = spark.sql("select user, num from table")
df.groupBy("user")
  .agg(MyAvg(col("num")).as("avg"))