如何获取单Spark Session下各Structured Streaming流的读写记录数？

我来帮你梳理下针对这个场景的可行方案，先聊聊你之前尝试的方法为什么踩坑，再给出靠谱的解决思路：

先解释下你遇到的问题根源 #

• SparkListener.onTaskEnd()失效：Structured Streaming的Sink实现很多时候不会在Task层面上报outputMetrics，而且Task事件本身无法直接关联到具体的StreamingQuery，所以拿不到对应流的写入数也很正常。
• 累加器统计的是待写入数：累加器在map阶段统计的是经过转换后的记录数，但这只是“准备写入”的数量，没法反映Sink实际成功写入的条数（比如Sink可能过滤部分数据、写入失败重试，或者有批量写入的原子性保证，只有全部成功才算写入）。
• StreamingQueryListener默认无写入指标：原生的Listener确实只暴露了numInputRows这类输入指标，因为写入指标依赖具体Sink的实现，Spark没法统一收集。

可行的解决方案 #

1. 自定义Sink+StreamingQueryListener：精准统计实际写入数 #

这是最可靠的方案，因为只有Sink本身知道到底成功写入了多少条数据。你可以包装原生Sink，在写入逻辑里统计记录数，再关联到对应的查询：

举个Scala的例子，包装FileSink来统计写入数：

import org.apache.spark.sql.execution.streaming.Sink
import org.apache.spark.sql.streaming.OutputMode
import org.apache.spark.sql.{DataFrame, SQLContext}
import org.apache.spark.SparkEnv

class CountingFileSink(delegate: Sink, queryName: String) extends Sink {
  // 用Spark Metrics注册一个计数器，关联查询名
  private val counter = SparkEnv.get.metricsSystem.counter(s"structured-streaming", s"${queryName}_records_written")

  override def addBatch(batchId: Long, data: DataFrame): Unit = {
    // 统计当前批次实际写入的记录数（根据Sink逻辑调整，这里以count为例）
    val writtenCount = data.count()
    counter.inc(writtenCount)
    // 调用原生Sink的写入逻辑
    delegate.addBatch(batchId, data)
  }
}

// 使用方式
import org.apache.spark.sql.execution.streaming.FileSink
import org.apache.spark.sql.streaming.Trigger

val fileSink = new FileSink(
  spark.sqlContext,
  "/path/to/s3",
  Map.empty,
  OutputMode.Append()
)

val query = df.writeStream
  .foreach(new CountingFileSink(fileSink, "write-s3"))
  .queryName("write-s3")
  .trigger(Trigger.ProcessingTime("1 minute"))
  .start()

然后结合StreamingQueryListener，在事件中关联查询与计数器，统一收集指标：

class QueryMetricsListener extends StreamingQueryListener {
  override def onQueryStarted(event: StreamingQueryListener.QueryStartedEvent): Unit = {}

  override def onQueryProgress(event: StreamingQueryListener.QueryProgressEvent): Unit = {
    val progress = event.progress
    val inputRows = progress.numInputRows
    // 从Metrics系统获取对应查询的写入计数器值
    val writtenCounter = SparkEnv.get.metricsSystem.getCounter(s"structured-streaming", s"${progress.name}_records_written")
    println(s"Query ${progress.name}: Input rows = $inputRows, Written rows = ${writtenCounter.value}")
  }

  override def onQueryTerminated(event: StreamingQueryListener.QueryTerminatedEvent): Unit = {}
}

// 注册Listener
spark.streams.addListener(new QueryMetricsListener)

2. 利用StreamingQueryProgress的自定义字段 #

如果你不想依赖Metrics系统，也可以在自定义Sink中把写入数存入一个与查询名绑定的全局Map，然后在StreamingQueryListener的onQueryProgress事件里读取这个Map的值：

// 全局存储查询的写入记录数
object QueryMetricsStore {
  private val writtenCounts = scala.collection.mutable.Map[String, Long]()

  def update(queryName: String, count: Long): Unit = {
    writtenCounts.update(queryName, writtenCounts.getOrElse(queryName, 0L) + count)
  }

  def get(queryName: String): Long = writtenCounts.getOrElse(queryName, 0L)
}

// 自定义Sink
class CountingJdbcSink(delegate: Sink, queryName: String) extends Sink {
  override def addBatch(batchId: Long, data: DataFrame): Unit = {
    delegate.addBatch(batchId, data)
    // 假设JdbcSink写入后返回成功条数，这里简化为count
    val writtenCount = data.count()
    QueryMetricsStore.update(queryName, writtenCount)
  }
}

// Listener中读取
class QueryMetricsListener extends StreamingQueryListener {
  override def onQueryProgress(event: StreamingQueryListener.QueryProgressEvent): Unit = {
    val progress = event.progress
    val inputRows = progress.numInputRows
    val writtenRows = QueryMetricsStore.get(progress.name)
    println(s"Query ${progress.name}: Input rows = $inputRows, Written rows = $writtenRows")
  }

  // 其他方法省略...
}

3. 针对特定Sink的原生指标 #

如果你用的是成熟的第三方Sink，它们通常会自带写入指标：

• Delta Lake：可以通过delta.sink.numRecordsWritten指标获取写入数，在Spark UI的Metrics页面或REST API中可查看。
• Kafka Sink：Kafka Producer本身会暴露record-sent-total这类指标，可通过Kafka的Metrics系统关联到对应的Spark查询。

总结 #

最通用、精准的方案是自定义Sink结合Spark Metrics和StreamingQueryListener，既能区分不同的流，又能准确获取实际写入的记录数。如果是特定Sink，优先用它原生的指标会更省心。

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