7

Is there any way to get the number of records written when using spark to save records? While I know it isn't in the spec currently, I'd like to be able to do something like:

val count = df.write.csv(path)

Alternatively, being able to do an inline count (preferably without just using a standard accumulator) of the results of a step would be (almost) as effective. i.e.:

dataset.countTo(count_var).filter({function}).countTo(filtered_count_var).collect()

Any ideas?

Jacek Laskowski
  • 72,696
  • 27
  • 242
  • 420
Loki
  • 6,205
  • 4
  • 24
  • 36

2 Answers2

6

I'd use SparkListener that can intercept onTaskEnd or onStageCompleted events that you could use to access task metrics.

Task metrics give you the accumulators Spark uses to display metrics in SQL tab (in Details for Query).

web UI / Details for Query

For example, the following query:

spark.
  read.
  option("header", true).
  csv("../datasets/people.csv").
  limit(10).
  write.
  csv("people")

gives exactly 10 output rows so Spark knows it (and you could too).

enter image description here


You could also explore Spark SQL's QueryExecutionListener:

The interface of query execution listener that can be used to analyze execution metrics.

You can register a QueryExecutionListener using ExecutionListenerManager that's available as spark.listenerManager.

scala> :type spark.listenerManager
org.apache.spark.sql.util.ExecutionListenerManager

scala> spark.listenerManager.
clear   clone   register   unregister

I think it's closer to the "bare metal", but haven't used that before.


@D3V (in the comments section) mentioned accessing the numOutputRows SQL metrics using QueryExecution of a structured query. Something worth considering.

scala> :type q
org.apache.spark.sql.DataFrame

scala> :type q.queryExecution.executedPlan.metrics
Map[String,org.apache.spark.sql.execution.metric.SQLMetric]

q.queryExecution.executedPlan.metrics("numOutputRows").value
Jacek Laskowski
  • 72,696
  • 27
  • 242
  • 420
  • 2
    Have you tested `SparkListener` with SQL? I've done some crude testing before, and it didn't capture SQL write metrics. `ExecutionListenerManager` is the way to go, and has nice example in the tests. – zero323 May 12 '17 at 12:51
  • 1
    No, I have not used either interface with SQL. Something I only know in theory...still. – Jacek Laskowski May 12 '17 at 12:59
  • 1
    If you're using SparkListenerTaskEnd.taskMetrics.outputMetrics.recordsWritten, keep in mind there was a [bug there solved only in Spark 2.3.0](https://issues.apache.org/jira/browse/SPARK-22605) – Eyal May 14 '19 at 10:56
  • Doesnt work. q.queryExecution.executedPlan.metrics is totally empty – beatrice Feb 14 '22 at 21:57
  • @beatrice I remember you have to execute the query first to access the metrics. They're accumulators so they filled in after tasks (of the `q` query) are done. – Jacek Laskowski Feb 17 '22 at 10:31
0

You could use an accumulator to count the rows as they are written out:

val count = df.sparkSession.sparkContext.longAccumulator("row count")
val counted = df.map(row => { count.add(1); row })(df.encoder)
counted.write.parquet("my file")
count.value

Since it has to decode/encode each row for us, I'm not sure this is faster than just checking the output:

df.sparkSession.read.parquet("my file").count

Parquet stores the row counts as metadata, so it might be fast enough to check.

Daniel Darabos
  • 26,991
  • 10
  • 102
  • 114