如何获取Spark JDBC写入场景下的Spark指标数据

• 运行环境：Scala 2.11，Spark 2.4.4
• 现象：自定义SparkListener通过onTaskEnd回调读取taskMetrics.recordsWritten获取写入记录数时，写入Hive、Parquet可正常获取结果，但使用df.write.format("jdbc")写入时计数始终为0

JDBC数据源的原生写入逻辑未实现TaskMetrics写入计数上报：写入流程执行完批量插入后，没有调用incRecordsWritten接口更新指标值，和Parquet、Hive等内置文件数据源的指标上报逻辑存在实现差异，因此直接读取TaskMetrics拿不到有效计数。

方案1：自定义JDBC数据源，主动上报写入计数（精度最高，推荐） #

对原生JDBC数据源写入逻辑做轻量改造，不改动原有业务代码逻辑：

• 继承原生JDBC数据源类，重写数据写入的核心方法，在每个Executor端的task执行完批量写入后，统计当前task实际写入的行数，主动调用TaskContext的指标接口上报
• 写入时format指定为自定义的JDBC数据源类，原有SparkListener的onTaskEnd逻辑不需要任何调整，即可正常读取到recordsWritten的准确值
  核心上报代码：

// 每个task批量写入完成后执行
val writeRowCount = statement.executeBatch().sum // 统计当前批次实际写入行数
TaskContext.get().taskMetrics().incRecordsWritten(writeRowCount) // 主动上报指标

方案2：注册QueryExecutionListener，从执行计划指标中统计写入行数 #

无需改造数据源，直接通过Spark内置的查询执行监听器采集：

• 初始化SparkSession后注册自定义QueryExecutionListener，在作业执行成功的回调中获取当前作业的物理执行计划
• 遍历执行计划节点，筛选出JDBC写入对应的节点，累加节点上的numOutputRows指标值，即可得到全作业的总写入行数
  监听器注册代码示例：

spark.listenerManager.register(new QueryExecutionListener {
  override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
    var totalJdbcWriteRows = 0L
    qe.executedPlan.foreach { node =>
      // 筛选JDBC写入对应的执行节点
      if (node.nodeName.contains("SaveIntoDataSourceCommand") && node.metadata.get("format").contains("jdbc")) {
        totalJdbcWriteRows += node.metrics.get("numOutputRows").map(_.value).getOrElse(0L)
      }
    }
    // 此处处理采集到的总写入行数即可
  }

  override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = {
    // 写入失败的异常逻辑可自行扩展
  }
})

注意：该方案需要做节点去重判断，避免AQE自适应执行或者计划节点重复遍历导致计数翻倍。

方案3：原有SparkListener扩展兜底逻辑 #

如果不想新增监听器、也不想改造数据源，可以在原有SparkListener的onTaskEnd逻辑中增加判断：如果当前task属于JDBC写入的最终stage，直接取该task的inputMetrics.recordsRead值作为该task的写入行数，累加所有对应task的数值即为总写入行数。该方案仅适用于写入前没有额外过滤、转换算子的场景，计数存在一定误差，适合精度要求不高的兜底场景。

提示：Spark 2.4.x版本的JDBC数据源未内置任何写入行数的原生指标，所有非侵入式方案都属于间接统计，对精度要求高的场景优先选择方案1。

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