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

[dwsmith1983]

Which issue does this PR close?

Closes #5336.

Rationale for this change

CometMetricNode.reportScanInputMetrics feeds the Input column on the Stages and Executors tabs. It registered its completion listener after super.compute had created the CometExecIterator, and Spark runs completion listeners in reverse registration order, so the report ran before the iterator's close published the final native metrics. A fully consumed iterator closes itself on exhaustion, which hid the problem for plain scans and joins. A block with a JVM input only publishes on spark.comet.metrics.updateInterval, and a consumer that stops early never exhausts the iterator: a broadcast join with LIMIT 3 reported zero bytes and zero rows at the default interval. The listener also used setBytesRead/setRecordsRead, replacing bytes a fallback Spark scan had accumulated in the same task.

What changes are included in this PR?

• Register the report before super.compute at the three sites (CometNativeScanExec, CometIcebergNativeScanExec, the CometExecRDD override in operators.scala), matching what the native shuffle writer and native write already do.
• Add to the task's counters with incBytesRead/incRecordsRead, and claim each accumulator once per task through the per-task identity registry reportSpillMetrics already uses (generalised from SeenSpillMetrics to SeenMetrics, with a shared claimMetricValue helper that also treats an unset size metric as zero).
• Failed attempts now report too, since the listener is registered before anything that can throw.

What this does not cover

• The scan-input gate in buildNativeContext still matches only CometNativeScanExec, so an Iceberg (or CSV, or contrib) scan fused into a parent native block registers no report at all. That is the gap fix: report Input column when native Iceberg scan is enabled or native shuffle is enabled #5265 closes by widening the match to every CometLeafExec; this PR leaves that line untouched so the two do not collide.
• Several native plan instances in one task, as CometCoalesceExec produces, set absolute values into the same accumulators and the last one wins, in the SQL metrics themselves as well as the input column. That needs a native-side change and is tracked in Native metrics from several plan instances in one task overwrite each other, so a coalesced scan reports only its last partition #5879.
• When a coalesced task runs a fallback Spark scan partition before the native one, FileScanRDD registers its close listener first, so it runs last and sets bytesRead to the value it snapshotted at construction plus its own reads, dropping the native side's bytes. Records still add up because FileScanRDD increments those. TaskContext offers no ordering control, so the fallback test pins this in its reversed-arm half. A fallback scan reaching a native block through CometSparkToColumnarExec is not affected, since the block registers before it pulls its inputs.

How are these changes tested?

New tests in CometTaskMetricsSuite, asserting against vanilla Spark's numbers, against each side measured on its own, or against the plan shape:

• Broadcast join plus LIMIT 3 over a native scan, at update interval -1 and the default. Reported 0 bytes before this change.
• Native scan block plus LIMIT 3 at both intervals, and the same shape over an Iceberg table in CometIcebergNativeSuite.
• A coalesced UNION ALL of a native scan and a fallback JSON scan in one task, in both arm orders. Native first must cover the sum of the two sides measured alone; fallback first pins the ordering gap above. Reported 5000 of 10000 rows before this change.
• A coalesced union of a native scan and a cached input read from the block manager, at least the sum of the sides.
• A coalesced union of two native scans, each block reporting its own bytes, at least the sum of the sides. Reports 5000 of 10000 rows with set semantics.
• A failed attempt (ANSI division by zero after several batches) still reports the bytes and rows scanned.
• Unit tests on TaskContext.empty(): overlapping trees count each accumulator once and keep counters a Spark scan already set, the ordering contract, and scans that never ran report nothing.

Each guard was checked by reverting one part of the fix at a time: the ordering, the incrementing setters, and the once-per-task claim each make a distinct test fail, and with the records assertions muted the sum-based bytes assertions fail on their own. The full suite passes on Spark 3.5 and 4.0, along with the Iceberg native scan input metrics tests.

[andygrove]

I pulled this down and ran CometTaskMetricsSuite locally on 4.1, all 24 pass. I also reverted each half of the fix in turn to check the guards. Reverting the listener ordering fails native scan left unconsumed by a limit still reports task input metrics with bytesRead should be > 0 at interval -1, got 0, and reverting inc back to set fails four tests. So the change does what it says on the tin. Two things came out of that exercise, both left inline.

One point that falls outside the diff so I could not anchor it: reportNativeWriteOutputMetrics just below still uses setBytesWritten / setRecordsWritten and does not go through the registry. Is that deliberate because a task can only ever have one writer, or just out of scope here? A short note either way would save the next reader the trip.

[andygrove]

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]

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.

[andygrove]

When I reverted inc back to set to check what the new tests catch, all three integration failures were on recordsRead (5000 did not equal 10000), none on bytesRead. assertCometBytesReadInRange allows a 0.7 to 1.3 ratio on 3.5 and 4.0, and merely cometBytes >= sparkBytes on 4.1+. In this fixture the parquet side is small enough relative to the JSON side that losing it entirely still lands around 0.88, inside the band. So the bytes half of the change is not pinned by anything here, or in the coalesced-union and cached-input tests.

Could this measure the two sides separately first and then assert the coalesced run is at least their sum? That would make the bytes assertion fail on a set revert the same way the records one does.

[dwsmith1983]

Could this measure the two sides separately first and then assert the coalesced run is at least their sum?

Done for the fallback, coalesced-union and cached-input tests. Each side is measured on its own with Comet enabled and the coalesced run must be at least the sum. With inc reverted to set and the records assertions muted, the bytes assertions fail on their own: the fallback test reports comet=20033 against native=20033 plus fallback=146148, the cached test comet=45277 against parquet=45277 plus cached=82280, and the two-scan union drops to a 0.47 ratio.

[andygrove]

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]

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.

[andygrove]

Worth a line in the suite noting that these tests depend on markTaskCompleted running. Every TaskContext.empty() is constructed with taskAttemptId == 0, so they all share one seenMetricsByTask entry, and the cleanup listener is the only thing keeping them isolated from each other. A future test that registers a report and forgets markTaskCompleted would silently zero the next test's claims rather than failing.

[dwsmith1983]

Worth a line in the suite noting that these tests depend on markTaskCompleted running.

Added above the first unit test.

[dwsmith1983]

Is that deliberate because a task can only ever have one writer, or just out of scope here?

Deliberate. The native writer is the root of its task's plan and runs once per task, so nothing else in the task writes the output counters and set is exact. Added that to the scaladoc on reportNativeWriteOutputMetrics.