Spark Streaming读取Kafka高效化：批量行数统计致重复消费问题

问题场景 #

我开发了一个从Kafka获取数据并保存至数据库的应用，初始代码如下：

spark.readStream
  .format("kafka")
  .options(options)
  .load()
  .writeStream
  .trigger(Trigger.ProcessingTime(20000))
  .foreachBatch({ (batch: DataFrame, _: Long) =>
    val rowsCount = batch.count
    saveBatch(batch)
    println(s"Saved $rowsCount rows")
  })
  .start()

在Spark UI的Structured Streaming标签中，流处理速率为100K行/秒。

若移除行数统计，代码修改为：

.foreachBatch({ (batch: DataFrame, _: Long) =>
    saveBatch(batch)
  })
  .start()

此时处理速率降至50K行/秒。由于未使用缓存，推测统计批量行数与保存操作导致从Kafka重复读取数据（仅保存时仅读取一次）。我难以相信行数统计这类简单操作会引发重复读取，请问是否存在无需缓存即可获取批量行数的方法？

分析原因 #

Spark DataFrame采用惰性求值机制，count()和saveBatch()会触发两次独立的作业执行——第一次作业读取Kafka数据并统计行数，第二次作业再次读取Kafka数据并写入数据库，这就是性能减半的核心原因。

无需缓存的解决方案 #

方案1：写入时同步统计行数 #

通过mapPartitions在分区级别同时完成数据写入和行数统计，再汇总总行数，整个过程仅触发一次作业：

// 封装写入+统计的逻辑
def saveAndCount(batch: DataFrame): Long = {
    batch.rdd.mapPartitions { iter =>
        val partitionCount = iter.size
        // 执行当前分区的数据写入逻辑（比如批量插入数据库）
        savePartition(iter)
        Iterator(partitionCount)
    }.sum().toLong
}

// 在foreachBatch中调用
.foreachBatch({ (batch: DataFrame, _: Long) =>
    val rowsCount = saveAndCount(batch)
    println(s"Saved $rowsCount rows")
})

这种方式让数据读取、统计、写入在同一个作业中完成，避免了重复读取Kafka的开销。

方案2：利用Spark内置监控事件获取行数 #

注册StreamingQueryListener监听批次完成事件，直接从Spark的监控指标中获取输入行数，无需额外计算：

// 注册监听事件
spark.streams.addListener(new StreamingQueryListener() {
    override def onQueryStarted(event: QueryStartedEvent): Unit = {}
    override def onQueryTerminated(event: QueryTerminatedEvent): Unit = {}
    override def onQueryProgress(event: QueryProgressEvent): Unit = {
        val rowsCount = event.progress.numInputRows
        println(s"Saved $rowsCount rows")
    }
})

// 原写入逻辑无需统计代码
.foreachBatch({ (batch: DataFrame, _: Long) =>
    saveBatch(batch)
})

numInputRows是Spark Structured Streaming内置的指标，直接记录了当前批次从Kafka读取的行数，完全不需要额外触发计算。

补充说明 #

如果坚持要分离统计和写入逻辑，最直接的优化还是对DataFrame执行batch.cache()，缓存后两次操作会复用同一个RDD，避免重复读取。但如果不想使用缓存，上述两种方案更符合需求。

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