Skip to content

fix: report task input metrics after the native iterator closes and add to Spark's counters #5880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dwsmith1983 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: report task input metrics after the native iterator closes and add to Spark's counters #5880

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,9 +235,10 @@ case class CometIcebergNativeScanExec(
nativeMetrics = nativeMetrics,
subqueries = Seq.empty) {
override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {
val res = super.compute(split, context)
// Register before super.compute creates the CometExecIterator, so this listener runs
// after the iterator's close has published the final scan metrics.

@andygrove andygrove Sep 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Iceberg site gets the same ordering fix, but I do not think any existing test can fail without it. "task-level inputMetrics.bytesRead is populated for Iceberg native scan" is a plain SELECT * with no JVM input and no early stop, which is the batch-receiver shape where update_metrics fires per batch, so it passes either way. An Iceberg scan with a LIMIT at the default update interval would mirror your parquet test and give this line a guard.

@dwsmith1983 dwsmith1983 Sep 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An Iceberg scan with a LIMIT at the default update interval would mirror your parquet test and give this line a guard.

Added in CometIcebergNativeSuite at the default interval and at -1. One caveat: the Iceberg scan is always its own block with no JVM input, so it takes the batch-receiver path where every returned batch publishes metrics, and the test passes with or without the reorder. The shape that would fail, an Iceberg scan fused under a join with a broadcast input, registers no report at all until #5265 widens the gate. So this test covers the site rather than proving the order.

Option(context).foreach(nativeMetrics.reportScanInputMetrics)
res
super.compute(split, context)
}
}
}
Expand Down
85 changes: 52 additions & 33 deletions spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
seenMetrics: IdentityHashMap[SQLMetric, java.lang.Boolean]): Long = {
def sumFromNode(metricNode: CometMetricNode): Long = {
val nodeValue = metricNode.metrics.get(metricName).fold(0L) { metric =>
if (seenMetrics.put(metric, java.lang.Boolean.TRUE) == null) {
math.max(metric.value, 0L)
} else {
0L
}
CometMetricNode.claimMetricValue(metric, seenMetrics)
}
nodeValue + metricNode.children.iterator.map(sumFromNode).sum
}
Expand All @@ -105,24 +101,34 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
})

