排序窗口上的Spark Aggregator从不调用merge()？该用法是否可靠？

核心结论 #

针对你描述的场景：当使用org.apache.spark.sql.expressions.Aggregator实现自定义逻辑，并应用于unboundedPreceding到currentRow的排序窗口时，merge()函数确实不会被触发，聚合全程依赖reduce()。这种用法对于不支持merge操作的自定义算法是安全的，背后的机制可以从Spark窗口计算逻辑来解释。

机制原因 #

Spark处理unboundedPreceding到currentRow的有序前缀窗口时，会对每个分区内的数据按指定规则全量排序，之后以流式逐行累积的方式计算：

• 从分区第一行开始，将当前行输入与之前的累积状态通过reduce()合并，生成新状态
• 由于是严格按顺序单路径累积，不存在需要合并多个独立子状态的场景，因此merge()没有调用的必要

这和普通分组聚合（可能拆分分区并行计算后合并子结果）的逻辑完全不同，前缀窗口的特性决定了它只能是线性累积过程。

安全性与注意事项 #

1. 算法适配性：只要你的自定义逻辑是基于顺序依赖的逐行累积计算（比如滚动统计、状态机类逻辑、累加计算），完全可以只依赖reduce()完成全部计算——即使实现了merge()，在该场景下也不会被执行。
2. 数据规模验证：你提到的3亿行数据验证结果已经能佐证这一行为，Spark在该场景下不会触发merge逻辑，不会因merge未实现导致错误。
3. 潜在风险：
   • 必须严格保证窗口范围是unboundedPreceding到currentRow且窗口已排序，若后续调整为滑动窗口或其他范围，merge()可能会被触发，未实现的merge逻辑会直接报错。
   • 分区内数据量过大时需注意内存压力：聚合状态会持续持有到分区处理结束，若状态序列化体积过大，可能引发OOM，需要提前评估状态大小。

测试代码示例 #

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders, SparkSession}

case class InputRow(id: Int, value: Double)
case class AccumState(sum: Double, count: Int)

class RollingStatsAgg extends Aggregator[InputRow, AccumState, (Double, Double)] {
  override def zero: AccumState = AccumState(0.0, 0)

  override def reduce(acc: AccumState, input: InputRow): AccumState = {
    AccumState(acc.sum + input.value, acc.count + 1)
  }

  // 此方法在unboundedPreceding到currentRow窗口中不会被调用
  override def merge(acc1: AccumState, acc2: AccumState): AccumState = {
    throw new UnsupportedOperationException("Merge not supported for this aggregator")
  }

  override def finish(acc: AccumState): (Double, Double) = {
    (acc.sum, acc.sum / acc.count)
  }

  override def bufferEncoder: Encoder[AccumState] = Encoders.product[AccumState]
  override def outputEncoder: Encoder[(Double, Double)] = Encoders.tuple(Encoders.Double, Encoders.Double)
}

object TestWindowAgg {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("RollingTest").getOrCreate()
    import spark.implicits._

    val df = spark.range(0, 300000000).map(i => InputRow(i % 100, i.toDouble)).toDF()
    val rollingAgg = new RollingStatsAgg().toColumn.name("rolling_stats")

    val result = df
      .orderBy("id")
      .select("id", "value", rollingAgg)
    
    result.show()
    spark.stop()
  }
}

内容的提问来源于stack exchange，提问作者McPeanutbutter