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

core/src/main/java/org/apache/spark/shuffle/api/ShuffleDriverComponents.java

@@ -20,6 +20,7 @@
 import java.util.Map;
 
 import org.apache.spark.annotation.Private;
+import org.apache.spark.shuffle.api.metric.CustomShuffleMetric;
 
 /**
  * :: Private ::
@@ -70,4 +71,12 @@ default void removeShuffle(int shuffleId, boolean blocking) {}
   default boolean supportsReliableStorage() {
     return false;
   }
+
+  /**
+   * The custom shuffle-write metrics this shuffle component supports. Per-task values are reported
+   * by {@link ShuffleMapOutputWriter#currentMetricsValues()} and matched to these by name.
+   */
+  default CustomShuffleMetric[] supportedCustomMetrics() {
+    return new CustomShuffleMetric[0];
+  }
 }

core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java

@@ -21,6 +21,7 @@
 
 import org.apache.spark.annotation.Private;
 import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage;
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric;
 
 /**
  * :: Private ::
@@ -85,4 +86,11 @@ public interface ShuffleMapOutputWriter {
    * clean up temporary state if necessary.
    */
   void abort(Throwable error) throws IOException;
+
+  /**
+   * The values of the custom shuffle metrics for this map task.
+   */
+  default CustomShuffleTaskMetric[] currentMetricsValues() {
+    return new CustomShuffleTaskMetric[0];
+  }
 }

core/src/main/java/org/apache/spark/shuffle/api/SingleSpillShuffleMapOutputWriter.java

@@ -21,6 +21,7 @@
 import java.io.IOException;
 
 import org.apache.spark.annotation.Private;
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric;
 
 /**
  * Optional extension for partition writing that is optimized for transferring a single
@@ -36,4 +37,11 @@ void transferMapSpillFile(
       File mapOutputFile,
       long[] partitionLengths,
       long[] checksums) throws IOException;
+
+  /**
+   * The values of the custom shuffle metrics for this map task.
+   */
+  default CustomShuffleTaskMetric[] currentMetricsValues() {
+    return new CustomShuffleTaskMetric[0];
+  }
 }

core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.api.metric;
+
+import org.apache.spark.annotation.Private;
+
+/**
+ * :: Private ::
+ * A custom shuffle metric.
+ *
+ * @since 4.3.0
+ */
+@Private
+public interface CustomShuffleMetric {
+
+  /**
+   * The name of this metric. Must match the name reported by the corresponding
+   * {@link CustomShuffleTaskMetric} so per-task values can be matched to this declaration.
+   */
+  String name();
+
+  /**
+   * A human-readable description of this metric.
+   */
+  String description();
+
+  /**
+   * Defines how the per-task values are represented as a string in the UI. Per-task values are
+   * always aggregated as a plain sum across tasks. Must be one of:
+   * {@code "sum"} (raw count), {@code "size"} (bytes), {@code "timing"} (milliseconds), or
+   * {@code "nsTiming"} (nanoseconds). Any other value is rejected.
+   */
+  String metricType();
+}

core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleTaskMetric.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.api.metric;
+
+import org.apache.spark.annotation.Private;
+
+/**
+ * :: Private ::
+ * A per-task value for a {@link CustomShuffleMetric}.
+ *
+ * @since 4.3.0
+ */
+@Private
+public interface CustomShuffleTaskMetric {
+
+  /**
+   * The name of this metric. Must match the name of the declaring {@link CustomShuffleMetric}.
+   */
+  String name();
+
+  /**
+   * Returns the long value of custom shuffle task metric.
+   */
+  long value();
+}

core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java

@@ -48,6 +48,7 @@
 import org.apache.spark.shuffle.api.ShuffleMapOutputWriter;
 import org.apache.spark.shuffle.api.ShufflePartitionWriter;
 import org.apache.spark.shuffle.api.WritableByteChannelWrapper;
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric;
 import org.apache.spark.shuffle.checksum.ShuffleChecksumSupport;
 import org.apache.spark.internal.config.package$;
 import org.apache.spark.scheduler.MapStatus;
@@ -104,6 +105,7 @@ final class BypassMergeSortShuffleWriter<K, V>
   private FileSegment[] partitionWriterSegments;
   @Nullable private MapStatus mapStatus;
   private long[] partitionLengths;