/**
* Reports aggregated scan input metrics (bytesRead, recordsRead) to Spark's task metrics.
* Aggregates across all scan leaf nodes to handle plans with multiple scans (e.g., joins). Must
* be called in a TaskCompletionListener after the iterator is fully consumed.
* Reports the scan leaves' bytes and rows (summed across joins and unions) to Spark's task
* input metrics, which drive the Input column on the UI's Stages and Executors tabs.
*
* Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so its
* completion listener publishes final SQL metrics before this listener runs. A block with a JVM
* input only publishes on the metrics update interval, and a consumer that stops early, such as
* a limit, leaves the final publish to that close.
*
* Adds to the task's counters instead of replacing them, so bytes that a fallback Spark scan

@andygrove andygrove Sep 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the reversed-arm test and the description bullet cover this now, but the scaladoc here still has the unconditional version, "so bytes that a fallback Spark scan accumulated in the same task survive". That is the claim the reversed arm disproves, and it is the sentence someone reading this file will find rather than the PR description.

Could you carry a short version of the caveat up here? Something like: this survives only when the fallback scan registers its listener after Comet's, which a CometSparkToColumnarExec input always does but a coalesced Spark-scan partition computed first does not.

* accumulated in the same task survive. Trees registered on one task may share accumulators
* (see [[reportSpillMetrics]]), so each accumulator is counted once per task.
*/
def reportScanInputMetrics(ctx: TaskContext): Unit = {
val seenMetrics = CometMetricNode.taskSeenMetrics(ctx).scanInput
ctx.addTaskCompletionListener[Unit] { _ =>
val scanLeaves = leafNodes.filter(_.metrics.contains("bytes_scanned"))
if (scanLeaves.nonEmpty) {
val totalBytes = scanLeaves.map(_.metrics("bytes_scanned").value).sum
val totalRows = scanLeaves.map { leaf =>
val outputRows =
leaf.metrics.get("output_rows").map(_.value).getOrElse(0L)
val prunedRows =
leaf.metrics.get("pushdown_rows_pruned").map(_.value).getOrElse(0L)
outputRows + prunedRows
}.sum
ctx.taskMetrics().inputMetrics.setBytesRead(totalBytes)
ctx.taskMetrics().inputMetrics.setRecordsRead(totalRows)
def claimed(leaf: CometMetricNode, metricName: String): Long =
leaf.metrics.get(metricName).fold(0L)(CometMetricNode.claimMetricValue(_, seenMetrics))

val totalBytes = scanLeaves.map(claimed(_, "bytes_scanned")).sum
val totalRows = scanLeaves.map { leaf =>
claimed(leaf, "output_rows") + claimed(leaf, "pushdown_rows_pruned")
}.sum
if (totalBytes > 0L) {
ctx.taskMetrics().inputMetrics.incBytesRead(totalBytes)

@andygrove andygrove Sep 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do not think the inc change holds in the other arm order. FileScanRDD snapshots existingBytesRead when its iterator is constructed in compute() and its completion-time close() does setBytesRead(existingBytesRead + getBytesReadCallback()), so it replaces rather than adds. Comet's inc only survives when Spark's close listener runs first, which requires Spark's scan to have registered its listener after Comet's.

In "task input metrics keep bytes read by a fallback Spark scan in the same task" the parquet arm comes first, so Comet registers first and Spark's close runs first. If I swap the arms so the JSON scan is the leading one, reading exactly the same two files, Comet's bytesRead drops by the parquet side's entire bytes_scanned:

native-first:   sparkBytes=147770  cometBytes=166181
fallback-first: sparkBytes=147770  cometBytes=146148

recordsRead is correct in both orders, since FileScanRDD uses incRecordsRead there. The scope is narrow: a fallback scan reaching a native block through CometSparkToColumnarExec is always safe, because CometExecRDD.compute registers this listener before resolveInputObjects pulls the input RDDs. It takes a coalesce that merges a Spark-scan partition and a Comet-scan partition into one task with the Spark side first, which is the shape the new test uses.

This is not a regression, both orders are wrong on main today, and I do not see a clean fix given TaskContext has no listener ordering control. Would you add the reversed-arm query as a test that pins the current behaviour, and mention it under "What this does not cover" next to #5265 and #5879? As written the description reads as though the fallback bytes always survive.

@dwsmith1983 dwsmith1983 Sep 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you add the reversed-arm query as a test that pins the current behaviour, and mention it under "What this does not cover" next to #5265 and #5879?

Added both. The test now runs the union in both arm orders. Native first asserts bytesRead covers the sum of the two sides measured on their own. Fallback first asserts records still add up and bytesRead stays below that sum, with a comment on FileScanRDD's close setting the value it snapshotted at construction. The description has a bullet for it next to the other two.

}
if (totalRows > 0L) {
ctx.taskMetrics().inputMetrics.incRecordsRead(totalRows)
}
}
}
Expand All @@ -135,6 +141,10 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
* Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so
* Spark's completion listener stack invokes the iterator `close` (final SQL metric update)
* before this listener runs.
*
* The native writer is the root of its task's plan and runs once per task, so its values
* replace the task's output counters without the per-task registry that the scan input and
* spill reports share.
*/
def reportNativeWriteOutputMetrics(ctx: TaskContext): Unit = {
ctx.addTaskCompletionListener[Unit] { _ =>
Expand All @@ -161,7 +171,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
* per-task registry, so each accumulator is counted once while disjoint trees still all report.
*/
def reportSpillMetrics(ctx: TaskContext): Unit = {
val seenMetrics = CometMetricNode.taskSeenSpillMetrics(ctx)
val seenMetrics = CometMetricNode.taskSeenMetrics(ctx)
ctx.addTaskCompletionListener[Unit] { _ =>
val diskBytesSpilled = sumMetricValues("spilled_bytes", seenMetrics.disk)
if (diskBytesSpilled > 0L) {
Expand Down Expand Up @@ -225,30 +235,39 @@ object CometMetricNode {
private val aggregateMetricNames =
Set("spill_count", "spilled_bytes", "spilled_rows", "peak_mem_used")

private case class SeenSpillMetrics(
disk: IdentityHashMap[SQLMetric, java.lang.Boolean],
memory: IdentityHashMap[SQLMetric, java.lang.Boolean])
private type SeenMetricSet = IdentityHashMap[SQLMetric, java.lang.Boolean]

private case class SeenMetrics(
disk: SeenMetricSet,
memory: SeenMetricSet,
scanInput: SeenMetricSet)

// Per running task attempt: the spill accumulators already claimed by a reporting listener,
// one identity set per metric name (see reportSpillMetrics). The first registration installs
// a cleanup listener ahead of every reporting listener, so it runs last (reverse registration
// order) and removes the entry.
private val seenSpillMetricsByTask = new ConcurrentHashMap[Long, SeenSpillMetrics]()
// Per running task attempt: the accumulators already claimed by a reporting listener, one
// identity set per reported task metric (see reportSpillMetrics and reportScanInputMetrics).
// The first registration installs a cleanup listener ahead of every reporting listener, so it
// runs last (reverse registration order) and removes the entry.
private val seenMetricsByTask = new ConcurrentHashMap[Long, SeenMetrics]()

private def taskSeenSpillMetrics(ctx: TaskContext): SeenSpillMetrics = {
private def taskSeenMetrics(ctx: TaskContext): SeenMetrics = {
val attemptId = ctx.taskAttemptId()
val existing = seenSpillMetricsByTask.get(attemptId)
val existing = seenMetricsByTask.get(attemptId)
if (existing != null) {
existing
} else {
// The task thread is the only registrant for its attempt id, so there is no put race.
val created = SeenSpillMetrics(new IdentityHashMap(), new IdentityHashMap())
seenSpillMetricsByTask.put(attemptId, created)
ctx.addTaskCompletionListener[Unit](_ => seenSpillMetricsByTask.remove(attemptId))
val created =
SeenMetrics(new IdentityHashMap(), new IdentityHashMap(), new IdentityHashMap())
seenMetricsByTask.put(attemptId, created)
ctx.addTaskCompletionListener[Unit](_ => seenMetricsByTask.remove(attemptId))
created
}
}

/** The metric's value the first time it is claimed for a task, zero afterwards. */
private def claimMetricValue(metric: SQLMetric, seenMetrics: SeenMetricSet): Long =
if (seenMetrics.put(metric, java.lang.Boolean.TRUE) == null) math.max(metric.value, 0L)
else 0L

/**
* The baseline SQL metrics for DataFusion `BaselineMetrics`.
*/
Expand Down
8 changes: 3 additions & 5 deletions spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,12 +277,10 @@ case class CometNativeScanExec(
encryptedFilePaths,
perPartitionFilePaths = perPartitionFilePaths) {
override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {
val res = super.compute(split, context)

// Report scan input metrics after the iterator is fully consumed.
// Register before super.compute creates the CometExecIterator, so this listener runs
// after the iterator's close has published the final scan metrics.
Option(context).foreach(nativeMetrics.reportScanInputMetrics)

res
super.compute(split, context)
}
}
}
Expand Down
5 changes: 3 additions & 2 deletions spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -626,11 +626,12 @@ abstract class CometNativeExec extends CometExec {
ctx.encryptedFilePaths,
ctx.shuffleScanIndices) {
override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {
val res = super.compute(split, context)
// Register before super.compute creates the CometExecIterator, so this listener runs
// after the iterator's close has published the final scan metrics.
if (ctx.hasScanInput) {
Option(context).foreach(nativeMetrics.reportScanInputMetrics)
}
res
super.compute(split, context)
}
}
}
Expand Down
75 changes: 75 additions & 0 deletions spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3738,6 +3738,81 @@ class CometIcebergNativeSuite
}
}

test("Iceberg native scan left unconsumed by a limit still reports task input metrics") {
assume(icebergAvailable, "Iceberg not available in classpath")

withTempIcebergDir { warehouseDir =>
withSQLConf(
"spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog",
"spark.sql.catalog.test_cat.type" -> "hadoop",
"spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath,
CometConf.COMET_ENABLED.key -> "true",
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {

spark.sql("""
CREATE TABLE test_cat.db.task_metrics_limit_test (
id INT,
value DOUBLE
) USING iceberg
""")
spark
.range(20000)
.selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value")
.repartition(4)
.write
.format("iceberg")
.mode("append")
.saveAsTable("test_cat.db.task_metrics_limit_test")

val bytesReadValues = mutable.ArrayBuffer.empty[Long]
val recordsReadValues = mutable.ArrayBuffer.empty[Long]
val listener = new SparkListener {
override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
val im = taskEnd.taskMetrics.inputMetrics
bytesReadValues.synchronized {
bytesReadValues += im.bytesRead
recordsReadValues += im.recordsRead
}
}
}
spark.sparkContext.addSparkListener(listener)

try {
// The limit stops pulling before the scan is exhausted, so the final metric publish

@andygrove andygrove Sep 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I ran the ordering revert against both suites to see what this test buys. It passes with the ordering reverted, as you said it would, and so does native scan block left unconsumed by a limit .... Only the broadcast-join test fails. The reason is in jni_api.rs: a block with no JVM input is spawned onto the batch_receiver channel and calls update_metrics on every batch, while a block that has one takes the busy-poll stream path where update_metrics only fires on the update interval. So having a JVM input is what decides whether the ordering matters, not stopping early.

Both tests earn their keep as coverage for their sites. Could this comment say what the test actually does though? Right now it reads "the final metric publish happens in the iterator's completion-time close", which is the broadcast-join test's mechanism rather than this one. Something like: this Iceberg scan is its own block with no JVM input, so metrics publish per batch and this covers the site rather than the order, which stays unguarded until #5265.

// happens in the iterator's completion-time close.
val query = "SELECT * FROM test_cat.db.task_metrics_limit_test LIMIT 3"
Seq("-1", CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach {
interval =>
withSQLConf(CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval) {
CometListenerBusUtils.waitUntilEmpty(spark.sparkContext)
bytesReadValues.clear()
recordsReadValues.clear()
val df = spark.sql(query)
assert(
collectIcebergNativeScans(df.queryExecution.executedPlan).nonEmpty,
"Expected CometIcebergNativeScanExec in plan")
df.collect()
CometListenerBusUtils.waitUntilEmpty(spark.sparkContext)

val cometBytes = bytesReadValues.sum
val cometRecords = recordsReadValues.sum
assert(
cometBytes > 0,
s"bytesRead should be > 0 at interval $interval, got $cometBytes")
assert(
cometRecords >= 3 && cometRecords <= 20000,
s"recordsRead should cover at least the limit at interval $interval, got $cometRecords")
}
}
} finally {
spark.sparkContext.removeSparkListener(listener)
spark.sql("DROP TABLE test_cat.db.task_metrics_limit_test")
}
}
}
}

test("task-level inputMetrics.bytesRead is populated for Iceberg native scan") {
assume(icebergAvailable, "Iceberg not available in classpath")

Expand Down
Loading
Loading