Scala下Spark DataFrame替代RDD aggregateByKey实现分组相关性计算

核心思路 #

原RDD方案通过aggregateByKey维护自定义CovCounter累积分组内统计量，转换为DataFrame方案时，我们用**自定义聚合函数（Aggregator/UDAF）**替代aggregateByKey实现分组统计量累积，再通过UDF展开结果格式，完全规避RDD操作。

步骤1：定义统计量结构与计算逻辑 #

先定义保存协方差/相关系数所需统计量的类，包含单样本更新和分组合并逻辑：

case class CovStats(n: Long, sumX: Double, sumY: Double, sumXY: Double, sumX2: Double, sumY2: Double) {
  // 新增单一样本（自动忽略NaN值）
  def add(x: Double, y: Double): CovStats = {
    if (x.isNaN || y.isNaN) this
    else CovStats(
      n + 1,
      sumX + x,
      sumY + y,
      sumXY + x * y,
      sumX2 + x * x,
      sumY2 + y * y
    )
  }

  // 合并两个分组的统计量
  def merge(other: CovStats): CovStats = {
    CovStats(
      n + other.n,
      sumX + other.sumX,
      sumY + other.sumY,
      sumXY + other.sumXY,
      sumX2 + other.sumX2,
      sumY2 + other.sumY2
    )
  }

  // 计算协方差
  def covariance: Double = {
    if (n < 2) Double.NaN
    else (sumXY - sumX * sumY / n) / (n - 1)
  }

  // 计算皮尔逊相关系数
  def correlation: Double = {
    if (n < 2) Double.NaN
    else {
      val cov = covariance
      val stdX = math.sqrt((sumX2 - sumX * sumX / n) / (n - 1))
      val stdY = math.sqrt((sumY2 - sumY * sumY / n) / (n - 1))
      if (stdX == 0 || stdY == 0) Double.NaN
      else cov / (stdX * stdY)
    }
  }
}

步骤2：实现自定义Aggregator（类型安全UDAF） #

Spark 3.x+推荐用类型安全的Aggregator实现自定义聚合，对应原RDD的seqOp和combOp逻辑：

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

class CovCorrAggregator(columnCount: Int, useCorrelation: Boolean) extends Aggregator[Row, Map[(Int, Int), CovStats], Map[(Int, Int), Double]] {
  // 初始化零值：所有列对的统计量初始化为空状态
  override def zero: Map[(Int, Int), CovStats] = {
    (for (i <- 0 until columnCount; j <- i until columnCount) 
      yield (i, j) -> CovStats(0L, 0.0, 0.0, 0.0, 0.0, 0.0)).toMap
  }

  // 序列操作：处理单条数据，更新列对统计量
  override def reduce(buffer: Map[(Int, Int), CovStats], row: Row): Map[(Int, Int), CovStats] = {
    val values = (0 until columnCount).map(row.getDouble(_)).toArray
    buffer.map { case ((i, j), stats) =>
      (i, j) -> stats.add(values(i), values(j))
    }
  }

  // 合并操作：合并两个分组的统计量缓存
  override def merge(b1: Map[(Int, Int), CovStats], b2: Map[(Int, Int), CovStats]): Map[(Int, Int), CovStats] = {
    b1.map { case ((i, j), stats1) =>
      (i, j) -> stats1.merge(b2((i, j)))
    }
  }

  // 最终转换：将统计量转为协方差或相关系数
  override def finish(reduction: Map[(Int, Int), CovStats]): Map[(Int, Int), Double] = {
    reduction.map { case ((i, j), stats) =>
      (i, j) -> (if (useCorrelation) stats.correlation else stats.covariance)
    }
  }

  // 中间缓存的序列化编码器
  override def bufferEncoder: Encoder[Map[(Int, Int), CovStats]] = Encoders.kryo
  // 输出结果的序列化编码器
  override def outputEncoder: Encoder[Map[(Int, Int), Double]] = Encoders.kryo
}

步骤3：完整DataFlow实现 #

将Aggregator注册为UDAF，按分组聚合后展开结果：

import org.apache.spark.sql.functions.{udaf, struct, explode, lit}
import org.apache.spark.sql.Row

// 1. 过滤分组列含null的行（与原逻辑完全一致）
val cleanDF = selectedDF.na.drop("any", groupByColumns)

// 2. 注册UDAF，传入要计算的列数和是否计算相关系数
val covCorrUDAF = udaf(new CovCorrAggregator(columns.length, useCorrelation))

// 3. 分组聚合：将所有待计算列打包为struct传入UDAF
val valueColumns = columns.map(col)
val aggregatedDF = cleanDF
  .groupBy(groupByColumns.map(col): _*)
  .agg(covCorrUDAF(struct(valueColumns: _*)).alias("cov_corr_map"))

// 4. 定义UDF，将聚合后的统计量Map转换为目标输出行结构
val mapToRowsUDF = udf((covCorrMap: Map[(Int, Int), Double], groupValues: Row, columnNames: Array[String]) => {
  columnNames.indices.map { ind =>
    // 利用协方差/相关系数的对称性，获取当前列的所有配对结果
    val pairedValues = columnPairedIndicesAll(ind).map { case (i, j) =>
      val key = if (i <= j) (i, j) else (j, i)
      covCorrMap.getOrElse(key, Double.NaN)
    }
    Row.fromSeq(groupValues.toSeq ++ Seq(columnNames(ind)) ++ pairedValues)
  }
})

// 5. 展开结果并匹配目标Schema
val finalDF = aggregatedDF
  .select(
    struct(groupByColumns.map(col): _*).alias("group_values"),
    $"cov_corr_map",
    lit(columns).alias("column_names")
  )
  .withColumn("output_rows", mapToRowsUDF($"cov_corr_map", $"group_values", $"column_names"))
  .select(explode($"output_rows").alias("output_row"))
  .select($"output_row.*")
  .toDF(outputSchema.fieldNames: _*)

关键说明 #

• UDAF替代aggregateByKey：CovCorrAggregator的reduce对应原RDD的seqOp，merge对应combOp，zero对应初始的zeroCounters，逻辑完全对齐。
• 结果展开逻辑：通过UDF将聚合后的统计量Map转换为原代码要求的行结构，再用explode展开多行，最后匹配目标Schema。
• 性能与类型安全：使用Aggregator比传统UserDefinedAggregateFunction更类型安全，Spark会对UDAF进行优化，性能接近原RDD方案。

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