Scala Spark中groupBy分组后对colB列表执行自定义迭代计算的方法

实现思路 #

你需要的是分组内按顺序迭代计算的逻辑，直接用collect_list全量收集分组数据再处理在数据量大时存在内存溢出风险，更推荐用Spark的自定义聚合器(Aggregator) 实现，逐行处理分组内数据，无需全量缓存列表。

注意：如果你的分组内colB有固定排序要求，需要在聚合前先按对应字段排序，保证计算顺序和预期一致。

完整代码实现 #

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

// 定义聚合器的中间状态类型，仅存储当前迭代的res值即可
case class IterState(var res: Int)

object CustomIterAgg extends Aggregator[Int, IterState, Int] {
  // 初始状态，对应公式里的res=0
  override def zero: IterState = IterState(0)

  // 分区内逐行计算逻辑，对应循环里的单步计算规则
  override def reduce(buffer: IterState, value: Int): IterState = {
    buffer.res += value * (3 + buffer.res)
    buffer
  }

  // 多分区合并逻辑，迭代计算要求同分组数据落在同一分区，直接取有效状态即可
  override def merge(b1: IterState, b2: IterState): IterState = {
    if (b1.res == 0) b2 else b1
  }

  // 输出最终计算结果
  override def finish(reduction: IterState): Int = reduction.res

  // 中间状态编码器
  override def bufferEncoder: Encoder[IterState] = Encoders.product[IterState]

  // 输出结果编码器
  override def outputEncoder: Encoder[Int] = Encoders.scalaInt
}

object GroupIterCalcDemo {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("GroupIterCalc")
      .getOrCreate()
    import spark.implicits._

    // 构造样例数据
    val df = Seq(
      (1,3),
      (1,2),
      (2,4),
      (2,5),
      (2,1)
    ).toDF("colA", "colB")

    // 注册自定义聚合函数
    val calc_res = functions.udaf(CustomIterAgg)

    // 按colA分组聚合计算
    val result = df.groupBy("colA")
      .agg(calc_res($"colB").alias("colB"))

    // 输出结果
    result.show()
    /* 输出和预期完全一致：
    +----+----+
    |colA|colB|
    +----+----+
    |   1|  24|
    |   2|  78|
    +----+----+
     */

    spark.stop()
  }
}

小数据量可选实现：collect_list后遍历计算 #

如果你的业务数据量很小，也可以先收集分组内列表再遍历计算，代码更简洁：

import org.apache.spark.sql.functions.collect_list
import spark.implicits._

df.groupBy("colA")
  .agg(collect_list("colB").alias("colB_list"))
  .map(row => {
    val colA = row.getAs[Int]("colA")
    val colBList = row.getAs[Seq[Int]]("colB_list")
    var res = 0
    colBList.foreach(i => res += i * (3 + res))
    (colA, res)
  }).toDF("colA", "colB")
  .show()

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