fix: report native shuffle write metrics accurately

[sunchao]

Why are the changes needed?

When a shuffle becomes too large to keep in memory, Spark spills intermediate data to local disk. Operators use the Spark UI to answer two different questions about that work: how much memory was occupied by the data that spilled, and how much data was actually written to disk. Spark intentionally exposes these as separate task metrics because the on-disk representation may be compressed and can be much smaller than the in-memory representation.

Comet's native shuffle writer did not preserve that distinction. It copied the same on-disk byte count into both memoryBytesSpilled and diskBytesSpilled. Worse, its disk counter considered only bytes returned by immediate writes, even though Arrow's batch coalescer can hold those bytes until flush(). In that common case, a real spill could appear as zero memory spilled and zero disk spilled in the Spark UI.

For example, suppose a task spills 256 MiB of in-memory shuffle data and partition indices, which compress to 32 MiB on disk. These numbers are illustrative:

Situation	What actually happened	Spark UI before	Spark UI after
Compressed bytes are emitted during the write	256 MiB in memory, 32 MiB on disk	Memory: 32 MiB; disk: 32 MiB	Memory: 256 MiB; disk: 32 MiB
Compressed bytes are emitted only when the writer flushes	256 MiB in memory, 32 MiB on disk	Memory: 0; disk: 0	Memory: 256 MiB; disk: 32 MiB
The task spills successfully, then fails before the shuffle commit	256 MiB in memory, 32 MiB on disk	Failed-task memory: 0; disk: 0	Failed-task memory: 256 MiB; disk: 32 MiB

The failure case matters especially for debugging: a task that runs out of disk, encounters bad input, or fails after several spills should not look as though it never spilled at all.

Accurate in-memory accounting also has an Arrow-specific complication. One input batch can be split into many zero-copy slices that all reference the same underlying allocation. The regression test uses 16,384 Int64 values, backed by one 128 KiB allocation, and processes them as 16 slices of 1,024 rows. If every slice triggers a spill, simply summing each released memory reservation counts that same 128 KiB allocation 16 times: 2 MiB instead of 128 KiB, before even including the real per-slice partition-index allocations. The underlying allocation remains owned by the original input batch throughout the operation, so it must be counted once for that batch, not once per slice.

Finally, native shuffle already measures the time spent interleaving rows into output partitions, but that work was not visible in the SQL exchange metrics. This made it harder to distinguish repartitioning overhead from encoding and compression overhead when investigating slow shuffle stages.

Part of #3996. This PR addresses native shuffle write observability only; shuffle reads and mixed native/JVM scan-input accounting remain outside its scope. Follow-up #5382 tracks spills from native child operators inlined beneath the shuffle writer.

What changes were proposed in this PR?

The change gives native shuffle spill accounting two independent sources of truth and carries them consistently into both of Spark's observability surfaces.

For disk usage, the source of truth is the physical spill file itself. The native writer observes the file position before a spill and again after buffered Arrow data has been flushed. Their difference is the actual number of bytes written to disk, including data produced only during the final flush and data appended to an existing spill file. When shuffle compression is enabled, this is the compressed size; when spark.shuffle.compress=false, it is the uncompressed on-disk size.

For memory usage, the source of truth is the in-memory shuffle state being spilled: Arrow backing allocations plus the partition-index structures that organize its rows. The native writer records the memory released by the spill and also includes the current batch when the memory pool rejected its reservation after that batch had already been buffered. To avoid turning Arrow's zero-copy slicing into fictional memory growth, it remembers which backing allocations have already been counted for the lifetime of the original input batch. Separate input batches still contribute cumulatively, as Spark's spill metrics require.

Those two measurements are exposed separately in the SQL shuffle exchange and copied into their corresponding Spark task metrics only after the native execution plan has published its final values. The copy runs from a task-completion listener registered before the native iterator; Spark invokes completion listeners in reverse registration order, so native cleanup and final metric publication happen first. Because task-completion listeners also run when an attempt fails or is canceled, the same accounting remains available for the attempts that are most useful to diagnose.

The exchange also surfaces the native partition-interleaving timer, making the main stages of native shuffle work visible alongside the existing repartitioning, encoding, and spill metrics. Existing shuffle output semantics, compression behavior, and scan-input accounting are unchanged.

How was this PR tested?

Focused Rust tests exercise both spill triggers: an explicit maximum buffer size and a memory-pool reservation failure. They compare reported disk spill bytes against the actual spill-file sizes, verify zero spill metrics when no spill occurs, and cover the shared-allocation example above. They also verify that separate outer input batches are counted cumulatively even when they reference the same Arrow backing buffer.

The Spark integration tests validate the complete native-to-Spark reporting path on both Spark 3.5 and Spark 4.0:

• A successful shuffle writes 20,000 rows across four partitions with Zstandard compression and forced spills. The test verifies shuffle records, output bytes, write time, partition-interleaving time, separate memory/disk SQL metrics, and exact agreement between those SQL metrics and Spark's recorded stage/task metrics.
• A failing shuffle processes 8,192 valid rows before a later row triggers an ANSI divide-by-zero inside a native projection. Earlier batches have already spilled, and the test verifies that the failed ShuffleMapTask still reports both memory and disk spill bytes.

cd native
cargo fmt --all -- --check
cargo clippy -p datafusion-comet-shuffle --lib -- -D warnings
cargo test -p datafusion-comet-shuffle --lib spill -- --nocapture
cargo test -p datafusion-comet-shuffle --lib shared_backing_once_per_input_batch -- --nocapture

