能否在Spark DataFrame/Dataset API中结合累加器统计过滤记录数？

完全可以实现这个需求，核心是通过自定义累加器追踪每个过滤条件的通过记录数，下面是具体的实现方案：

实现思路 #

1. 自定义Map类型累加器：用键值对存储不同过滤条件对应的通过记录数，key为过滤条件标识，value为计数
2. 嵌入累加器更新逻辑：通过map算子先判断记录是否符合过滤条件，同步更新累加器，再用filter算子筛选出符合条件的记录
3. 触发Action操作：Spark是懒执行模型，必须执行count、write等Action才会触发累加器的实际更新

完整代码示例 #

1. 自定义累加器类

import org.apache.spark.util.AccumulatorV2
import scala.collection.mutable

class FilterCountAccumulator extends AccumulatorV2[(String, Boolean), Map[String, Long]] {
  private val countMap = mutable.HashMap[String, Long]().withDefaultValue(0L)

  override def isZero: Boolean = countMap.isEmpty
  override def copy(): AccumulatorV2[(String, Boolean), Map[String, Long]] = {
    val newAcc = new FilterCountAccumulator()
    newAcc.countMap ++= this.countMap
    newAcc
  }
  override def reset(): Unit = countMap.clear()
  override def add(v: (String, Boolean)): Unit = {
    if (v._2) countMap(v._1) += 1
  }
  override def merge(other: AccumulatorV2[(String, Boolean), Map[String, Long]]): Unit = {
    other.value.foreach { case (k, v) => countMap(k) += v }
  }
  override def value: Map[String, Long] = countMap.toMap
}

2. 结合DataFrame使用累加器

// 注册累加器
val filterAcc = new FilterCountAccumulator()
sparkSession.sparkContext.register(filterAcc, "FilterCountAccumulator")

// 定义过滤参数
val dateFromInclusive = "2024-01-01"
val dateToExclusive = "2024-02-01"
val origins = Array("A", "B", "C")

val resultDF = sparkSession.read
  .format("org.apache.spark.sql.delta.sources.DeltaDataSource")
  .load(path)
  // 处理日期过滤并统计
  .map(row => {
    val passDate = row.getAs[String]("ds") >= dateFromInclusive && row.getAs[String]("ds") < dateToExclusive
    filterAcc.add(("date_filter", passDate))
    (row, passDate)
  })
  .filter(_._2)
  .map(_._1)
  // 处理origin过滤并统计
  .map(row => {
    val passOrigin = origins.contains(row.getAs[String]("origin"))
    filterAcc.add(("origin_filter", passOrigin))
    (row, passOrigin)
  })
  .filter(_._2)
  .map(_._1)

// 触发Action，执行计算
resultDF.count()

// 输出统计结果
println("各过滤条件通过记录数：")
filterAcc.value.foreach { case (filterName, count) =>
  println(s"$filterName: $count")
}

注意事项 #

• 不要直接在where子句中更新累加器：where属于逻辑计划阶段，累加器更新需要在执行阶段触发，map+filter的组合更可靠
• 累加器是一次性累加：如果重复触发Action，累加器会重复计数，需要重新计算时记得调用reset()方法
• 自定义累加器必须正确实现merge方法：确保分布式环境下各个Executor的计数能正确合并

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