如何在Scala中按组计算二分类器的评估指标？

我有一个Spark DataFrame，存储了多个二分类任务的score和label数据，示例结构如下：

我的目标是按problem分组，计算每个组的二分类评估指标（比如AreaUnderROC），期望得到的结果格式如下：

| problem | areaUnderROC |
| a | 0.83 |
| b | 0.68 |
| ... | ... |

我想使用类似df.groupBy("problem").agg(getMetrics)的写法，但不确定如何基于Spark的Aggregators来实现这个getMetrics函数，请问有没有具体的实现建议？

针对你的需求，我们可以通过自定义Spark Aggregator来实现按分组计算二分类指标的功能。Aggregator允许我们自定义分组聚合逻辑，完美适配“按problem收集score和label，再计算AUC”的场景。下面是具体的实现步骤和代码示例：

1. 基础版：使用Spark ML的Evaluator计算AUC #

首先，我们可以借助Spark ML的BinaryClassificationEvaluator来封装AUC计算逻辑，实现一个基础版的Aggregator：

步骤1：导入依赖并定义数据结构 #

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
import org.apache.spark.ml.linalg.Vectors

// 定义封装score和label的样例类
case class ScoreLabel(score: Double, label: Boolean)

步骤2：实现自定义Aggregator #

class ROCAUCAggregator extends Aggregator[ScoreLabel, List[(Double, Double)], Double] {
  // 初始化缓冲区：空的(score, label)列表（label转成Double类型）
  override def zero: List[(Double, Double)] = List.empty

  // 将单条输入数据添加到缓冲区
  override def reduce(buffer: List[(Double, Double)], input: ScoreLabel): List[(Double, Double)] = {
    val labelAsDouble = if (input.label) 1.0 else 0.0
    (input.score, labelAsDouble) :: buffer
  }

  // 合并两个分组的缓冲区数据
  override def merge(b1: List[(Double, Double)], b2: List[(Double, Double)]): List[(Double, Double)] = b1 ++ b2

  // 基于缓冲区的所有数据计算AUC
  override def finish(reduction: List[(Double, Double)]): Double = {
    // 处理数据量不足的情况（比如只有正样本或负样本）
    if (reduction.size < 2) return 0.5

    // 分离scores和labels并转成Spark ML需要的Vector格式
    val (scores, labels) = reduction.unzip
    val scoreVector = Vectors.dense(scores.toArray)
    val labelVector = Vectors.dense(labels.toArray)

    // 使用BinaryClassificationEvaluator计算AUC
    val evaluator = new BinaryClassificationEvaluator()
      .setRawPredictionCol("score")
      .setLabelCol("label")
      .setMetricName("areaUnderROC")

    // 临时构造DataFrame完成评估
    val tempDF = spark.createDataFrame(Seq((scoreVector, labelVector))).toDF("score", "label")
    evaluator.evaluate(tempDF)
  }

  // 缓冲区的编码器（用Kryo序列化列表）
  override def bufferEncoder: Encoder[List[(Double, Double)]] = Encoders.kryo[List[(Double, Double)]]

  // 输出结果的编码器（Double类型）
  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

步骤3：使用自定义聚合函数 #

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

// 注册自定义聚合函数
val rocAUCAgg = udaf(new ROCAUCAggregator())

// 按problem分组计算AUC
val resultDF = df.groupBy("problem")
  .agg(rocAUCAgg(struct($"score", $"label")).alias("areaUnderROC"))

// 查看结果
resultDF.show()

2. 优化版：手动计算AUC提升性能 #

基础版中每次计算都要创建临时DataFrame，在大数据场景下会有性能开销。我们可以手动实现AUC的计算逻辑，避免依赖Spark ML的Evaluator：

class OptimizedROCAUCAggregator extends Aggregator[ScoreLabel, List[(Double, Double)], Double] {
  override def zero: List[(Double, Double)] = List.empty

  override def reduce(buffer: List[(Double, Double)], input: ScoreLabel): List[(Double, Double)] = {
    val labelAsDouble = if (input.label) 1.0 else 0.0
    (input.score, labelAsDouble) :: buffer
  }

  override def merge(b1: List[(Double, Double)], b2: List[(Double, Double)]): List[(Double, Double)] = b1 ++ b2

  override def finish(reduction: List[(Double, Double)]): Double = {
    val sortedData = reduction.sortBy(-_._1) // 按score降序排序
    val totalPos = sortedData.count(_._2 == 1.0)
    val totalNeg = sortedData.size - totalPos

    // 没有正样本或负样本时，返回随机猜测的AUC值0.5
    if (totalPos == 0 || totalNeg == 0) return 0.5

    var tp = 0.0
    var fp = 0.0
    var prevScore = Double.NaN
    var auc = 0.0
    var lastTp = 0.0
    var lastFp = 0.0

    // 遍历排序后的数据计算AUC
    for ((score, label) <- sortedData) {
      if (score != prevScore) {
        // 累加当前区间的AUC贡献
        auc += (fp - lastFp) * (tp + lastTp) / 2.0
        prevScore = score
        lastTp = tp
        lastFp = fp
      }
      if (label == 1.0) tp += 1.0 else fp += 1.0
    }
    // 加上最后一个区间的贡献
    auc += (fp - lastFp) * (tp + lastTp) / 2.0
    // 归一化到0-1范围
    auc / (totalPos * totalNeg)
  }

  override def bufferEncoder: Encoder[List[(Double, Double)]] = Encoders.kryo[List[(Double, Double)]]
  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}

使用优化版的方式和基础版一致：

val optimizedRocAUCAgg = udaf(new OptimizedROCAUCAggregator())
val resultDF = df.groupBy("problem")
  .agg(optimizedRocAUCAgg(struct($"score", $"label")).alias("areaUnderROC"))

关键注意事项 #

• 数据边界处理：当分组内没有正样本或负样本时，AUC返回0.5（符合随机分类器的预期）
• 性能考量：优化版手动计算AUC避免了Spark ML的额外开销，适合数据量较大的场景
• 扩展性：如果需要计算其他二分类指标（比如Precision-Recall AUC），可以修改finish方法中的计算逻辑

内容的提问来源于stack exchange，提问作者K.Chen