Question

我正在使用onTaskEnd Spark侦听器来获取写入文件的记录数，如下所示：

import spark.implicits._
import org.apache.spark.sql._
import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}

var recordsWritten: Long = 0L

val rowCountListener: SparkListener = new SparkListener() {
  override def onTaskEnd(taskEnd: SparkListenerTaskEnd) {
    synchronized {
      recordsWritten += taskEnd.taskMetrics.outputMetrics.recordsWritten
    }
  }
}

def rowCountOf(proc: => Unit): Long = {
  recordsWritten = 0L
  spark.sparkContext.addSparkListener(rowCountListener)
  try {
    proc
  } finally {
    spark.sparkContext.removeSparkListener(rowCountListener)
  }
  recordsWritten
}

val rc = rowCountOf { (1 to 100).toDF.write.csv(s"test.csv") }
println(rc)

=> 100

但是，尝试在线程中运行多个操作显然会中断：

Seq(1, 2, 3).par.foreach { i =>
  val rc = rowCountOf { (1 to 100).toDF.write.csv(s"test${i}.csv") }
  println(rc)
}

=> 600
=> 700
=> 750

我可以让每个线程声明其自己的变量，但是spark上下文仍然共享，并且我无法确认特定的SparkListenerTaskEnd事件属于哪个线程。有什么办法可以使它工作？

（对，也许我可以将其单独做为火花作业。但这只是程序的一部分，所以为了简单起见，我宁愿留在线程中。在最坏的情况下，我只会执行它串行或忘记计数记录...）

Answer 1

有点黑，但是您可以将累加器用作过滤的副作用

val acc = spark.sparkContext.longAccumulator("write count")
df.filter { _ =>
  acc.add(1)
  true
}.write.csv(...)
println(s"rows written ${acc.count}")

使用线程计数Spark中的书面记录

1 个答案: