From 73f3ea0f9e32d7a5fca9d4e332d2ff8b63b29f39 Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Wed, 15 Jul 2026 21:46:32 -0700 Subject: [PATCH 1/6] Custom shuffle metrics --- .../shuffle/api/ShuffleDriverComponents.java | 9 ++ .../shuffle/api/ShuffleMapOutputWriter.java | 8 + .../api/metric/CustomShuffleMetric.java | 49 ++++++ .../api/metric/CustomShuffleTaskMetric.java | 40 +++++ .../sort/BypassMergeSortShuffleWriter.java | 15 +- .../shuffle/sort/UnsafeShuffleWriter.java | 15 +- .../spark/shuffle/ShuffleWriteProcessor.scala | 7 + .../apache/spark/shuffle/ShuffleWriter.scala | 7 + .../shuffle/sort/SortShuffleWriter.scala | 6 + .../sort/UnsafeShuffleWriterSuite.java | 33 +++- .../BypassMergeSortShuffleWriterSuite.scala | 45 ++++++ .../sort/CustomMetricShuffleTestUtils.scala | 75 +++++++++ .../shuffle/sort/SortShuffleWriterSuite.scala | 30 ++++ .../exchange/ShuffleExchangeExec.scala | 30 +++- .../metric/CustomShuffleMetrics.scala | 61 +++++++ ...CustomShuffleMetricsIntegrationSuite.scala | 151 ++++++++++++++++++ .../metric/CustomShuffleMetricsSuite.scala | 86 ++++++++++ 17 files changed, 655 insertions(+), 12 deletions(-) create mode 100644 core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java create mode 100644 core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleTaskMetric.java create mode 100644 core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala diff --git a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleDriverComponents.java b/core/src/main/java/org/apache/spark/shuffle/api/ShuffleDriverComponents.java index 5c4c1eff9f59..fd169fad413a 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleDriverComponents.java +++ b/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]; + } } diff --git a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java b/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java index 2237ec052cac..267cf4028496 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java +++ b/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]; + } } diff --git a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java new file mode 100644 index 000000000000..9513c7b660d4 --- /dev/null +++ b/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(); +} diff --git a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleTaskMetric.java b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleTaskMetric.java new file mode 100644 index 000000000000..7e96be930323 --- /dev/null +++ b/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(); +} diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java index 5acc66e12063..aae8746ea138 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java +++ b/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 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> 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( diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java index e3ecfed32348..4d1ac3b2bd21 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java +++ b/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 extends ShuffleWriter { @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 maybeSingleFileWriter = shuffleExecutorComponents.createSingleFileMapOutputWriter(shuffleId, mapId); @@ -348,6 +353,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 +571,9 @@ public void close() throws IOException { public long[] getPartitionLengths() { return partitionLengths; } + + @Override + public CustomShuffleTaskMetric[] currentMetricsValues() { + return customMetricsValues; + } } diff --git a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala index 14d2f3e9b11e..846bf793be21 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriteProcessor.scala +++ b/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) { diff --git a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriter.scala b/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriter.scala index a279b4c8f42f..47ac9c622134 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/ShuffleWriter.scala +++ b/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 } diff --git a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala b/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala index a7ac20016a0e..4c91e4c2762b 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala +++ b/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 { diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java index 3bddc8afcd5f..7da2f164ad14 100644 --- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java +++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java @@ -50,6 +50,8 @@ import org.apache.spark.scheduler.MapStatus; import org.apache.spark.security.CryptoStreamUtils; import org.apache.spark.serializer.*; +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.shuffle.IndexShuffleBlockResolver; import org.apache.spark.shuffle.sort.io.LocalDiskShuffleExecutorComponents; @@ -198,6 +200,13 @@ private UnsafeShuffleWriter createWriter(boolean transferToEnabl private UnsafeShuffleWriter createWriter( boolean transferToEnabled, IndexShuffleBlockResolver blockResolver) throws SparkException { + return createWriter(transferToEnabled, + new LocalDiskShuffleExecutorComponents(conf, blockManager, blockResolver)); + } + + private UnsafeShuffleWriter createWriter( + boolean transferToEnabled, + ShuffleExecutorComponents executorComponents) throws SparkException { conf.set(package$.MODULE$.SHUFFLE_MERGE_PREFER_NIO().key(), String.valueOf(transferToEnabled)); return new UnsafeShuffleWriter<>( @@ -208,7 +217,7 @@ private UnsafeShuffleWriter createWriter( taskContext, conf, taskContext.taskMetrics().shuffleWriteMetrics(), - new LocalDiskShuffleExecutorComponents(conf, blockManager, blockResolver)); + executorComponents); } private void assertSpillFilesWereCleanedUp() { @@ -289,6 +298,28 @@ public void writeEmptyIterator() throws Exception { assertEquals(0, taskMetrics.memoryBytesSpilled()); } + @Test + public void reportsNoCustomShuffleMetricsByDefault() throws Exception { + final UnsafeShuffleWriter writer = createWriter(true); + writer.write(Collections.emptyIterator()); + writer.stop(true); + assertEquals(0, writer.currentMetricsValues().length); + } + + @Test + public void exposesCustomShuffleMetricsReportedByMapOutputWriter() throws Exception { + CustomShuffleTaskMetric[] reported = new CustomShuffleTaskMetric[] { + new TestCustomShuffleTaskMetric("s3BytesUploaded", 4096L) + }; + final UnsafeShuffleWriter writer = createWriter(true, + new CustomMetricReportingExecutorComponents( + new LocalDiskShuffleExecutorComponents(conf, blockManager, shuffleBlockResolver), + reported)); + writer.write(Collections.emptyIterator()); + writer.stop(true); + assertArrayEquals(reported, writer.currentMetricsValues()); + } + @Test public void writeWithoutSpilling() throws Exception { // In this example, each partition should have exactly one record: diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala index c9b951cf0369..4a22bda5f8a7 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala @@ -37,6 +37,7 @@ import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper import org.apache.spark.serializer.{JavaSerializer, SerializerInstance, SerializerManager} import org.apache.spark.shuffle.{IndexShuffleBlockResolver, ShuffleChecksumTestHelper} import org.apache.spark.shuffle.api.ShuffleExecutorComponents +import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric import org.apache.spark.shuffle.sort.io.LocalDiskShuffleExecutorComponents import org.apache.spark.storage._ import org.apache.spark.util.Utils @@ -157,6 +158,50 @@ class BypassMergeSortShuffleWriterSuite when(dependency.rowBasedChecksums).thenReturn(rowBasedChecksums) } + test("reports no custom shuffle metrics by default") { + val writer = new BypassMergeSortShuffleWriter[Int, Int]( + blockManager, + shuffleHandle, + 0L, // MapId + conf, + taskContext.taskMetrics().shuffleWriteMetrics, + shuffleExecutorComponents) + writer.write(Iterator((1, 1), (2, 2), (3, 3))) + writer.stop( /* success = */ true) + assert(writer.currentMetricsValues().isEmpty) + } + + test("exposes custom shuffle metrics reported by the map output writer") { + val reported = Array[CustomShuffleTaskMetric]( + TestCustomShuffleTaskMetric("s3BytesUploaded", 2048L)) + val writer = new BypassMergeSortShuffleWriter[Int, Int]( + blockManager, + shuffleHandle, + 0L, // MapId + conf, + taskContext.taskMetrics().shuffleWriteMetrics, + new CustomMetricReportingExecutorComponents(shuffleExecutorComponents, reported)) + writer.write(Iterator((1, 1), (2, 2), (3, 3))) + writer.stop( /* success = */ true) + assert(writer.currentMetricsValues() === reported) + } + + test("exposes custom shuffle metrics reported by the map output writer " + + "on the empty-iterator path") { + val reported = Array[CustomShuffleTaskMetric]( + TestCustomShuffleTaskMetric("s3BlockUploads", 0L)) + val writer = new BypassMergeSortShuffleWriter[Int, Int]( + blockManager, + shuffleHandle, + 0L, // MapId + conf, + taskContext.taskMetrics().shuffleWriteMetrics, + new CustomMetricReportingExecutorComponents(shuffleExecutorComponents, reported)) + writer.write(Iterator.empty) + writer.stop( /* success = */ true) + assert(writer.currentMetricsValues() === reported) + } + test("write empty iterator") { val writer = new BypassMergeSortShuffleWriter[Int, Int]( blockManager, diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala new file mode 100644 index 000000000000..c81181b82a3b --- /dev/null +++ b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala @@ -0,0 +1,75 @@ +/* + * 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.sort + +import java.util.Optional + +import org.apache.spark.shuffle.api.{ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter, SingleSpillShuffleMapOutputWriter} +import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage +import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric + +/** + * A [[CustomShuffleTaskMetric]] with a fixed name and value, for tests. + */ +private[spark] case class TestCustomShuffleTaskMetric(metricName: String, metricValue: Long) + extends CustomShuffleTaskMetric { + override def name(): String = metricName + override def value(): Long = metricValue +} + +/** + * Wraps a [[ShuffleExecutorComponents]] so every [[ShuffleMapOutputWriter]] it creates reports the + * given custom metric values from `currentMetricsValues()`, delegating all real writing to the + * wrapped components. Lets tests assert that the concrete + * [[org.apache.spark.shuffle.ShuffleWriter]] implementations stash and expose what the map output + * writer reports. + */ +private[spark] class CustomMetricReportingExecutorComponents( + delegate: ShuffleExecutorComponents, + reportedMetrics: Array[CustomShuffleTaskMetric]) + extends ShuffleExecutorComponents { + + override def initializeExecutor( + appId: String, execId: String, extraConfigs: java.util.Map[String, String]): Unit = + delegate.initializeExecutor(appId, execId, extraConfigs) + + override def createMapOutputWriter( + shuffleId: Int, mapTaskId: Long, numPartitions: Int): ShuffleMapOutputWriter = + new CustomMetricReportingMapOutputWriter( + delegate.createMapOutputWriter(shuffleId, mapTaskId, numPartitions), reportedMetrics) + + override def createSingleFileMapOutputWriter( + shuffleId: Int, mapId: Long): Optional[SingleSpillShuffleMapOutputWriter] = + delegate.createSingleFileMapOutputWriter(shuffleId, mapId) +} + +private[spark] class CustomMetricReportingMapOutputWriter( + delegate: ShuffleMapOutputWriter, + reportedMetrics: Array[CustomShuffleTaskMetric]) + extends ShuffleMapOutputWriter { + + override def getPartitionWriter(reducePartitionId: Int): ShufflePartitionWriter = + delegate.getPartitionWriter(reducePartitionId) + + override def commitAllPartitions(checksums: Array[Long]): MapOutputCommitMessage = + delegate.commitAllPartitions(checksums) + + override def abort(error: Throwable): Unit = delegate.abort(error) + + override def currentMetricsValues(): Array[CustomShuffleTaskMetric] = reportedMetrics +} diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/SortShuffleWriterSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/SortShuffleWriterSuite.scala index 9d4b0625f762..4ef58fd357be 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/sort/SortShuffleWriterSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/sort/SortShuffleWriterSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.memory.MemoryTestingUtils import org.apache.spark.serializer.JavaSerializer import org.apache.spark.shuffle.{BaseShuffleHandle, IndexShuffleBlockResolver, ShuffleChecksumTestHelper} import org.apache.spark.shuffle.api.ShuffleExecutorComponents +import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric import org.apache.spark.shuffle.sort.io.LocalDiskShuffleExecutorComponents import org.apache.spark.storage.BlockManager import org.apache.spark.util.Utils @@ -128,6 +129,35 @@ class SortShuffleWriterSuite assert(records.size === writeMetrics.recordsWritten) } + test("reports no custom shuffle metrics by default") { + val context = MemoryTestingUtils.fakeTaskContext(sc.env) + val writer = new SortShuffleWriter[Int, Int, Int]( + shuffleHandle, + mapId = 3, + context, + context.taskMetrics().shuffleWriteMetrics, + shuffleExecutorComponents) + writer.write(List[(Int, Int)]((1, 2), (2, 3)).iterator) + writer.stop(success = true) + assert(writer.currentMetricsValues().isEmpty) + } + + test("exposes custom shuffle metrics reported by the map output writer") { + val context = MemoryTestingUtils.fakeTaskContext(sc.env) + val reported = Array[CustomShuffleTaskMetric]( + TestCustomShuffleTaskMetric("s3BytesUploaded", 1024L), + TestCustomShuffleTaskMetric("s3BlockUploads", 2L)) + val writer = new SortShuffleWriter[Int, Int, Int]( + shuffleHandle, + mapId = 4, + context, + context.taskMetrics().shuffleWriteMetrics, + new CustomMetricReportingExecutorComponents(shuffleExecutorComponents, reported)) + writer.write(List[(Int, Int)]((1, 2), (2, 3), (4, 4)).iterator) + writer.stop(success = true) + assert(writer.currentMetricsValues() === reported) + } + test("Row-based checksums are independent of input row order") { val shuffleBlockResolver = new IndexShuffleBlockResolver(conf) val context = MemoryTestingUtils.fakeTaskContext(sc.env) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index dd829e697df6..cc3b8f9ec44b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -28,6 +28,7 @@ import org.apache.spark.internal.config import org.apache.spark.rdd.{RDD, RDDOperationScope} import org.apache.spark.serializer.Serializer import org.apache.spark.shuffle.{ShuffleWriteMetricsReporter, ShuffleWriteProcessor} +import org.apache.spark.shuffle.api.metric.CustomShuffleTaskMetric import org.apache.spark.shuffle.sort.SortShuffleManager import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{ @@ -39,7 +40,7 @@ import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} +import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.util.{MutablePair, ThreadUtils} import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator} @@ -198,10 +199,15 @@ case class ShuffleExchangeExec( SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) private[sql] lazy val readMetrics = SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) + + private lazy val customWriteMetrics: Map[String, SQLMetric] = + CustomShuffleMetrics.createMetrics( + sparkContext, sparkContext.shuffleDriverComponents.supportedCustomMetrics()) + override lazy val metrics = Map( "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"), "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions") - ) ++ readMetrics ++ writeMetrics + ) ++ readMetrics ++ writeMetrics ++ customWriteMetrics override def nodeName: String = "Exchange" @@ -252,7 +258,8 @@ case class ShuffleExchangeExec( child.output, outputPartitioning, serializer, - writeMetrics) + writeMetrics, + customWriteMetrics) metrics("numPartitions").set(dep.partitioner.numPartitions) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates( @@ -346,7 +353,8 @@ object ShuffleExchangeExec { outputAttributes: Seq[Attribute], newPartitioning: Partitioning, serializer: Serializer, - writeMetrics: Map[String, SQLMetric]) + writeMetrics: Map[String, SQLMetric], + customWriteMetrics: Map[String, SQLMetric] = Map.empty) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -541,7 +549,7 @@ object ShuffleExchangeExec { rddWithPartitionIds, new PartitionIdPassthrough(part.numPartitions), serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics, customWriteMetrics), rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, checksumMismatchQueryLevelRollbackEnabled = @@ -552,14 +560,22 @@ object ShuffleExchangeExec { /** * Create a customized [[ShuffleWriteProcessor]] for SQL which wrap the default metrics reporter - * with [[SQLShuffleWriteMetricsReporter]] as new reporter for [[ShuffleWriteProcessor]]. + * with [[SQLShuffleWriteMetricsReporter]] as new reporter for [[ShuffleWriteProcessor]], and + * reports any custom shuffle metrics into `customWriteMetrics`. */ - def createShuffleWriteProcessor(metrics: Map[String, SQLMetric]): ShuffleWriteProcessor = { + def createShuffleWriteProcessor( + metrics: Map[String, SQLMetric], + customWriteMetrics: Map[String, SQLMetric] = Map.empty): ShuffleWriteProcessor = { new ShuffleWriteProcessor { override protected def createMetricsReporter( context: TaskContext): ShuffleWriteMetricsReporter = { new SQLShuffleWriteMetricsReporter(context.taskMetrics().shuffleWriteMetrics, metrics) } + + override protected def reportCustomMetrics( + metricsValues: Array[CustomShuffleTaskMetric]): Unit = { + CustomShuffleMetrics.updateMetrics(metricsValues, customWriteMetrics) + } } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala new file mode 100644 index 000000000000..3e8b59c93ecb --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala @@ -0,0 +1,61 @@ +/* + * 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.sql.execution.metric + +import org.apache.spark.SparkContext +import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} +import org.apache.spark.util.MetricUtils + +object CustomShuffleMetrics { + + /** + * Creates [[SQLMetric]]s for the given declared custom shuffle metrics, keyed by metric name. + */ + def createMetrics( + sc: SparkContext, + metrics: Array[CustomShuffleMetric]): Map[String, SQLMetric] = { + metrics.map { metric => + val label = metric.description() + // AVERAGE_METRIC is intentionally unsupported: its per-task values must be stored pre-scaled + // via SQLMetric.set(Double), whereas custom task metrics report a plain Long. + val acc = metric.metricType() match { + case MetricUtils.SUM_METRIC => SQLMetrics.createMetric(sc, label) + case MetricUtils.SIZE_METRIC => SQLMetrics.createSizeMetric(sc, label) + case MetricUtils.TIMING_METRIC => SQLMetrics.createTimingMetric(sc, label) + case MetricUtils.NS_TIMING_METRIC => SQLMetrics.createNanoTimingMetric(sc, label) + case other => throw new IllegalArgumentException( + s"Unsupported custom shuffle metric type '$other' for metric '${metric.name()}'. " + + s"Supported types: ${MetricUtils.SUM_METRIC}, ${MetricUtils.SIZE_METRIC}, " + + s"${MetricUtils.TIMING_METRIC}, ${MetricUtils.NS_TIMING_METRIC}.") + } + metric.name() -> acc + }.toMap + } + + /** + * Updates the custom-metric [[SQLMetric]]s with the per-task reported values, matching by name. + * Reported values with no matching declaration are ignored. + */ + def updateMetrics( + taskMetricsValues: Array[CustomShuffleTaskMetric], + customMetrics: Map[String, SQLMetric]): Unit = { + taskMetricsValues.foreach { metric => + customMetrics.get(metric.name()).foreach(_.set(metric.value())) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala new file mode 100644 index 000000000000..4daa86ea88a9 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala @@ -0,0 +1,151 @@ +/* + * 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.sql.execution.metric + +import java.util.{Map => JMap} + +import org.apache.spark.{SparkConf, SparkFunSuite} +import org.apache.spark.internal.config.SHUFFLE_IO_PLUGIN_CLASS +import org.apache.spark.internal.config.UI.UI_ENABLED +import org.apache.spark.shuffle.api.{ShuffleDataIO, ShuffleDriverComponents, ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter, SingleSpillShuffleMapOutputWriter} +import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage +import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} +import org.apache.spark.shuffle.sort.io.LocalDiskShuffleDataIO +import org.apache.spark.sql.{LocalSparkSession, SparkSession} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.util.MetricUtils + +/** + * Integration coverage for the shuffle custom-metrics SPI: a toy [[ShuffleDataIO]] declares a + * custom metric on its driver components and reports a fixed per-task value from every map output + * writer, and we assert the value surfaces on the [[ShuffleExchangeExec]] operator node after a + * real shuffle. This exercises the full path (driver declaration, task reporting, accumulator + * propagation, and the SQL-side bridge) without a `ThreadLocal` side-channel. + */ +class CustomShuffleMetricsIntegrationSuite + extends SparkFunSuite with LocalSparkSession with AdaptiveSparkPlanHelper { + + private def withPluginSession(pluginClass: Option[String])(f: SparkSession => Unit): Unit = { + val conf = new SparkConf(false) + .setMaster("local[2]") + .setAppName("custom-shuffle-metrics") + .set(UI_ENABLED, false) + pluginClass.foreach(conf.set(SHUFFLE_IO_PLUGIN_CLASS, _)) + spark = SparkSession.builder().config(conf).getOrCreate() + LocalSparkSession.withSparkSession(spark)(f) + } + + private def runShuffleAndGetExchange(spark: SparkSession): ShuffleExchangeExec = { + import spark.implicits._ + val df = spark.range(0, 100, 1, 4).map(id => (id % 8, id)).toDF("key", "value") + .groupBy("key").count() + df.collect() + collectFirst(df.queryExecution.executedPlan) { + case e: ShuffleExchangeExec => e + }.getOrElse(fail("query plan had no ShuffleExchangeExec")) + } + + test("a declared custom shuffle metric surfaces on the exchange node after a shuffle") { + withPluginSession(Some(classOf[TestCustomMetricShuffleDataIO].getName)) { spark => + val exchange = runShuffleAndGetExchange(spark) + assert(exchange.metrics.contains(TestCustomMetricShuffleDataIO.MetricName)) + // Four map tasks each report a fixed value; the metric aggregates as a plain sum. + assert(exchange.metrics(TestCustomMetricShuffleDataIO.MetricName).value === + 4 * TestCustomMetricShuffleDataIO.PerTaskValue) + } + } + + test("the built-in local-disk plugin declares no custom metrics (regression guard)") { + withPluginSession(None) { spark => + val exchange = runShuffleAndGetExchange(spark) + assert(!exchange.metrics.contains(TestCustomMetricShuffleDataIO.MetricName)) + // Only the built-in exchange metrics are present; no custom key leaks in. + assert(exchange.metrics.keySet.forall(!_.startsWith("test"))) + } + } +} + +object TestCustomMetricShuffleDataIO { + val MetricName: String = "testShuffleBytesUploaded" + val PerTaskValue: Long = 1024L +} + +/** + * A [[ShuffleDataIO]] that delegates all real writing to the local-disk implementation but declares + * one custom metric on the driver and reports a fixed per-task value from every map output writer. + */ +class TestCustomMetricShuffleDataIO(sparkConf: SparkConf) extends ShuffleDataIO { + private val delegate = new LocalDiskShuffleDataIO(sparkConf) + + override def driver(): ShuffleDriverComponents = + new TestCustomMetricDriverComponents(delegate.driver()) + + override def executor(): ShuffleExecutorComponents = + new TestCustomMetricExecutorComponents(delegate.executor()) +} + +class TestCustomMetricDriverComponents(delegate: ShuffleDriverComponents) + extends ShuffleDriverComponents { + + override def initializeApplication(): JMap[String, String] = delegate.initializeApplication() + + override def cleanupApplication(): Unit = delegate.cleanupApplication() + + override def supportedCustomMetrics(): Array[CustomShuffleMetric] = Array( + new CustomShuffleMetric { + override def name(): String = TestCustomMetricShuffleDataIO.MetricName + override def description(): String = "test shuffle bytes uploaded" + override def metricType(): String = MetricUtils.SIZE_METRIC + }) +} + +class TestCustomMetricExecutorComponents(delegate: ShuffleExecutorComponents) + extends ShuffleExecutorComponents { + + override def initializeExecutor( + appId: String, execId: String, extraConfigs: JMap[String, String]): Unit = + delegate.initializeExecutor(appId, execId, extraConfigs) + + override def createMapOutputWriter( + shuffleId: Int, mapTaskId: Long, numPartitions: Int): ShuffleMapOutputWriter = + new TestCustomMetricMapOutputWriter( + delegate.createMapOutputWriter(shuffleId, mapTaskId, numPartitions)) + + override def createSingleFileMapOutputWriter( + shuffleId: Int, mapId: Long): java.util.Optional[SingleSpillShuffleMapOutputWriter] = + delegate.createSingleFileMapOutputWriter(shuffleId, mapId) +} + +class TestCustomMetricMapOutputWriter(delegate: ShuffleMapOutputWriter) + extends ShuffleMapOutputWriter { + + override def getPartitionWriter(reducePartitionId: Int): ShufflePartitionWriter = + delegate.getPartitionWriter(reducePartitionId) + + override def commitAllPartitions(checksums: Array[Long]): MapOutputCommitMessage = + delegate.commitAllPartitions(checksums) + + override def abort(error: Throwable): Unit = delegate.abort(error) + + override def currentMetricsValues(): Array[CustomShuffleTaskMetric] = Array( + new CustomShuffleTaskMetric { + override def name(): String = TestCustomMetricShuffleDataIO.MetricName + override def value(): Long = TestCustomMetricShuffleDataIO.PerTaskValue + }) +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala new file mode 100644 index 000000000000..e3ca50ad15e7 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala @@ -0,0 +1,86 @@ +/* + * 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.sql.execution.metric + +import org.apache.spark.SharedSparkContext +import org.apache.spark.SparkFunSuite +import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} +import org.apache.spark.util.MetricUtils + +class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext { + + test("createMetrics maps each declared metric to a SQLMetric of the declared type") { + val metrics = CustomShuffleMetrics.createMetrics(sc, Array( + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), + shuffleMetric("s3BlockUploads", "s3 block uploads", MetricUtils.SUM_METRIC), + shuffleMetric("s3FirstByteLatency", "s3 first byte latency", MetricUtils.TIMING_METRIC))) + + assert(metrics.keySet === Set("s3BytesUploaded", "s3BlockUploads", "s3FirstByteLatency")) + assert(metrics("s3BytesUploaded").metricType === MetricUtils.SIZE_METRIC) + assert(metrics("s3BlockUploads").metricType === MetricUtils.SUM_METRIC) + assert(metrics("s3FirstByteLatency").metricType === MetricUtils.TIMING_METRIC) + } + + test("createMetrics rejects an unsupported metric type") { + val e = intercept[IllegalArgumentException] { + CustomShuffleMetrics.createMetrics(sc, Array( + shuffleMetric("s3Avg", "s3 avg", MetricUtils.AVERAGE_METRIC))) + } + assert(e.getMessage.contains("s3Avg")) + assert(e.getMessage.contains(MetricUtils.AVERAGE_METRIC)) + } + + test("updateMetrics folds reported values into the matching SQLMetrics by name") { + val metrics = CustomShuffleMetrics.createMetrics(sc, Array( + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), + shuffleMetric("s3BlockUploads", "s3 block uploads", MetricUtils.SUM_METRIC))) + + CustomShuffleMetrics.updateMetrics( + Array(taskMetric("s3BytesUploaded", 1024L), taskMetric("s3BlockUploads", 3L)), metrics) + + assert(metrics("s3BytesUploaded").value === 1024L) + assert(metrics("s3BlockUploads").value === 3L) + } + + test("updateMetrics ignores reported values with no matching declaration") { + val metrics = CustomShuffleMetrics.createMetrics(sc, Array( + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC))) + + CustomShuffleMetrics.updateMetrics( + Array(taskMetric("s3BytesUploaded", 512L), taskMetric("undeclared", 99L)), metrics) + + assert(metrics("s3BytesUploaded").value === 512L) + assert(!metrics.contains("undeclared")) + } + + private def shuffleMetric( + metricName: String, desc: String, mType: String): CustomShuffleMetric = { + new CustomShuffleMetric { + override def name(): String = metricName + override def description(): String = desc + override def metricType(): String = mType + } + } + + private def taskMetric(metricName: String, v: Long): CustomShuffleTaskMetric = { + new CustomShuffleTaskMetric { + override def name(): String = metricName + override def value(): Long = v + } + } +} From b2d8b8178deaa5a2cfdddb7eada41be2bd459363 Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Sun, 19 Jul 2026 19:27:47 -0700 Subject: [PATCH 2/6] Address review comments --- .../SingleSpillShuffleMapOutputWriter.java | 8 ++++ .../shuffle/sort/UnsafeShuffleWriter.java | 6 ++- .../sort/UnsafeShuffleWriterSuite.java | 19 ++++++++++ .../sort/CustomMetricShuffleTestUtils.scala | 16 ++++++++ .../exchange/ShuffleExchangeExec.scala | 24 +++++++++--- ...CustomShuffleMetricsIntegrationSuite.scala | 38 ++++++++++++++++++- 6 files changed, 102 insertions(+), 9 deletions(-) diff --git a/core/src/main/java/org/apache/spark/shuffle/api/SingleSpillShuffleMapOutputWriter.java b/core/src/main/java/org/apache/spark/shuffle/api/SingleSpillShuffleMapOutputWriter.java index ba3d5a603e05..f690bb4af345 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/SingleSpillShuffleMapOutputWriter.java +++ b/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]; + } } diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java index 4d1ac3b2bd21..9305137a5560 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java +++ b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java @@ -304,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); } diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java index 7da2f164ad14..47466b8b5c44 100644 --- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java +++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java @@ -320,6 +320,25 @@ public void exposesCustomShuffleMetricsReportedByMapOutputWriter() throws Except assertArrayEquals(reported, writer.currentMetricsValues()); } + @Test + public void exposesCustomShuffleMetricsFromSingleSpillWriter() throws Exception { + CustomShuffleTaskMetric[] reported = new CustomShuffleTaskMetric[] { + new TestCustomShuffleTaskMetric("s3BytesUploaded", 4096L) + }; + final UnsafeShuffleWriter writer = createWriter(true, + new CustomMetricReportingExecutorComponents( + new LocalDiskShuffleExecutorComponents(conf, blockManager, shuffleBlockResolver), + reported)); + final ArrayList> dataToWrite = new ArrayList<>(); + for (int i = 0; i < NUM_PARTITIONS; i++) { + dataToWrite.add(new Tuple2<>(i, i)); + } + writer.write(dataToWrite.iterator()); + writer.stop(true); + assertEquals(1, spillFilesCreated.size()); + assertArrayEquals(reported, writer.currentMetricsValues()); + } + @Test public void writeWithoutSpilling() throws Exception { // In this example, each partition should have exactly one record: diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala index c81181b82a3b..a7a569fb5e30 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala @@ -17,6 +17,7 @@ package org.apache.spark.shuffle.sort +import java.io.File import java.util.Optional import org.apache.spark.shuffle.api.{ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter, SingleSpillShuffleMapOutputWriter} @@ -56,6 +57,21 @@ private[spark] class CustomMetricReportingExecutorComponents( override def createSingleFileMapOutputWriter( shuffleId: Int, mapId: Long): Optional[SingleSpillShuffleMapOutputWriter] = delegate.createSingleFileMapOutputWriter(shuffleId, mapId) + .map[SingleSpillShuffleMapOutputWriter] { writer => + new CustomMetricReportingSingleSpillMapOutputWriter(writer, reportedMetrics) + } +} + +private[spark] class CustomMetricReportingSingleSpillMapOutputWriter( + delegate: SingleSpillShuffleMapOutputWriter, + reportedMetrics: Array[CustomShuffleTaskMetric]) + extends SingleSpillShuffleMapOutputWriter { + + override def transferMapSpillFile( + mapSpillFile: File, partitionLengths: Array[Long], checksums: Array[Long]): Unit = + delegate.transferMapSpillFile(mapSpillFile, partitionLengths, checksums) + + override def currentMetricsValues(): Array[CustomShuffleTaskMetric] = reportedMetrics } private[spark] class CustomMetricReportingMapOutputWriter( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index cc3b8f9ec44b..a09ea4c94985 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -24,6 +24,7 @@ import scala.collection.mutable import scala.concurrent.{ExecutionContext, Future, Promise} import org.apache.spark._ +import org.apache.spark.internal.LogKeys.METRIC_NAME import org.apache.spark.internal.config import org.apache.spark.rdd.{RDD, RDDOperationScope} import org.apache.spark.serializer.Serializer @@ -200,14 +201,25 @@ case class ShuffleExchangeExec( private[sql] lazy val readMetrics = SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - private lazy val customWriteMetrics: Map[String, SQLMetric] = - CustomShuffleMetrics.createMetrics( - sparkContext, sparkContext.shuffleDriverComponents.supportedCustomMetrics()) - - override lazy val metrics = Map( + private lazy val sparkMetrics: Map[String, SQLMetric] = Map( "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"), "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions") - ) ++ readMetrics ++ writeMetrics ++ customWriteMetrics + ) ++ readMetrics ++ writeMetrics + + // Drop plugin metrics whose names collide with Spark-owned ones so they can't shadow the real + // accumulators (e.g. `shuffleRecordsWritten`, which feeds AQE). + private lazy val customWriteMetrics: Map[String, SQLMetric] = { + val (reserved, allowed) = sparkContext.shuffleDriverComponents.supportedCustomMetrics() + .partition(metric => sparkMetrics.contains(metric.name())) + if (reserved.nonEmpty) { + logWarning(log"Ignoring custom shuffle metrics whose names collide with Spark-owned " + + log"Exchange metrics: " + + log"${MDC(METRIC_NAME, reserved.map(_.name()).sorted.mkString(", "))}.") + } + CustomShuffleMetrics.createMetrics(sparkContext, allowed) + } + + override lazy val metrics = sparkMetrics ++ customWriteMetrics override def nodeName: String = "Exchange" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala index 4daa86ea88a9..bfbc0b6c4a7c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala @@ -29,6 +29,7 @@ import org.apache.spark.shuffle.sort.io.LocalDiskShuffleDataIO import org.apache.spark.sql.{LocalSparkSession, SparkSession} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN import org.apache.spark.util.MetricUtils /** @@ -41,7 +42,7 @@ import org.apache.spark.util.MetricUtils class CustomShuffleMetricsIntegrationSuite extends SparkFunSuite with LocalSparkSession with AdaptiveSparkPlanHelper { - private def withPluginSession(pluginClass: Option[String])(f: SparkSession => Unit): Unit = { + private def withPluginSession[T](pluginClass: Option[String])(f: SparkSession => T): T = { val conf = new SparkConf(false) .setMaster("local[2]") .setAppName("custom-shuffle-metrics") @@ -79,6 +80,13 @@ class CustomShuffleMetricsIntegrationSuite assert(exchange.metrics.keySet.forall(!_.startsWith("test"))) } } + + test("a custom metric colliding with a Spark-owned name is dropped and cannot shadow it") { + withPluginSession(Some(classOf[TestCollidingMetricShuffleDataIO].getName)) { spark => + val exchange = runShuffleAndGetExchange(spark) + assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value > 0) + } + } } object TestCustomMetricShuffleDataIO { @@ -149,3 +157,31 @@ class TestCustomMetricMapOutputWriter(delegate: ShuffleMapOutputWriter) override def value(): Long = TestCustomMetricShuffleDataIO.PerTaskValue }) } + +/** + * A [[ShuffleDataIO]] that delegates all real work to the local-disk implementation but declares a + * custom driver metric using a Spark-owned name (`shuffleRecordsWritten`). + */ +class TestCollidingMetricShuffleDataIO(sparkConf: SparkConf) extends ShuffleDataIO { + private val delegate = new LocalDiskShuffleDataIO(sparkConf) + + override def driver(): ShuffleDriverComponents = + new TestCollidingMetricDriverComponents(delegate.driver()) + + override def executor(): ShuffleExecutorComponents = delegate.executor() +} + +class TestCollidingMetricDriverComponents(delegate: ShuffleDriverComponents) + extends ShuffleDriverComponents { + + override def initializeApplication(): JMap[String, String] = delegate.initializeApplication() + + override def cleanupApplication(): Unit = delegate.cleanupApplication() + + override def supportedCustomMetrics(): Array[CustomShuffleMetric] = Array( + new CustomShuffleMetric { + override def name(): String = SHUFFLE_RECORDS_WRITTEN + override def description(): String = "colliding shuffle records written" + override def metricType(): String = MetricUtils.SUM_METRIC + }) +} From 92a08ba67ecba5087dd112cbfbe0b91e216ca455 Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Tue, 21 Jul 2026 16:36:22 -0700 Subject: [PATCH 3/6] Address review comments: Fix missed custom metrics in limitExec --- .../exchange/ShuffleExchangeExec.scala | 19 ++----- .../apache/spark/sql/execution/limit.scala | 53 +++++++++++-------- .../metric/CustomShuffleMetrics.scala | 24 ++++++++- ...CustomShuffleMetricsIntegrationSuite.scala | 16 ++++++ 4 files changed, 73 insertions(+), 39 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index a09ea4c94985..d63019d3af8d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -24,7 +24,6 @@ import scala.collection.mutable import scala.concurrent.{ExecutionContext, Future, Promise} import org.apache.spark._ -import org.apache.spark.internal.LogKeys.METRIC_NAME import org.apache.spark.internal.config import org.apache.spark.rdd.{RDD, RDDOperationScope} import org.apache.spark.serializer.Serializer @@ -206,18 +205,8 @@ case class ShuffleExchangeExec( "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions") ) ++ readMetrics ++ writeMetrics - // Drop plugin metrics whose names collide with Spark-owned ones so they can't shadow the real - // accumulators (e.g. `shuffleRecordsWritten`, which feeds AQE). - private lazy val customWriteMetrics: Map[String, SQLMetric] = { - val (reserved, allowed) = sparkContext.shuffleDriverComponents.supportedCustomMetrics() - .partition(metric => sparkMetrics.contains(metric.name())) - if (reserved.nonEmpty) { - logWarning(log"Ignoring custom shuffle metrics whose names collide with Spark-owned " + - log"Exchange metrics: " + - log"${MDC(METRIC_NAME, reserved.map(_.name()).sorted.mkString(", "))}.") - } - CustomShuffleMetrics.createMetrics(sparkContext, allowed) - } + private lazy val customWriteMetrics: Map[String, SQLMetric] = + CustomShuffleMetrics.createFilteredMetrics(sparkContext, sparkMetrics.keySet) override lazy val metrics = sparkMetrics ++ customWriteMetrics @@ -366,7 +355,7 @@ object ShuffleExchangeExec { newPartitioning: Partitioning, serializer: Serializer, writeMetrics: Map[String, SQLMetric], - customWriteMetrics: Map[String, SQLMetric] = Map.empty) + customWriteMetrics: Map[String, SQLMetric]) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -577,7 +566,7 @@ object ShuffleExchangeExec { */ def createShuffleWriteProcessor( metrics: Map[String, SQLMetric], - customWriteMetrics: Map[String, SQLMetric] = Map.empty): ShuffleWriteProcessor = { + customWriteMetrics: Map[String, SQLMetric]): ShuffleWriteProcessor = { new ShuffleWriteProcessor { override protected def createMetricsReporter( context: TaskContext): ShuffleWriteMetricsReporter = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala index c0fb1c37b210..5cde0311f818 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGe import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.catalyst.util.truncatedString import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec -import org.apache.spark.sql.execution.metric.{SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} +import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.execution.python.HybridRowQueue import org.apache.spark.util.collection.Utils @@ -38,6 +38,24 @@ trait LimitExec extends UnaryExecNode { def limit: Int } +/** + * Provides the shuffle read/write metric plumbing shared by the operators that collapse their + * input to a single partition via [[ShuffleExchangeExec.prepareShuffleDependency]]. Exposes the + * built-in read/write metrics plus collision-filtered plugin metrics as `customWriteMetrics`, and + * publishes all of them through `metrics`. Mixing this in wires the custom shuffle metrics in one + * place so a shuffle-writing operator can't silently drop them. + */ +trait ShuffleMetricsSupport { self: SparkPlan => + protected lazy val writeMetrics = + SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) + protected lazy val readMetrics = + SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) + private lazy val shuffleMetrics = readMetrics ++ writeMetrics + protected lazy val customWriteMetrics = + CustomShuffleMetrics.createFilteredMetrics(sparkContext, shuffleMetrics.keySet) + override lazy val metrics = shuffleMetrics ++ customWriteMetrics +} + /** * Take the first `limit` elements, collect them to a single partition and then to drop the * first `offset` elements. @@ -45,7 +63,8 @@ trait LimitExec extends UnaryExecNode { * This operator will be used when a logical `Limit` and/or `Offset` operation is the final operator * in an logical plan, which happens when the user is collecting results back to the driver. */ -case class CollectLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0) extends LimitExec { +case class CollectLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0) + extends LimitExec with ShuffleMetricsSupport { assert(limit >= 0 || (limit == -1 && offset > 0)) override def output: Seq[Attribute] = child.output @@ -62,11 +81,6 @@ case class CollectLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0) } } private val serializer: Serializer = new UnsafeRowSerializer(child.output.size) - private lazy val writeMetrics = - SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) - private lazy val readMetrics = - SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - override lazy val metrics = readMetrics ++ writeMetrics protected override def doExecute(): RDD[InternalRow] = { val childRDD = child.execute() if (childRDD.getNumPartitions == 0 || limit == 0) { @@ -86,7 +100,8 @@ case class CollectLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0) child.output, SinglePartition, serializer, - writeMetrics), + writeMetrics, + customWriteMetrics), readMetrics) } if (limit >= 0) { @@ -118,18 +133,14 @@ case class CollectLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0) * This operator will be used when a logical `Tail` operation is the final operator in an * logical plan, which happens when the user is collecting results back to the driver. */ -case class CollectTailExec(limit: Int, child: SparkPlan) extends LimitExec { +case class CollectTailExec(limit: Int, child: SparkPlan) + extends LimitExec with ShuffleMetricsSupport { assert(limit >= 0) override def output: Seq[Attribute] = child.output override def outputPartitioning: Partitioning = SinglePartition override def executeCollect(): Array[InternalRow] = child.executeTail(limit) private val serializer: Serializer = new UnsafeRowSerializer(child.output.size) - private lazy val writeMetrics = - SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) - private lazy val readMetrics = - SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - override lazy val metrics = readMetrics ++ writeMetrics protected override def doExecute(): RDD[InternalRow] = { val childRDD = child.execute() if (childRDD.getNumPartitions == 0 || limit == 0) { @@ -145,7 +156,8 @@ case class CollectTailExec(limit: Int, child: SparkPlan) extends LimitExec { child.output, SinglePartition, serializer, - writeMetrics), + writeMetrics, + customWriteMetrics), readMetrics) } singlePartitionRDD.mapPartitionsInternal(takeRight) @@ -312,7 +324,7 @@ case class TakeOrderedAndProjectExec( sortOrder: Seq[SortOrder], projectList: Seq[NamedExpression], child: SparkPlan, - offset: Int = 0) extends OrderPreservingUnaryExecNode { + offset: Int = 0) extends OrderPreservingUnaryExecNode with ShuffleMetricsSupport { override def output: Seq[Attribute] = { projectList.map(_.toAttribute) @@ -338,12 +350,6 @@ case class TakeOrderedAndProjectExec( private val serializer: Serializer = new UnsafeRowSerializer(child.output.size) - private lazy val writeMetrics = - SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) - private lazy val readMetrics = - SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - override lazy val metrics = readMetrics ++ writeMetrics - protected override def doExecute(): RDD[InternalRow] = { val orderingSatisfies = SortOrder.orderingSatisfies(child.outputOrdering, sortOrder) val ord = new LazilyGeneratedOrdering(sortOrder, child.output) @@ -368,7 +374,8 @@ case class TakeOrderedAndProjectExec( child.output, SinglePartition, serializer, - writeMetrics), + writeMetrics, + customWriteMetrics), readMetrics) } singlePartitionRDD.mapPartitionsWithIndexInternal { (idx, iter) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala index 3e8b59c93ecb..2699f9a2a9aa 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala @@ -18,10 +18,32 @@ package org.apache.spark.sql.execution.metric import org.apache.spark.SparkContext +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.METRIC_NAME import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} import org.apache.spark.util.MetricUtils -object CustomShuffleMetrics { +object CustomShuffleMetrics extends Logging { + + /** + * Creates collision-filtered custom shuffle metric [[SQLMetric]]s for an operator that already + * exposes the given Spark-owned metric names. Plugin metrics whose names collide with a + * Spark-owned name are dropped so they can't shadow the real accumulators (e.g. + * `shuffleRecordsWritten`, which feeds AQE). Operators pass the reserved names they own; the + * declared metrics come from the shuffle driver components. + */ + def createFilteredMetrics( + sc: SparkContext, + reservedNames: Set[String]): Map[String, SQLMetric] = { + val (reserved, allowed) = sc.shuffleDriverComponents.supportedCustomMetrics() + .partition(metric => reservedNames.contains(metric.name())) + if (reserved.nonEmpty) { + logWarning(log"Ignoring custom shuffle metrics whose names collide with Spark-owned " + + log"Exchange metrics: " + + log"${MDC(METRIC_NAME, reserved.map(_.name()).sorted.mkString(", "))}.") + } + createMetrics(sc, allowed) + } /** * Creates [[SQLMetric]]s for the given declared custom shuffle metrics, keyed by metric name. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala index bfbc0b6c4a7c..23878e9ec221 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala @@ -27,6 +27,7 @@ import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} import org.apache.spark.shuffle.sort.io.LocalDiskShuffleDataIO import org.apache.spark.sql.{LocalSparkSession, SparkSession} +import org.apache.spark.sql.execution.CollectLimitExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN @@ -87,6 +88,21 @@ class CustomShuffleMetricsIntegrationSuite assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value > 0) } } + + test("a declared custom shuffle metric surfaces on a multi-partition limit shuffle") { + withPluginSession(Some(classOf[TestCustomMetricShuffleDataIO].getName)) { spark => + import spark.implicits._ + val df = spark.range(0, 100, 1, 4).map(id => (id, id)).toDF("key", "value").limit(50) + val limitExec = collectFirst(df.queryExecution.executedPlan) { + case c: CollectLimitExec => c + }.getOrElse(fail("query plan had no CollectLimitExec")) + limitExec.execute().count() + assert(limitExec.metrics.contains(TestCustomMetricShuffleDataIO.MetricName), + "CollectLimitExec did not expose the declared custom shuffle metric") + assert(limitExec.metrics(TestCustomMetricShuffleDataIO.MetricName).value === + 4 * TestCustomMetricShuffleDataIO.PerTaskValue) + } + } } object TestCustomMetricShuffleDataIO { From e9922f588dd9e1fa54c7ef4326d3bf6be23961b9 Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Wed, 22 Jul 2026 23:21:45 -0700 Subject: [PATCH 4/6] Address review comments: binary+source compat, refactor --- .../sort/CustomMetricShuffleTestUtils.scala | 5 +- .../exchange/ShuffleExchangeExec.scala | 17 ++++++ ...CustomShuffleMetricsIntegrationSuite.scala | 53 ++++--------------- 3 files changed, 27 insertions(+), 48 deletions(-) diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala index a7a569fb5e30..3ff46a862d27 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/sort/CustomMetricShuffleTestUtils.scala @@ -35,10 +35,7 @@ private[spark] case class TestCustomShuffleTaskMetric(metricName: String, metric /** * Wraps a [[ShuffleExecutorComponents]] so every [[ShuffleMapOutputWriter]] it creates reports the - * given custom metric values from `currentMetricsValues()`, delegating all real writing to the - * wrapped components. Lets tests assert that the concrete - * [[org.apache.spark.shuffle.ShuffleWriter]] implementations stash and expose what the map output - * writer reports. + * given custom metric values from `currentMetricsValues()`, delegating all real writing. */ private[spark] class CustomMetricReportingExecutorComponents( delegate: ShuffleExecutorComponents, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index d63019d3af8d..155a78398ba8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -344,6 +344,18 @@ object ShuffleExchangeExec { } } + /** Backward-compatible overload for callers that do not report custom shuffle metrics. */ + def prepareShuffleDependency( + rdd: RDD[InternalRow], + outputAttributes: Seq[Attribute], + newPartitioning: Partitioning, + serializer: Serializer, + writeMetrics: Map[String, SQLMetric]) + : ShuffleDependency[Int, InternalRow, InternalRow] = { + prepareShuffleDependency( + rdd, outputAttributes, newPartitioning, serializer, writeMetrics, Map.empty) + } + /** * Returns a [[ShuffleDependency]] that will partition rows of its child based on * the partitioning scheme defined in `newPartitioning`. Those partitions of @@ -559,6 +571,11 @@ object ShuffleExchangeExec { dependency } + /** Backward-compatible overload for callers that do not report custom shuffle metrics. */ + def createShuffleWriteProcessor(metrics: Map[String, SQLMetric]): ShuffleWriteProcessor = { + createShuffleWriteProcessor(metrics, Map.empty) + } + /** * Create a customized [[ShuffleWriteProcessor]] for SQL which wrap the default metrics reporter * with [[SQLShuffleWriteMetricsReporter]] as new reporter for [[ShuffleWriteProcessor]], and diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala index 23878e9ec221..f466cde16dc6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala @@ -22,9 +22,9 @@ import java.util.{Map => JMap} import org.apache.spark.{SparkConf, SparkFunSuite} import org.apache.spark.internal.config.SHUFFLE_IO_PLUGIN_CLASS import org.apache.spark.internal.config.UI.UI_ENABLED -import org.apache.spark.shuffle.api.{ShuffleDataIO, ShuffleDriverComponents, ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter, SingleSpillShuffleMapOutputWriter} -import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage -import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} +import org.apache.spark.shuffle.api.{ShuffleDataIO, ShuffleDriverComponents, ShuffleExecutorComponents} +import org.apache.spark.shuffle.api.metric.CustomShuffleMetric +import org.apache.spark.shuffle.sort.{CustomMetricReportingExecutorComponents, TestCustomShuffleTaskMetric} import org.apache.spark.shuffle.sort.io.LocalDiskShuffleDataIO import org.apache.spark.sql.{LocalSparkSession, SparkSession} import org.apache.spark.sql.execution.CollectLimitExec @@ -35,10 +35,7 @@ import org.apache.spark.util.MetricUtils /** * Integration coverage for the shuffle custom-metrics SPI: a toy [[ShuffleDataIO]] declares a - * custom metric on its driver components and reports a fixed per-task value from every map output - * writer, and we assert the value surfaces on the [[ShuffleExchangeExec]] operator node after a - * real shuffle. This exercises the full path (driver declaration, task reporting, accumulator - * propagation, and the SQL-side bridge) without a `ThreadLocal` side-channel. + * custom metric and reports a fixed per-task value, and we assert it surfaces on the operator node. */ class CustomShuffleMetricsIntegrationSuite extends SparkFunSuite with LocalSparkSession with AdaptiveSparkPlanHelper { @@ -73,7 +70,7 @@ class CustomShuffleMetricsIntegrationSuite } } - test("the built-in local-disk plugin declares no custom metrics (regression guard)") { + test("no custom shuffle metrics surface when no plugin declares any") { withPluginSession(None) { spark => val exchange = runShuffleAndGetExchange(spark) assert(!exchange.metrics.contains(TestCustomMetricShuffleDataIO.MetricName)) @@ -121,7 +118,10 @@ class TestCustomMetricShuffleDataIO(sparkConf: SparkConf) extends ShuffleDataIO new TestCustomMetricDriverComponents(delegate.driver()) override def executor(): ShuffleExecutorComponents = - new TestCustomMetricExecutorComponents(delegate.executor()) + new CustomMetricReportingExecutorComponents( + delegate.executor(), + Array(TestCustomShuffleTaskMetric( + TestCustomMetricShuffleDataIO.MetricName, TestCustomMetricShuffleDataIO.PerTaskValue))) } class TestCustomMetricDriverComponents(delegate: ShuffleDriverComponents) @@ -139,41 +139,6 @@ class TestCustomMetricDriverComponents(delegate: ShuffleDriverComponents) }) } -class TestCustomMetricExecutorComponents(delegate: ShuffleExecutorComponents) - extends ShuffleExecutorComponents { - - override def initializeExecutor( - appId: String, execId: String, extraConfigs: JMap[String, String]): Unit = - delegate.initializeExecutor(appId, execId, extraConfigs) - - override def createMapOutputWriter( - shuffleId: Int, mapTaskId: Long, numPartitions: Int): ShuffleMapOutputWriter = - new TestCustomMetricMapOutputWriter( - delegate.createMapOutputWriter(shuffleId, mapTaskId, numPartitions)) - - override def createSingleFileMapOutputWriter( - shuffleId: Int, mapId: Long): java.util.Optional[SingleSpillShuffleMapOutputWriter] = - delegate.createSingleFileMapOutputWriter(shuffleId, mapId) -} - -class TestCustomMetricMapOutputWriter(delegate: ShuffleMapOutputWriter) - extends ShuffleMapOutputWriter { - - override def getPartitionWriter(reducePartitionId: Int): ShufflePartitionWriter = - delegate.getPartitionWriter(reducePartitionId) - - override def commitAllPartitions(checksums: Array[Long]): MapOutputCommitMessage = - delegate.commitAllPartitions(checksums) - - override def abort(error: Throwable): Unit = delegate.abort(error) - - override def currentMetricsValues(): Array[CustomShuffleTaskMetric] = Array( - new CustomShuffleTaskMetric { - override def name(): String = TestCustomMetricShuffleDataIO.MetricName - override def value(): Long = TestCustomMetricShuffleDataIO.PerTaskValue - }) -} - /** * A [[ShuffleDataIO]] that delegates all real work to the local-disk implementation but declares a * custom driver metric using a Spark-owned name (`shuffleRecordsWritten`). From 8d8aaa02225d429249cbd8b4ea87a750bd3a0f55 Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Fri, 24 Jul 2026 16:39:36 -0700 Subject: [PATCH 5/6] Address review comments: Drop unsupported metric instead failing --- .../org/apache/spark/internal/LogKeys.java | 1 + .../api/metric/CustomShuffleMetric.java | 3 +- .../metric/CustomShuffleMetrics.scala | 33 +++++++++++-------- .../metric/CustomShuffleMetricsSuite.scala | 17 +++++++--- 4 files changed, 35 insertions(+), 19 deletions(-) diff --git a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java index c4e1c5b96b1d..411576fd2a89 100644 --- a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java +++ b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java @@ -409,6 +409,7 @@ public enum LogKeys implements LogKey { METHOD_PARAM_TYPES, METRICS_JSON, METRIC_NAME, + METRIC_TYPE, MINI_BATCH_FRACTION, MIN_COMPACTION_BATCH_ID, MIN_NUM_FREQUENT_PATTERN, diff --git a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java index 9513c7b660d4..9c58e1ecfbb5 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java +++ b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java @@ -43,7 +43,8 @@ public interface CustomShuffleMetric { * 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. + * {@code "nsTiming"} (nanoseconds). A metric declaring any other value is dropped with a + * warning rather than failing the query. */ String metricType(); } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala index 2699f9a2a9aa..87a34be71fd4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala @@ -19,12 +19,18 @@ package org.apache.spark.sql.execution.metric import org.apache.spark.SparkContext import org.apache.spark.internal.Logging -import org.apache.spark.internal.LogKeys.METRIC_NAME +import org.apache.spark.internal.LogKeys.{METRIC_NAME, METRIC_TYPE} import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} import org.apache.spark.util.MetricUtils object CustomShuffleMetrics extends Logging { + private val SupportedMetricTypes = Seq( + MetricUtils.SUM_METRIC, + MetricUtils.SIZE_METRIC, + MetricUtils.TIMING_METRIC, + MetricUtils.NS_TIMING_METRIC).mkString(", ") + /** * Creates collision-filtered custom shuffle metric [[SQLMetric]]s for an operator that already * exposes the given Spark-owned metric names. Plugin metrics whose names collide with a @@ -47,25 +53,26 @@ object CustomShuffleMetrics extends Logging { /** * Creates [[SQLMetric]]s for the given declared custom shuffle metrics, keyed by metric name. + * A metric with an unsupported [[CustomShuffleMetric.metricType]] is dropped with a warning. */ def createMetrics( sc: SparkContext, metrics: Array[CustomShuffleMetric]): Map[String, SQLMetric] = { - metrics.map { metric => + metrics.flatMap { metric => val label = metric.description() - // AVERAGE_METRIC is intentionally unsupported: its per-task values must be stored pre-scaled - // via SQLMetric.set(Double), whereas custom task metrics report a plain Long. val acc = metric.metricType() match { - case MetricUtils.SUM_METRIC => SQLMetrics.createMetric(sc, label) - case MetricUtils.SIZE_METRIC => SQLMetrics.createSizeMetric(sc, label) - case MetricUtils.TIMING_METRIC => SQLMetrics.createTimingMetric(sc, label) - case MetricUtils.NS_TIMING_METRIC => SQLMetrics.createNanoTimingMetric(sc, label) - case other => throw new IllegalArgumentException( - s"Unsupported custom shuffle metric type '$other' for metric '${metric.name()}'. " + - s"Supported types: ${MetricUtils.SUM_METRIC}, ${MetricUtils.SIZE_METRIC}, " + - s"${MetricUtils.TIMING_METRIC}, ${MetricUtils.NS_TIMING_METRIC}.") + case MetricUtils.SUM_METRIC => Some(SQLMetrics.createMetric(sc, label)) + case MetricUtils.SIZE_METRIC => Some(SQLMetrics.createSizeMetric(sc, label)) + case MetricUtils.TIMING_METRIC => Some(SQLMetrics.createTimingMetric(sc, label)) + case MetricUtils.NS_TIMING_METRIC => Some(SQLMetrics.createNanoTimingMetric(sc, label)) + case other => + logWarning(log"Ignoring custom shuffle metric " + + log"${MDC(METRIC_NAME, metric.name())} with unsupported type " + + log"${MDC(METRIC_TYPE, other)}. Supported types: " + + log"${MDC(METRIC_TYPE, SupportedMetricTypes)}.") + None } - metric.name() -> acc + acc.map(metric.name() -> _) }.toMap } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala index e3ca50ad15e7..b05af6b09a86 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala @@ -36,13 +36,20 @@ class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext { assert(metrics("s3FirstByteLatency").metricType === MetricUtils.TIMING_METRIC) } - test("createMetrics rejects an unsupported metric type") { - val e = intercept[IllegalArgumentException] { - CustomShuffleMetrics.createMetrics(sc, Array( + test("createMetrics drops a metric with an unsupported type instead of failing") { + val logAppender = new LogAppender("unsupported custom shuffle metric type") + var metrics: Map[String, SQLMetric] = Map.empty + withLogAppender(logAppender) { + metrics = CustomShuffleMetrics.createMetrics(sc, Array( + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), shuffleMetric("s3Avg", "s3 avg", MetricUtils.AVERAGE_METRIC))) } - assert(e.getMessage.contains("s3Avg")) - assert(e.getMessage.contains(MetricUtils.AVERAGE_METRIC)) + + assert(metrics.keySet === Set("s3BytesUploaded")) + assert(logAppender.loggingEvents.exists { event => + val message = event.getMessage.getFormattedMessage + message.contains("s3Avg") && message.contains(MetricUtils.AVERAGE_METRIC) + }) } test("updateMetrics folds reported values into the matching SQLMetrics by name") { From 84e62548066ea09a7580df338be087fe9d6ee15f Mon Sep 17 00:00:00 2001 From: Karuppayya Rajendran Date: Wed, 29 Jul 2026 18:42:49 -0700 Subject: [PATCH 6/6] Address review comments --- .../org/apache/spark/internal/LogKeys.java | 1 - .../shuffle/api/ShuffleMapOutputWriter.java | 3 +- .../api/metric/CustomShuffleMetric.java | 19 +++++-- .../exchange/ShuffleExchangeExec.scala | 18 ++---- .../apache/spark/sql/execution/limit.scala | 19 ++++--- .../metric/CustomShuffleMetrics.scala | 29 +++------- ...CustomShuffleMetricsIntegrationSuite.scala | 6 +- .../metric/CustomShuffleMetricsSuite.scala | 55 +++++++++---------- 8 files changed, 65 insertions(+), 85 deletions(-) diff --git a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java index 411576fd2a89..c4e1c5b96b1d 100644 --- a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java +++ b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java @@ -409,7 +409,6 @@ public enum LogKeys implements LogKey { METHOD_PARAM_TYPES, METRICS_JSON, METRIC_NAME, - METRIC_TYPE, MINI_BATCH_FRACTION, MIN_COMPACTION_BATCH_ID, MIN_NUM_FREQUENT_PATTERN, diff --git a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java b/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java index 267cf4028496..7c0bbbad9047 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java +++ b/core/src/main/java/org/apache/spark/shuffle/api/ShuffleMapOutputWriter.java @@ -88,7 +88,8 @@ public interface ShuffleMapOutputWriter { void abort(Throwable error) throws IOException; /** - * The values of the custom shuffle metrics for this map task. + * The values of the custom shuffle metrics for this map task. Values are surfaced only when the + * {@link org.apache.spark.shuffle.ShuffleWriteProcessor} consumes them; otherwise discarded. */ default CustomShuffleTaskMetric[] currentMetricsValues() { return new CustomShuffleTaskMetric[0]; diff --git a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java index 9c58e1ecfbb5..ab7b5f28ee6b 100644 --- a/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java +++ b/core/src/main/java/org/apache/spark/shuffle/api/metric/CustomShuffleMetric.java @@ -28,6 +28,17 @@ @Private public interface CustomShuffleMetric { + enum MetricType { + /** Count rendered as a plain number. */ + SUM, + /** A byte size, rendered in human-readable units. */ + SIZE, + /** A duration in milliseconds. */ + TIMING, + /** A duration in nanoseconds. */ + NS_TIMING + } + /** * 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. @@ -40,11 +51,7 @@ public interface CustomShuffleMetric { 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). A metric declaring any other value is dropped with a - * warning rather than failing the query. + * Selects how the per-task values, once summed across tasks, are rendered in the UI. */ - String metricType(); + MetricType metricType(); } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index 155a78398ba8..5a815cb4af4c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} +import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLMetric, SQLMetrics, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.util.{MutablePair, ThreadUtils} import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator} @@ -193,22 +193,12 @@ case class ShuffleExchangeExec( child: SparkPlan, shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS, advisoryPartitionSize: Option[Long] = None) - extends ShuffleExchangeLike { + extends ShuffleExchangeLike with ShuffleMetricsSupport { - private lazy val writeMetrics = - SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) - private[sql] lazy val readMetrics = - SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - - private lazy val sparkMetrics: Map[String, SQLMetric] = Map( + override protected def extraReservedMetrics: Map[String, SQLMetric] = Map( "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"), "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions") - ) ++ readMetrics ++ writeMetrics - - private lazy val customWriteMetrics: Map[String, SQLMetric] = - CustomShuffleMetrics.createFilteredMetrics(sparkContext, sparkMetrics.keySet) - - override lazy val metrics = sparkMetrics ++ customWriteMetrics + ) override def nodeName: String = "Exchange" diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala index 5cde0311f818..451d11fd9055 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGe import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.catalyst.util.truncatedString import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec -import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} +import org.apache.spark.sql.execution.metric.{CustomShuffleMetrics, SQLMetric, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.execution.python.HybridRowQueue import org.apache.spark.util.collection.Utils @@ -39,21 +39,22 @@ trait LimitExec extends UnaryExecNode { } /** - * Provides the shuffle read/write metric plumbing shared by the operators that collapse their - * input to a single partition via [[ShuffleExchangeExec.prepareShuffleDependency]]. Exposes the - * built-in read/write metrics plus collision-filtered plugin metrics as `customWriteMetrics`, and - * publishes all of them through `metrics`. Mixing this in wires the custom shuffle metrics in one - * place so a shuffle-writing operator can't silently drop them. + * Shuffle read/write metric plumbing for operators that shuffle via + * [[ShuffleExchangeExec.prepareShuffleDependency]]. Publishes the built-in read/write metrics plus + * collision-filtered plugin metrics (`customWriteMetrics`) through `metrics`. Operators with + * additional Spark-owned metrics reserve those names against plugin collisions via + * [[extraReservedMetrics]]. */ trait ShuffleMetricsSupport { self: SparkPlan => protected lazy val writeMetrics = SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) protected lazy val readMetrics = SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) - private lazy val shuffleMetrics = readMetrics ++ writeMetrics + protected def extraReservedMetrics: Map[String, SQLMetric] = Map.empty + private lazy val reservedMetrics = extraReservedMetrics ++ readMetrics ++ writeMetrics protected lazy val customWriteMetrics = - CustomShuffleMetrics.createFilteredMetrics(sparkContext, shuffleMetrics.keySet) - override lazy val metrics = shuffleMetrics ++ customWriteMetrics + CustomShuffleMetrics.createFilteredMetrics(sparkContext, reservedMetrics.keySet) + override lazy val metrics = reservedMetrics ++ customWriteMetrics } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala index 87a34be71fd4..543aeca17001 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetrics.scala @@ -19,18 +19,11 @@ package org.apache.spark.sql.execution.metric import org.apache.spark.SparkContext import org.apache.spark.internal.Logging -import org.apache.spark.internal.LogKeys.{METRIC_NAME, METRIC_TYPE} +import org.apache.spark.internal.LogKeys.METRIC_NAME import org.apache.spark.shuffle.api.metric.{CustomShuffleMetric, CustomShuffleTaskMetric} -import org.apache.spark.util.MetricUtils object CustomShuffleMetrics extends Logging { - private val SupportedMetricTypes = Seq( - MetricUtils.SUM_METRIC, - MetricUtils.SIZE_METRIC, - MetricUtils.TIMING_METRIC, - MetricUtils.NS_TIMING_METRIC).mkString(", ") - /** * Creates collision-filtered custom shuffle metric [[SQLMetric]]s for an operator that already * exposes the given Spark-owned metric names. Plugin metrics whose names collide with a @@ -53,26 +46,20 @@ object CustomShuffleMetrics extends Logging { /** * Creates [[SQLMetric]]s for the given declared custom shuffle metrics, keyed by metric name. - * A metric with an unsupported [[CustomShuffleMetric.metricType]] is dropped with a warning. */ def createMetrics( sc: SparkContext, metrics: Array[CustomShuffleMetric]): Map[String, SQLMetric] = { - metrics.flatMap { metric => + metrics.map { metric => val label = metric.description() val acc = metric.metricType() match { - case MetricUtils.SUM_METRIC => Some(SQLMetrics.createMetric(sc, label)) - case MetricUtils.SIZE_METRIC => Some(SQLMetrics.createSizeMetric(sc, label)) - case MetricUtils.TIMING_METRIC => Some(SQLMetrics.createTimingMetric(sc, label)) - case MetricUtils.NS_TIMING_METRIC => Some(SQLMetrics.createNanoTimingMetric(sc, label)) - case other => - logWarning(log"Ignoring custom shuffle metric " + - log"${MDC(METRIC_NAME, metric.name())} with unsupported type " + - log"${MDC(METRIC_TYPE, other)}. Supported types: " + - log"${MDC(METRIC_TYPE, SupportedMetricTypes)}.") - None + case CustomShuffleMetric.MetricType.SUM => SQLMetrics.createMetric(sc, label) + case CustomShuffleMetric.MetricType.SIZE => SQLMetrics.createSizeMetric(sc, label) + case CustomShuffleMetric.MetricType.TIMING => SQLMetrics.createTimingMetric(sc, label) + case CustomShuffleMetric.MetricType.NS_TIMING => + SQLMetrics.createNanoTimingMetric(sc, label) } - acc.map(metric.name() -> _) + metric.name() -> acc }.toMap } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala index f466cde16dc6..2749a01b7a08 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsIntegrationSuite.scala @@ -31,7 +31,6 @@ import org.apache.spark.sql.execution.CollectLimitExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN -import org.apache.spark.util.MetricUtils /** * Integration coverage for the shuffle custom-metrics SPI: a toy [[ShuffleDataIO]] declares a @@ -135,7 +134,8 @@ class TestCustomMetricDriverComponents(delegate: ShuffleDriverComponents) new CustomShuffleMetric { override def name(): String = TestCustomMetricShuffleDataIO.MetricName override def description(): String = "test shuffle bytes uploaded" - override def metricType(): String = MetricUtils.SIZE_METRIC + override def metricType(): CustomShuffleMetric.MetricType = + CustomShuffleMetric.MetricType.SIZE }) } @@ -163,6 +163,6 @@ class TestCollidingMetricDriverComponents(delegate: ShuffleDriverComponents) new CustomShuffleMetric { override def name(): String = SHUFFLE_RECORDS_WRITTEN override def description(): String = "colliding shuffle records written" - override def metricType(): String = MetricUtils.SUM_METRIC + override def metricType(): CustomShuffleMetric.MetricType = CustomShuffleMetric.MetricType.SUM }) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala index b05af6b09a86..d1b6b36b7d5e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/CustomShuffleMetricsSuite.scala @@ -24,38 +24,26 @@ import org.apache.spark.util.MetricUtils class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext { - test("createMetrics maps each declared metric to a SQLMetric of the declared type") { + test("each metric type maps to a SQLMetric that renders the summed value accordingly") { val metrics = CustomShuffleMetrics.createMetrics(sc, Array( - shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), - shuffleMetric("s3BlockUploads", "s3 block uploads", MetricUtils.SUM_METRIC), - shuffleMetric("s3FirstByteLatency", "s3 first byte latency", MetricUtils.TIMING_METRIC))) + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", CustomShuffleMetric.MetricType.SIZE), + shuffleMetric("s3BlockUploads", "s3 block uploads", CustomShuffleMetric.MetricType.SUM), + shuffleMetric("s3Latency", "s3 first byte latency", CustomShuffleMetric.MetricType.TIMING), + shuffleMetric( + "s3LatencyNs", "s3 first byte latency ns", CustomShuffleMetric.MetricType.NS_TIMING))) - assert(metrics.keySet === Set("s3BytesUploaded", "s3BlockUploads", "s3FirstByteLatency")) + assert(metrics.keySet === + Set("s3BytesUploaded", "s3BlockUploads", "s3Latency", "s3LatencyNs")) assert(metrics("s3BytesUploaded").metricType === MetricUtils.SIZE_METRIC) assert(metrics("s3BlockUploads").metricType === MetricUtils.SUM_METRIC) - assert(metrics("s3FirstByteLatency").metricType === MetricUtils.TIMING_METRIC) - } - - test("createMetrics drops a metric with an unsupported type instead of failing") { - val logAppender = new LogAppender("unsupported custom shuffle metric type") - var metrics: Map[String, SQLMetric] = Map.empty - withLogAppender(logAppender) { - metrics = CustomShuffleMetrics.createMetrics(sc, Array( - shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), - shuffleMetric("s3Avg", "s3 avg", MetricUtils.AVERAGE_METRIC))) - } - - assert(metrics.keySet === Set("s3BytesUploaded")) - assert(logAppender.loggingEvents.exists { event => - val message = event.getMessage.getFormattedMessage - message.contains("s3Avg") && message.contains(MetricUtils.AVERAGE_METRIC) - }) + assert(metrics("s3Latency").metricType === MetricUtils.TIMING_METRIC) + assert(metrics("s3LatencyNs").metricType === MetricUtils.NS_TIMING_METRIC) } test("updateMetrics folds reported values into the matching SQLMetrics by name") { val metrics = CustomShuffleMetrics.createMetrics(sc, Array( - shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC), - shuffleMetric("s3BlockUploads", "s3 block uploads", MetricUtils.SUM_METRIC))) + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", CustomShuffleMetric.MetricType.SIZE), + shuffleMetric("s3BlockUploads", "s3 block uploads", CustomShuffleMetric.MetricType.SUM))) CustomShuffleMetrics.updateMetrics( Array(taskMetric("s3BytesUploaded", 1024L), taskMetric("s3BlockUploads", 3L)), metrics) @@ -66,7 +54,7 @@ class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext { test("updateMetrics ignores reported values with no matching declaration") { val metrics = CustomShuffleMetrics.createMetrics(sc, Array( - shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", MetricUtils.SIZE_METRIC))) + shuffleMetric("s3BytesUploaded", "s3 bytes uploaded", CustomShuffleMetric.MetricType.SIZE))) CustomShuffleMetrics.updateMetrics( Array(taskMetric("s3BytesUploaded", 512L), taskMetric("undeclared", 99L)), metrics) @@ -76,18 +64,25 @@ class CustomShuffleMetricsSuite extends SparkFunSuite with SharedSparkContext { } private def shuffleMetric( - metricName: String, desc: String, mType: String): CustomShuffleMetric = { + name: String, + description: String, + metricType: CustomShuffleMetric.MetricType): CustomShuffleMetric = { + val metricName = name + val metricDescription = description + val metricTypeValue = metricType new CustomShuffleMetric { override def name(): String = metricName - override def description(): String = desc - override def metricType(): String = mType + override def description(): String = metricDescription + override def metricType(): CustomShuffleMetric.MetricType = metricTypeValue } } - private def taskMetric(metricName: String, v: Long): CustomShuffleTaskMetric = { + private def taskMetric(name: String, value: Long): CustomShuffleTaskMetric = { + val metricName = name + val metricValue = value new CustomShuffleTaskMetric { override def name(): String = metricName - override def value(): Long = v + override def value(): Long = metricValue } } }