cd ..
make core
./mvnw -Pspark-3.5 test -Dtest=none \
  '-Dsuites=org.apache.spark.sql.comet.CometTaskMetricsSuite memory and disk spill metrics'
./mvnw -Pspark-4.0 test -Dtest=none \
  '-Dsuites=org.apache.spark.sql.comet.CometTaskMetricsSuite memory and disk spill metrics'

[andygrove]

Thanks for this. The disk-side fix is a genuine bug with a clear root cause, and I like that writer_stream_position measures the file rather than re-deriving it, since BufBatchWriter::flush guarantees everything has reached the file before the second reading.

I traced the shared-buffer memory bookkeeping through several scenarios and it holds up:

• Clearing spill_accounted_input_buffers per insert_batch is what protects against allocator address reuse across outer batches. Since the outer batch is held alive for the whole slicing loop, an address cannot be freed and reused within one input batch, so the set cannot produce a false "already counted" hit. That is a subtle hazard the code gets right.
• repeated_spill_buffer_bytes deliberately persisting across insert_batch boundaries is correct. The bytes it refers to are still pinned by leftover slices, so they must stay subtracted until the next spill drains them.
• repeated is a subset of new_buffer_bytes, which is a subset of mem_growth, and mem_growth always lands in the pot either via a successful try_grow or via unreserved_bytes, so the subtraction never over-subtracts.
• In the common shape where input batches arrive at exactly batch_size, repeated_spill_buffer_bytes stays zero throughout and the whole mechanism reduces to reservation.free() + unreserved_bytes. The complexity only engages for oversized inputs, which seems like the right tradeoff.

Marking only buffered_batches.last() rather than every buffered batch also looks fine to me, since RecordBatch::slice is offset-based across every Arrow array type and all slices of one outer batch expose identical buffer addresses.

The Rust tests are well targeted. Comparing one input batch against two in max_buffer_spills_charge_shared_backing_once_per_input_batch is exactly the assertion that catches per-slice double counting, and checking spilled_bytes against actual fs::metadata file lengths pins the disk metric to ground truth rather than to itself.

A few things I would like to see addressed.

1. The task metric bridge belongs on CometMetricNode

CometMetricNode already has two methods that do precisely this job: reportScanInputMetrics(ctx) and reportNativeWriteOutputMetrics(ctx). The second one even carries the doc comment explaining the listener ordering requirement that the new inline block in CometNativeShuffleWriter.write() re-explains. write() already calls nativeMetrics.reportScanInputMetrics(...) a few lines above.

Would a reportSpillMetrics(ctx) next to those two work here? That keeps all three task metric bridges in one place under one documented ordering contract. A future change to how Comet publishes final native metrics would then have one place to fix rather than three.

For what it is worth I checked the ordering claim against Spark master. TaskContextImpl keeps completion listeners in a Stack and pops them, so registration-before does mean invocation-after. The reasoning in the comment is right.

2. The disk metric label is now ambiguous

The whole point of the change is separating the two figures, but the UI ends up showing spilled bytes next to memory spilled bytes. Someone looking at those two rows has no way to tell that the first is the compressed on-disk number. Could the existing label become something like disk spilled bytes? The metric key stays spilled_bytes, so nothing on the native side changes.

3. docs/source/user-guide/latest/metrics.md needs updating

That page is hand-edited, and its Exchange section currently lists only native shuffle time, repartition time, memory pool time, and encoding and compression time. This PR adds two user-visible metrics and, more importantly, changes what Spark's memoryBytesSpilled means for native shuffle. Documenting partition interleaving time and the memory versus disk spill distinction there would help people comparing Comet's numbers against vanilla Spark.

4. Spills from operators inlined under the shuffle writer still do not reach task metrics

shuffleWriterSQLMetrics filters detailedMetrics out of the writer node's own map. When spec.childNativeOp is a rich native subtree, a CometSortExec or CometSortMergeJoinExec inside it reports its own spill_count and spilled_bytes into spec.childMetricNode, and those never make it into taskMetrics. So the Stages tab still under-reports spill for those plans.

That is not a regression, it was equally true before. But it is directly adjacent to what this PR fixes, and the listener you just added is the natural place to handle it. Since the description scopes this PR to the writer, could you file a tracking issue and link it here? Otherwise it will not get picked up.

5. COMET_SHUFFLE_JVM_BATCH_SIZE looks like a no-op in the new test

In CometTaskMetricsSuite, the failed-attempt test sets COMET_SHUFFLE_JVM_BATCH_SIZE, but that config is only read by CometDiskBlockWriter and the JVM SpillWriter, both on the columnar shuffle path. The test runs with COMET_SHUFFLE_MODE=native, so it should have no effect, and COMET_BATCH_SIZE is doing the real work. Was it left over from an earlier iteration? Worth dropping so a future reader does not assume it matters for reproducing the failure at row 8192.

[sunchao]

Thanks for the detailed review, @andygrove — addressed all five points:

1. Moved the task-metric bridge into CometMetricNode.reportSpillMetrics, preserving completion-listener ordering for both successful and failed attempts.
2. Renamed the SQL UI label to disk spilled bytes while keeping the native spilled_bytes metric key unchanged.
3. Updated the metrics guide to cover partition interleaving, disk-versus-memory spill metrics, Spark task-metric semantics, and compression-disabled shuffles.
4. Opened Report native child-operator spill metrics in Spark task metrics for unified shuffle plans #5382 to track spill accounting for native operators inlined beneath the shuffle writer.
5. Removed the unused COMET_SHUFFLE_JVM_BATCH_SIZE setting from the failed-attempt regression.

The focused successful- and failed-attempt spill regressions pass against both Spark 3.5 and Spark 4.0.