fix: report native child spill metrics in shuffle tasks

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,26 +124,23 @@ 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
    * shuffle attempt fails.
    */
   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)
       }
     }
   }

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(

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,

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))
   }
 
   /**

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

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)
     }