+  private CustomShuffleTaskMetric[] customMetricsValues = new CustomShuffleTaskMetric[0];
   /** Checksum calculator for each partition. Empty when shuffle checksum disabled. */
   private final Checksum[] partitionChecksums;
   /**
@@ -153,6 +155,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
       if (!records.hasNext()) {
         partitionLengths = mapOutputWriter.commitAllPartitions(
           ShuffleChecksumHelper.EMPTY_CHECKSUM_VALUE).getPartitionLengths();
+        customMetricsValues = mapOutputWriter.currentMetricsValues();
         mapStatus = MapStatus$.MODULE$.apply(
           blockManager.shuffleServerId(), partitionLengths, mapId, getAggregatedChecksumValue());
         return;
@@ -213,6 +216,11 @@ public long[] getPartitionLengths() {
     return partitionLengths;
   }
 
+  @Override
+  public CustomShuffleTaskMetric[] currentMetricsValues() {
+    return customMetricsValues;
+  }
+
   // For test only.
   @VisibleForTesting
   RowBasedChecksum[] getRowBasedChecksums() {
@@ -261,8 +269,11 @@ private long[] writePartitionedData(ShuffleMapOutputWriter mapOutputWriter) thro
       }
       partitionWriters = null;
     }
-    return mapOutputWriter.commitAllPartitions(getChecksumValues(partitionChecksums))
-      .getPartitionLengths();
+    long[] committedPartitionLengths =
+      mapOutputWriter.commitAllPartitions(getChecksumValues(partitionChecksums))
+        .getPartitionLengths();
+    customMetricsValues = mapOutputWriter.currentMetricsValues();
+    return committedPartitionLengths;
   }
 
   private void writePartitionedDataWithChannel(

core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java

@@ -59,6 +59,7 @@
 import org.apache.spark.shuffle.api.ShufflePartitionWriter;
 import org.apache.spark.shuffle.api.SingleSpillShuffleMapOutputWriter;
 import org.apache.spark.shuffle.api.WritableByteChannelWrapper;
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric;
 import org.apache.spark.shuffle.checksum.RowBasedChecksum;
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.storage.TimeTrackingOutputStream;
@@ -93,6 +94,7 @@ public class UnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V> {
   @Nullable private MapStatus mapStatus;
   @Nullable private ShuffleExternalSorter sorter;
   @Nullable private long[] partitionLengths;
+  private CustomShuffleTaskMetric[] customMetricsValues = new CustomShuffleTaskMetric[0];
   private long peakMemoryUsedBytes = 0;
 
   private ExposedBufferByteArrayOutputStream serBuffer;
@@ -288,8 +290,11 @@ private long[] mergeSpills(SpillInfo[] spills) throws IOException {
     if (spills.length == 0) {
       final ShuffleMapOutputWriter mapWriter = shuffleExecutorComponents
           .createMapOutputWriter(shuffleId, mapId, partitioner.numPartitions());
-      return mapWriter.commitAllPartitions(
-        ShuffleChecksumHelper.EMPTY_CHECKSUM_VALUE).getPartitionLengths();
+      long[] emptyPartitionLengths =
+        mapWriter.commitAllPartitions(ShuffleChecksumHelper.EMPTY_CHECKSUM_VALUE)
+          .getPartitionLengths();
+      customMetricsValues = mapWriter.currentMetricsValues();
+      return emptyPartitionLengths;
     } else if (spills.length == 1) {
       Optional<SingleSpillShuffleMapOutputWriter> maybeSingleFileWriter =
           shuffleExecutorComponents.createSingleFileMapOutputWriter(shuffleId, mapId);
@@ -299,8 +304,10 @@ private long[] mergeSpills(SpillInfo[] spills) throws IOException {
         partitionLengths = spills[0].partitionLengths;
         logger.debug("Merge shuffle spills for mapId {} with length {}", mapId,
             partitionLengths.length);
-        maybeSingleFileWriter.get()
-          .transferMapSpillFile(spills[0].file, partitionLengths, sorter.getChecksums());
+        SingleSpillShuffleMapOutputWriter singleFileWriter = maybeSingleFileWriter.get();
+        singleFileWriter.transferMapSpillFile(spills[0].file, partitionLengths,
+            sorter.getChecksums());
+        customMetricsValues = singleFileWriter.currentMetricsValues();
       } else {
         partitionLengths = mergeSpillsUsingStandardWriter(spills);
       }
@@ -348,6 +355,7 @@ private long[] mergeSpillsUsingStandardWriter(SpillInfo[] spills) throws IOExcep
         mergeSpillsWithFileStream(spills, mapWriter, compressionCodec);
       }
       partitionLengths = mapWriter.commitAllPartitions(sorter.getChecksums()).getPartitionLengths();
+      customMetricsValues = mapWriter.currentMetricsValues();
     } catch (Exception e) {
       try {
         mapWriter.abort(e);
@@ -565,4 +573,9 @@ public void close() throws IOException {
   public long[] getPartitionLengths() {
     return partitionLengths;
   }
+
+  @Override
+  public CustomShuffleTaskMetric[] currentMetricsValues() {
+    return customMetricsValues;
+  }
 }

core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala

@@ -21,6 +21,7 @@ import org.apache.spark.{ShuffleDependency, SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.LogKeys.{NUM_MERGER_LOCATIONS, SHUFFLE_ID, STAGE_ID}
 import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric
 import org.apache.spark.util.PostStatusUpdateListener
 
 /**
@@ -37,6 +38,11 @@ private[spark] class ShuffleWriteProcessor extends Serializable with Logging {
     context.taskMetrics().shuffleWriteMetrics
   }
 
+  /**
+   * Report the custom shuffle metrics of a [[ShuffleWriter]] for this task.
+   */
+  protected def reportCustomMetrics(metricsValues: Array[CustomShuffleTaskMetric]): Unit = {}
+
   /**
    * The write process for particular partition, it controls the life circle of [[ShuffleWriter]]
    * get from [[ShuffleManager]] finally return the [[MapStatus]] for this task.
@@ -57,6 +63,7 @@ private[spark] class ShuffleWriteProcessor extends Serializable with Logging {
         createMetricsReporter(context))
       writer.write(inputs.asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
       val mapStatus = writer.stop(success = true)
+      reportCustomMetrics(writer.currentMetricsValues())
       if (mapStatus.isDefined) {
         // Check if sufficient shuffle mergers are available now for the ShuffleMapTask to push
         if (dep.shuffleMergeAllowed && dep.getMergerLocs.isEmpty) {

core/src/main/scala/org/apache/spark/shuffle/ShuffleWriter.scala

@@ -20,6 +20,7 @@ package org.apache.spark.shuffle
 import java.io.IOException
 
 import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric
 
 /**
  * Obtained inside a map task to write out records to the shuffle system.
@@ -34,4 +35,10 @@ private[spark] abstract class ShuffleWriter[K, V] {
 
   /** Get the lengths of each partition */
   def getPartitionLengths(): Array[Long]
