diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 6075dcc34dd..a0eb94e83b0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet +import java.util.IdentityHashMap + import scala.jdk.CollectionConverters._ import org.apache.spark.{SparkContext, TaskContext} @@ -61,6 +63,23 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM else children.flatMap(_.leafNodes) } + private[comet] def sumMetricValues(metricName: String): Long = { + val seenMetrics = new IdentityHashMap[SQLMetric, java.lang.Boolean]() + + def sumFromNode(metricNode: CometMetricNode): Long = { + val nodeValue = metricNode.metrics.get(metricName).fold(0L) { metric => + if (seenMetrics.put(metric, java.lang.Boolean.TRUE) == null) { + math.max(metric.value, 0L) + } else { + 0L + } + } + nodeValue + metricNode.children.iterator.map(sumFromNode).sum + } + + sumFromNode(this) + } + /** * Reports aggregated scan input metrics (bytesRead, recordsRead) to Spark's task metrics. * Aggregates across all scan leaf nodes to handle plans with multiple scans (e.g., joins). Must @@ -105,8 +124,8 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM } /** - * Reports this node's native shuffle spill metrics to Spark's task metrics, preserving the - * distinction between on-disk bytes and uncompressed in-memory bytes. + * Reports this node's and its descendants' native shuffle spill metrics to Spark's task + * metrics, preserving the distinction between on-disk bytes and uncompressed in-memory bytes. * * Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so its * completion listener publishes final SQL metrics before this listener runs, including when the @@ -114,17 +133,14 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM */ def reportSpillMetrics(ctx: TaskContext): Unit = { ctx.addTaskCompletionListener[Unit] { _ => - metrics.get("spilled_bytes").foreach { metric => - val spilledBytes = metric.value - if (spilledBytes > 0L) { - ctx.taskMetrics().incDiskBytesSpilled(spilledBytes) - } + val diskBytesSpilled = sumMetricValues("spilled_bytes") + if (diskBytesSpilled > 0L) { + ctx.taskMetrics().incDiskBytesSpilled(diskBytesSpilled) } - metrics.get("memory_spilled_bytes").foreach { metric => - val spilledBytes = metric.value - if (spilledBytes > 0L) { - ctx.taskMetrics().incMemoryBytesSpilled(spilledBytes) - } + + val memoryBytesSpilled = sumMetricValues("memory_spilled_bytes") + if (memoryBytesSpilled > 0L) { + ctx.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled) } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index b79fb9458c6..41d1d6cc6a4 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.RDD -import org.apache.spark.sql.comet.CometExecRDD +import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometShuffleBlockIterator @@ -39,6 +39,7 @@ private[shuffle] class CometNativeShuffleInputRDD( var inputRDDs: Seq[RDD[_]], numPartitionsParam: Int, shuffleScanIndices: Set[Int], + spillMetricNode: CometMetricNode, @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty) extends RDD[Product2[Int, ColumnarBatch]]( sc, @@ -63,6 +64,7 @@ private[shuffle] class CometNativeShuffleInputRDD( override def compute( split: Partition, context: TaskContext): Iterator[Product2[Int, ColumnarBatch]] = { + spillMetricNode.reportSpillMetrics(context) val partition = split.asInstanceOf[CometNativeShuffleInputPartition] val (inputObjects, shuffleBlockIters) = CometExecRDD.resolveInputObjects( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 3245cc3429b..7c64e963aa9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -137,8 +137,6 @@ class CometNativeShuffleWriter[K, V]( Option(context).foreach(nativeMetrics.reportScanInputMetrics) } - Option(context).foreach(nativeMetrics.reportSpillMetrics) - val cometIter = new CometExecIterator( CometExec.newIterId, inputObjects, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 84f313ad37c..2e12fb17b43 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -115,6 +115,9 @@ case class CometShuffleExchangeExec( case _ => None } + @transient private lazy val nativeChildMetricNode: CometMetricNode = + CometMetricNode.fromCometPlan(child) + @transient lazy val inputRDD: RDD[_] = if (shuffleType == CometNativeShuffle) { nativeChildContext match { case Some(ctx) => @@ -123,6 +126,7 @@ case class CometShuffleExchangeExec( ctx.inputs, ctx.numPartitions, ctx.shuffleScanIndices, + CometMetricNode(metrics, Seq(nativeChildMetricNode)), ctx.perPartitionByKey) case None => // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets @@ -189,10 +193,7 @@ case class CometShuffleExchangeExec( outputPartitioning, serializer, metrics, - NativeShuffleSpec( - nativeChild.nativeOp, - CometMetricNode.fromCometPlan(nativeChild), - ctx)) + NativeShuffleSpec(nativeChild.nativeOp, nativeChildMetricNode, ctx)) case None => CometShuffleExchangeExec.prepareShuffleDependency( inputRDD.asInstanceOf[RDD[ColumnarBatch]], @@ -717,11 +718,13 @@ object CometShuffleExchangeExec CometArrowStream.NATIVE_TIMEZONE, "ShuffleWriterInput") + val childMetricNode = CometMetricNode(Map.empty) val thinRDD = new CometNativeShuffleInputRDD( rdd.sparkContext, Seq(streamRDD), rdd.getNumPartitions, - shuffleScanIndices = Set.empty) + shuffleScanIndices = Set.empty, + spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode))) val ctx = NativeExecContext( inputs = Seq(streamRDD), @@ -743,7 +746,7 @@ object CometShuffleExchangeExec outputPartitioning, serializer, metrics, - NativeShuffleSpec(scanOp, CometMetricNode(Map.empty), ctx)) + NativeShuffleSpec(scanOp, childMetricNode, ctx)) } /** diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index ed94e3a696b..aa4c26396c8 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -35,6 +35,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -48,6 +49,35 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { import testImplicits._ + test("spill metric tree counts nested shared accumulators once") { + def metric(name: String, value: Long): SQLMetric = { + val sqlMetric = new SQLMetric(name) + sqlMetric.set(value) + sqlMetric + } + + val writerDisk = metric("writerDisk", 5L) + val writerMemory = metric("writerMemory", 3L) + val childDisk = metric("childDisk", 7L) + val nestedDisk = metric("nestedDisk", 11L) + val sharedDisk = metric("sharedDisk", 13L) + val sharedMemory = metric("sharedMemory", 17L) + val metricTree = CometMetricNode( + Map("spilled_bytes" -> writerDisk, "memory_spilled_bytes" -> writerMemory), + Seq( + CometMetricNode( + Map("spilled_bytes" -> childDisk), + Seq( + CometMetricNode( + Map("spilled_bytes" -> nestedDisk, "memory_spilled_bytes" -> sharedMemory)))), + CometMetricNode( + Map("spilled_bytes" -> sharedDisk, "memory_spilled_bytes" -> sharedMemory), + Seq(CometMetricNode(Map("spilled_bytes" -> sharedDisk)))))) + + assert(metricTree.sumMetricValues("spilled_bytes") == 36L) + assert(metricTree.sumMetricValues("memory_spilled_bytes") == 20L) + } + test("per-task native shuffle metrics") { withParquetTable((0 until 10000).map(i => (i, (i + 1).toLong)), "tbl") { val df = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc) @@ -160,6 +190,59 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("native shuffle task metrics include existing child sort spill metrics once") { + val expectedRecords = 20000L + val compressibleValue = "native-child-sort-spill-metrics-" * 8 + withParquetTable( + (0 until expectedRecords.toInt).map(index => (index, compressibleValue)), + "tbl") { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.key -> "zstd", + CometConf.COMET_SHUFFLE_NATIVE_MAX_BUFFER_BYTES.key -> "32k", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.spill_compression" -> "zstd", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536", + SQLConf.SHUFFLE_PARTITIONS.key -> "4") { + val shuffled = sql("SELECT * FROM tbl") + .sortWithinPartitions($"_1".desc) + .repartition(4, $"_1") + val store = spark.sparkContext.statusStore + spark.sparkContext.listenerBus.waitUntilEmpty() + val stagesBefore = store.stageList(null).map(_.stageId).toSet + + assert(shuffled.collect().length == expectedRecords) + spark.sparkContext.listenerBus.waitUntilEmpty() + + val exchange = collectFirst(shuffled.queryExecution.executedPlan) { + case native: CometShuffleExchangeExec if native.shuffleType == CometNativeShuffle => + native + }.getOrElse(fail("Expected a native shuffle exchange")) + val childSorts = collect(exchange.child) { case sort: CometSortExec => sort } + assert(childSorts.nonEmpty, s"Expected a native child sort:\n${exchange.treeString}") + + val writerDiskSpilled = exchange.metrics("spilled_bytes").value + val writerMemorySpilled = exchange.metrics("memory_spilled_bytes").value + val childDiskSpilled = childSorts.map(_.metrics("spilled_bytes").value).sum + assert(childDiskSpilled > 0L, "Native child sort did not spill") + assert(childSorts.forall(!_.metrics.contains("memory_spilled_bytes"))) + + val shuffleWriteStages = store + .stageList(null) + .filter(stage => + !stagesBefore.contains(stage.stageId) && stage.shuffleWriteRecords > 0L) + + assert(shuffleWriteStages.nonEmpty, "No native shuffle write stage was recorded") + assert( + shuffleWriteStages.map(_.diskBytesSpilled).sum == + writerDiskSpilled + childDiskSpilled) + assert(shuffleWriteStages.map(_.memoryBytesSpilled).sum == writerMemorySpilled) + } + } + } + test("failed native shuffle attempts preserve memory and disk spill metrics") { val failureRow = 8192 val compressibleValue = "native-shuffle-failed-spill-metrics-" * 8 diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index 55194e80a27..c3a19c8790b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -19,11 +19,12 @@ package org.apache.spark.sql.comet.execution.shuffle -import org.apache.spark.HashPartitioner +import org.apache.spark.{HashPartitioner, Partition, TaskContext} +import org.apache.spark.rdd.RDD import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} -import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.serde.OperatorOuterClass.Operator @@ -44,6 +45,59 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ class CometNativeShuffleInputRDDSuite extends CometTestBase { + test("spill reporting is registered before native shuffle input producers") { + Seq(None, Some(new IllegalStateException("failed native shuffle"))).foreach { failure => + val writerDisk = new SQLMetric("writerDisk") + val writerMemory = new SQLMetric("writerMemory") + val childDisk = new SQLMetric("childDisk") + val childMemory = new SQLMetric("childMemory") + val childMetrics = + CometMetricNode(Map("spilled_bytes" -> childDisk, "memory_spilled_bytes" -> childMemory)) + val taskContext = TaskContext.empty() + val nestedInput = new RDD[AnyRef](spark.sparkContext, Nil) { + override protected def getPartitions: Array[Partition] = Array(new Partition { + override def index: Int = 0 + }) + + override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] = { + context.addTaskCompletionListener[Unit] { _ => + childDisk.set(19L) + childMemory.set(37L) + } + Iterator.single(null) + } + } + val writerMetrics = + Map("spilled_bytes" -> writerDisk, "memory_spilled_bytes" -> writerMemory) + val inputRDD = new CometNativeShuffleInputRDD( + spark.sparkContext, + Seq(nestedInput), + 1, + Set.empty, + CometMetricNode(writerMetrics, Seq(childMetrics))) + + inputRDD.iterator(inputRDD.partitions.head, taskContext) + new CometNativeShuffleWriter[Int, Any]( + NativeShuffleSpec(null, childMetrics, null), + null, + Nil, + writerMetrics, + 1, + 0, + 0L, + taskContext, + null) + taskContext.addTaskCompletionListener[Unit] { _ => + writerDisk.set(23L) + writerMemory.set(41L) + } + taskContext.markTaskCompleted(failure) + + assert(taskContext.taskMetrics.diskBytesSpilled == 42L) + assert(taskContext.taskMetrics.memoryBytesSpilled == 78L) + } + } + test("serialized (rdd, dep) task binary size is independent of partition count") { val sc = spark.sparkContext val ser = new JavaSerializer(sc.getConf).newInstance() @@ -56,11 +110,16 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]) = { val perPartitionByKey = Map("scan-0" -> Array.fill(numPartitions)(new Array[Byte](1024))) + val childMetricNode = CometMetricNode(Map.empty) + val writerMetrics = Map( + "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "disk spilled bytes"), + "memory_spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "memory spilled bytes")) val rdd = new CometNativeShuffleInputRDD( sc, inputRDDs = Seq.empty, numPartitionsParam = numPartitions, shuffleScanIndices = Set.empty, + spillMetricNode = CometMetricNode(writerMetrics, Seq(childMetricNode)), perPartitionByKey = perPartitionByKey) val execContext = NativeExecContext( inputs = Seq.empty, @@ -72,12 +131,12 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { perPartitionByKey = perPartitionByKey, shuffleScanIndices = Set.empty, hasScanInput = false) - val spec = - NativeShuffleSpec(Operator.getDefaultInstance, CometMetricNode(Map.empty), execContext) + val spec = NativeShuffleSpec(Operator.getDefaultInstance, childMetricNode, execContext) val dep = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( _rdd = rdd, partitioner = new HashPartitioner(numPartitions), decodeTime = SQLMetrics.createMetric(sc, "decode time"), + shuffleWriteMetrics = writerMetrics, nativeShuffleSpec = Some(spec)) (rdd, dep) }