apache / apache/datafusion-comet
Task input metrics are unreliable when a native block mixes a native scan with a JVM input
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
## Describe the bug
`CometMetricNode.reportScanInputMetrics` (`spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala:69`) populates Spark's task-level `inputMetrics`, which drive the `Input` column on the UI's Stages and Executors tabs. It is unreliable when a native block contains a native scan *alongside* a JVM input, for example an Iceberg or Parquet scan joined against a shuffled side, or against a fallback Spark scan reaching the block through `CometSparkToColumnarExec`.
There are two independent problems. Both are long-standing. They were found while reviewing #5265, which does not introduce either one, but which does widen the set of plans that reach this reporting path from Parquet-only to any `CometLeafExec`.
## Problem 1: the completion listener is registered too late
Spark's `TaskContextImpl` keeps task completion listeners in a `Stack` and invokes them in reverse registration order:
> Using a stack causes us to process listeners in reverse order of registration. As listeners are invoked, they are popped from the stack.
`CometNativeExec.executeColumnarWithContext` (`spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:626-629`) calls `reportScanInputMetrics` *after* `super.compute(split, context)`. That is where the `CometExecIterator` is constructed and registers its own close listener. So the reporting listener runs *before* the iterator's `close()`, and `close()` is where `nativeLib.releasePlan` performs the final `update_metrics` (`native/core/src/execution/jni_api.rs:965`).
The same ordering appears at `CometNativeScanExec.scala:282` and `CometIcebergNativeScanExec.scala:239`.
Whether this actually loses data depends on which native execution path the block takes:
- **No JVM data sources.** `jni_api.rs:822` takes the `batch_receiver` branch, which calls `update_metrics` on every batch (`jni_api.rs:872`). The metric is current when the listener runs, so reporting is correct.
- **With a JVM data source.** Execution takes the busy-poll branch, where `update_metrics` only fires when `spark.comet.metrics.updateInterval` has elapsed (`jni_api.rs:903-905`). The default is 3000ms, so a task finishing inside that window reads a stale or zero `bytes_scanned`. `CometConf.COMET_METRICS_UPDATE_INTERVAL` documents that a negative interval means metrics are updated only on task completion, which would report zero every time.
The codebase already documents the correct ordering elsewhere. `CometNativeShuffleWriter.scala:130` registers before constructing its iterator, and `CometNativeWriteExec.scala:223` does the same with an explicit comment. The scaladoc on `CometMetricNode.reportNativeWriteOutputMetrics` states the requirement outright. The scan-input path is the one site that does not follow it.
## Problem 2: `setBytesRead` overwrites rather than accumulates
`reportScanInputMetrics` uses `setBytesRead` / `setRecordsRead` (`CometMetricNode.scala:81-82`) rather than the incrementing variants. If a fallback Spark scan reaches the same native block through `CometSparkToColumnarExec`, Spark's `FileScanRDD` has already accumulated bytes for that side into the task's `inputMetrics`. Comet's completion listener then discards it.
## How the two compound
Together they turn a missing number into a wrong one. For a block with an Iceberg scan and a fallback Spark scan, Spark's accumulated `bytesRead` is overwritten by a native value that Problem 1 can leave at zero, so the `Input` column goes from partially correct to empty.
## Suggested fix
Move the `reportScanInputMetrics` call above `super.compute(...)` at the three sites listed above, so reverse-order invocation puts it after the iterator's `close()`. This matches `CometNativeShuffleWriter` and `CometNativeWriteExec`. Separately, consider whether the incrementing metric setters are more appropriate than `setBytesRead` / `setRecordsRead`.
## Test coverage gap
Existing coverage exercises only the pure-native shape, where per-batch metric updates mask Problem 1. That includes the tests added in #5265 and `"native shuffle reports task input metrics for its scan child"` in `CometTaskMetricsSuite`. A test for a native block combining a native scan with a JVM input would cover both problems.
## Additional context
Found while reviewing #5265, which fixes a related but distinct bug where `hasScanInput` matched only `CometNativeScanExec` and so skipped Iceberg, CSV and contrib scans entirely.
Contributor guide
Research direction
Start with CometMetricNode.scala:69-82 and compare the listener registration in CometNativeExec, CometNativeScanExec, and CometIcebergNativeScanExec with CometNativeShuffleWriter.scala:130 and CometNativeWriteExec.scala:223. Read the related metric update paths in native/core/src/execution/jni_api.rs, then run CometTaskMetricsSuite and add coverage for a native scan combined with a JVM input; done means task input metrics retain both sources and include final native values.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust, scala
- Domain
- distributed-systems, observability, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100