+
+  /**
+   * The custom shuffle metrics reported by the underlying `ShuffleMapOutputWriter` for this task,
+   * available after a successful [[write]]. By default returns an empty array.
+   */
+  def currentMetricsValues(): Array[CustomShuffleTaskMetric] = Array.empty
 }

core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala

@@ -23,6 +23,7 @@ import org.apache.spark.scheduler.MapStatus
 import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleWriter}
 import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
 import org.apache.spark.shuffle.api.ShuffleExecutorComponents
+import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric
 import org.apache.spark.shuffle.checksum.RowBasedChecksum
 import org.apache.spark.util.collection.ExternalSorter
 
@@ -49,6 +50,8 @@ private[spark] class SortShuffleWriter[K, V, C](
 
   private var partitionLengths: Array[Long] = _
 
+  private var customMetricsValues: Array[CustomShuffleTaskMetric] = Array.empty
+
   def getRowBasedChecksums: Array[RowBasedChecksum] = {
     if (sorter != null) {
       sorter.getRowBasedChecksums
@@ -84,6 +87,7 @@ private[spark] class SortShuffleWriter[K, V, C](
       dep.shuffleId, mapId, dep.partitioner.numPartitions)
     sorter.writePartitionedMapOutput(dep.shuffleId, mapId, mapOutputWriter, writeMetrics)
     partitionLengths = mapOutputWriter.commitAllPartitions(sorter.getChecksums).getPartitionLengths
+    customMetricsValues = mapOutputWriter.currentMetricsValues()
     mapStatus =
       MapStatus(blockManager.shuffleServerId, partitionLengths, mapId, getAggregatedChecksumValue)
   }
@@ -112,6 +116,8 @@ private[spark] class SortShuffleWriter[K, V, C](
   }
 
   override def getPartitionLengths(): Array[Long] = partitionLengths
+
+  override def currentMetricsValues(): Array[CustomShuffleTaskMetric] = customMetricsValues
 }
 
 private[spark] object SortShuffleWriter {