[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.

[sunchao]

LGTM

[uros-b]

The default empty customWriteMetrics means the three callers in limit.scala (CollectLimitExec, CollectTailExec, and TakeOrderedAndProjectExec) silently discard metrics returned by the shuffle plugin.

These operators already expose built-in shuffle metrics, so they should also create collision-filtered custom metrics, include them in metrics, and pass them to prepareShuffleDependency.

Could you also add a regression test covering a multi-partition limit shuffle?

[karuppayya]

Good catch, thanks @uros-b — fixed.

• All three operators now create collision-filtered custom metrics, expose them in metrics, and pass them to prepareShuffleDependency.
• Extracted a shared ShuffleMetricsSupport trait + CustomShuffleMetrics.createFilteredMetrics helper
  to avoid repeating the plumbing.
• Removed the Map.empty default on prepareShuffleDependency / createShuffleWriteProcessor so a
  future caller that forgets to wire metrics gets a compile error .
• Added a regression test for a multi-partition limit shuffle.

[dongjoon-hyun]

Thank you, @karuppayya .

However, we need to be more careful for this part.

Although ShuffleExchangeExec is an internal class of Apache Spark, this PR is a breaking change for Spark-accelerator projects like Apache Gluten and Apache DataFusion Comet.

For example, Apache DataFusion Comet has a dependency for the following.

• prepareShuffleDependency

$ git grep prepareShuffleDependency
docs/source/contributor-guide/native_shuffle.md:     convenience `prepareShuffleDependency(rdd, ...)` overload (used by
spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala:        val dep = CometShuffleExchangeExec.prepareShuffleDependency(
spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala:        val dep = CometShuffleExchangeExec.prepareShuffleDependency(
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala: * [[CometShuffleExchangeExec.prepareShuffleDependency]] convenience overload). Same handling
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:   * [[CometShuffleExchangeExec.prepareShuffleDependency]] convenience overload (synthetic Scan
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:          CometShuffleExchangeExec.prepareShuffleDependency(
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:  def prepareShuffleDependency(
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:    // ShuffleExchangeExec::prepareShuffleDependency
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:  // the prepareShuffleDependency convenience-overload path, which builds its own NativeExecContext.

• createShuffleWriteProcessor

$ git grep createShuffleWriteProcessor
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:      shuffleWriterProcessor = ShuffleExchangeExec.createShuffleWriteProcessor(metrics),
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:        shuffleWriterProcessor = ShuffleExchangeExec.createShuffleWriteProcessor(writeMetrics),

To avoid unnecessary overhead for downstream projects, please revise this PR into a backward compatible one by preserving the legacy APIs.

cc @andygrove, @kazuyukitanimura , @peter-toth

[sunchao]

One compatibility blocker remains on this head.

[karuppayya]

Kept the original signatures as overloads that delegate to the metrics-aware versions, so external callers are unaffected.

• minor test refactoring.

@sunchao @dongjoon-hyun would appreciate another review round.

[dongjoon-hyun]

This looks like an inconsistent error handling between name collisions and invalid metric types.

createFilteredMetrics handles the two failure modes differently: a plugin metric whose
name collides with a Spark-owned metric is silently dropped with a warning, but an invalid
metricType() throws an IllegalArgumentException.

Since the exception is thrown while evaluating the operator's lazy metrics val (i.e. during
planning / UI access), a single misconfigured plugin would fail every query that contains a
shuffle, rather than just losing the bad metric.

Could we make the two cases consistent? Two options:

1. Treat an invalid metric type the same as a collision: log a warning and drop the metric,
   so a buggy plugin degrades gracefully instead of breaking all queries.
2. If we do want to fail fast on an invalid type, use an error-class-based exception
   (e.g. SparkIllegalArgumentException) instead of a plain IllegalArgumentException,
   per the current error framework convention.

I'd lean toward option 1, since custom metrics are purely observability and shouldn't be able
to take down query execution.