[SPARK-58206][CORE][SQL] Add custom write metrics to the pluggable shuffle framework

[karuppayya]

What changes were proposed in this pull request?

Adds the ability for a pluggable shuffle storage implementation to declare custom write metrics and report their values. The declared metrics are surfaced on the Exchange operator node in the SQL UI.
(This is similar to the custom-metrics support that DSV2 sources already have.)

This covers shuffle write metrics only. Read-side custom metrics will be a follow-up.

Why are the changes needed?

Spark supports pluggable shuffle storage, so shuffle data can be written to substrates like a remote/external shuffle service, a distributed filesystem (HDFS), or the default local disk.
However, plugins have no way to expose substrate-specific write metrics to users — e.g. bytes pushed to a remote shuffle service, or spill-file counts and fsync counts for local disk — so these remain invisible in the SQL UI.

Does this PR introduce any user-facing change?

Yes, indirectly.
A shuffle storage plugin can now declare custom write metrics, which appear on the Shuffle exchange operator node in the SQL UI.
Plugins that dont declare custome metrics see no change.

How was this patch tested?

• Unit tests
• Manual testing from SQL UI

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

[karuppayya]

@cloud-fan @viirya @sunchao @dongjoon-hyun can you please help review? Thanks!

[sunchao]

Summary

Prior state and problem

Spark's pluggable shuffle-storage API lets implementations redirect shuffle output to local disks, remote services, or distributed filesystems, but it has no standard way to surface implementation-specific write metrics on SQL Exchange nodes. This makes important substrate behavior such as uploaded bytes, remote requests, spill-file counts, and fsync activity invisible to users.

Design approach

This PR adds driver-side metric declarations through ShuffleDriverComponents.supportedCustomMetrics(), per-task values through ShuffleMapOutputWriter.currentMetricsValues(), and storage in the built-in sort shuffle writers. ShuffleWriteProcessor then bridges the task values into SQLMetric accumulators exposed by ShuffleExchangeExec.

Correctness / compatibility analysis

The new Java API methods have defaults, so existing shuffle plugins remain source and binary compatible. The supported sum, size, millisecond-timing, and nanosecond-timing types also reuse established SQL metric behavior. Two correctness gaps remain, however: custom names can currently replace Spark-owned Exchange metrics, including the row-count statistic used by AQE, and the common Unsafe single-spill fast path cannot return any of the new custom values.

Key design decisions

Metric declarations live on the driver while task values are matched by name on executors, avoiding serialization of plugin-owned declaration objects. The feature intentionally covers shuffle writes only, aggregates per-task values by sum, rejects unsupported metric types, and ignores task values that were not declared by the driver.

Implementation sketch

Each Exchange creates one SQLMetric per declared custom metric and includes those accumulators in its shuffle writer processor. After a built-in shuffle writer commits its map output, it caches the values reported by the underlying map-output writer; after successful writer shutdown, the processor transfers those values into the task's SQL accumulators.

Behavioral changes worth calling out

Plugins that declare no metrics retain the previous behavior. Plugins with unique names work on the Bypass, Sort, empty-Unsafe, and standard Unsafe merge paths. As written, a colliding name can corrupt a built-in metric rather than merely hiding the custom value, while an optimized one-spill Unsafe write silently contributes no custom metrics.

Suggested improvements

Reserve or reject every Spark-owned Exchange metric name (or otherwise guarantee that Spark-owned metrics take precedence), and provide a metric handoff for SingleSpillShuffleMapOutputWriter. Regression coverage should include a shuffleRecordsWritten collision and a non-empty Unsafe write whose executor components return the optional single-file writer.

[sunchao]

[P1] Keep custom metrics from replacing Spark-owned entries

Because this map is appended last, a plugin can declare a metric such as shuffleRecordsWritten and silently replace the Spark accumulator exposed through metrics. The real reporter still updates the separate writeMetrics instance, while runtimeStatistics reads metrics("shuffleRecordsWritten"); a custom value (or default) of zero therefore makes a materialized non-empty shuffle report rowCount = 0. AQEPropagateEmptyRelation treats that as definitively empty and can rewrite a non-empty branch to an empty relation, changing query results. Please reject/reserve collisions, or make all Spark-owned keys take precedence as SupportsCustomDriverMetrics does, and add a collision regression test.

[sunchao]

[P2] Capture metrics from the single-spill fast path

This handles the zero-spill case, but the following one-spill branch uses SingleSpillShuffleMapOutputWriter and never assigns customMetricsValues. A non-empty serialized shuffle that fits in memory still produces one final spill, so whenever a plugin exposes the optional single-file writer this common path completes successfully with the custom array left empty. The added Unsafe test uses an empty iterator and cannot catch this. Please expose values through the single-spill API (or a common task-level hook), capture them after transfer, and cover the non-empty one-spill path.

[karuppayya]

Thanks @sunchao for the review

Keep custom metrics from replacing Spark-owned entries

i have added the filtering logic similar to custom read/write metrics in DSV2.

Capture metrics from the single-spill fast path

Handled this case which was initially missed.