You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/27 20:25:44 UTC
(datafusion-comet) branch main updated: fix: report native operator spill metrics in Spark task metrics for non-shuffle stages (#5497)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new eabb5d4773 fix: report native operator spill metrics in Spark task metrics for non-shuffle stages (#5497)
eabb5d4773 is described below
commit eabb5d4773091b983d8fce713f0e34b1cf93f877
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Fri Aug 28 04:25:39 2026 +0800
fix: report native operator spill metrics in Spark task metrics for non-shuffle stages (#5497)
Spill-capable native operators running in non-shuffle stages (e.g. a sort
in a result stage executed via CometExecRDD) reported spilled_bytes only
as a SQL operator metric, so the Spark Stages/task view showed zero spill
for the task -- the same SQL-view vs task-view discrepancy that #5382
fixed for unified shuffle plans.
Register the task-level spill reporting callback in CometExecRDD.compute,
aggregating spilled_bytes / memory_spilled_bytes over the stage's
CometMetricNode tree as the shuffle path does: shared accumulators are
counted once, disk and memory totals stay separate, and no memory value
is inferred from disk bytes. The callback is registered before input
producers are resolved and before the native iterator is created, so
reverse-ordered task completion listeners publish final native metrics
first, including for failed attempts, matching the ordering established
in #5445.
Several native executions can run inside one Spark task with overlapping
metric trees: CometMetricNode.fromCometPlan descends through native input
boundaries, and a coalesced partition computes the same tree once per
parent partition. All reportSpillMetrics registrations in a task
(including the native shuffle writer path) now share a per-task identity
seen-set keyed by task attempt id, so each spill accumulator contributes
to the task totals exactly once while disjoint trees still contribute
separately.
Closes #5447.
Co-authored-by: Claude Fable 5 <no...@anthropic.com>
---
.../org/apache/spark/sql/comet/CometExecRDD.scala | 5 +
.../apache/spark/sql/comet/CometMetricNode.scala | 53 +++++-
.../spark/sql/comet/CometTaskMetricsSuite.scala | 199 +++++++++++++++++++++
3 files changed, 250 insertions(+), 7 deletions(-)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
index e8d2ef0006..d9e0bf3a4c 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala
@@ -98,6 +98,11 @@ private[spark] class CometExecRDD(
}
override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = {
+ // Must precede resolveInputObjects and the CometExecIterator: completion listeners run in
+ // reverse registration order, so registering first means this listener runs last, after
+ // nested native blocks and the iterator have published their final metric values.
+ Option(context).foreach(nativeMetrics.reportSpillMetrics)
+
val partition = split.asInstanceOf[CometExecPartition]
val (inputObjects, shuffleBlockIters) =
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 f67ff17c6d..c1c05a0e6e 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
@@ -20,6 +20,7 @@
package org.apache.spark.sql.comet
import java.util.IdentityHashMap
+import java.util.concurrent.ConcurrentHashMap
import scala.jdk.CollectionConverters._
@@ -63,9 +64,17 @@ 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]()
+ /**
+ * Sums `metricName` across this metric tree. A shared accumulator reachable through multiple
+ * paths in the tree contributes only once, and unset size metrics (initialized to -1) count as
+ * zero.
+ */
+ private[comet] def sumMetricValues(metricName: String): Long =
+ sumMetricValues(metricName, new IdentityHashMap[SQLMetric, java.lang.Boolean]())
+ private def sumMetricValues(
+ metricName: String,
+ seenMetrics: IdentityHashMap[SQLMetric, java.lang.Boolean]): Long = {
def sumFromNode(metricNode: CometMetricNode): Long = {
val nodeValue = metricNode.metrics.get(metricName).fold(0L) { metric =>
if (seenMetrics.put(metric, java.lang.Boolean.TRUE) == null) {
@@ -139,21 +148,27 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
}
/**
- * 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.
+ * Reports this node's and its descendants' native 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.
+ * attempt fails.
+ *
+ * A task may register several trees that share accumulators: an outer tree spans nested native
+ * blocks behind union/coalesce input boundaries, and a coalesced partition registers the same
+ * tree once per parent partition. All of a task's listeners claim accumulators from a shared
+ * per-task registry, so each accumulator is counted once while disjoint trees still all report.
*/
def reportSpillMetrics(ctx: TaskContext): Unit = {
+ val seenMetrics = CometMetricNode.taskSeenSpillMetrics(ctx)
ctx.addTaskCompletionListener[Unit] { _ =>
- val diskBytesSpilled = sumMetricValues("spilled_bytes")
+ val diskBytesSpilled = sumMetricValues("spilled_bytes", seenMetrics.disk)
if (diskBytesSpilled > 0L) {
ctx.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
}
- val memoryBytesSpilled = sumMetricValues("memory_spilled_bytes")
+ val memoryBytesSpilled = sumMetricValues("memory_spilled_bytes", seenMetrics.memory)
if (memoryBytesSpilled > 0L) {
ctx.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
}
@@ -210,6 +225,30 @@ object CometMetricNode {
private val aggregateMetricNames =
Set("spill_count", "spilled_bytes", "spilled_rows", "peak_mem_used")
+ private case class SeenSpillMetrics(
+ disk: IdentityHashMap[SQLMetric, java.lang.Boolean],
+ memory: IdentityHashMap[SQLMetric, java.lang.Boolean])
+
+ // Per running task attempt: the spill accumulators already claimed by a reporting listener,
+ // one identity set per metric name (see reportSpillMetrics). The first registration installs
+ // a cleanup listener ahead of every reporting listener, so it runs last (reverse registration
+ // order) and removes the entry.
+ private val seenSpillMetricsByTask = new ConcurrentHashMap[Long, SeenSpillMetrics]()
+
+ private def taskSeenSpillMetrics(ctx: TaskContext): SeenSpillMetrics = {
+ val attemptId = ctx.taskAttemptId()
+ val existing = seenSpillMetricsByTask.get(attemptId)
+ if (existing != null) {
+ existing
+ } else {
+ // The task thread is the only registrant for its attempt id, so there is no put race.
+ val created = SeenSpillMetrics(new IdentityHashMap(), new IdentityHashMap())
+ seenSpillMetricsByTask.put(attemptId, created)
+ ctx.addTaskCompletionListener[Unit](_ => seenSpillMetricsByTask.remove(attemptId))
+ created
+ }
+ }
+
/**
* The baseline SQL metrics for DataFusion `BaselineMetrics`.
*/
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 bd3123447e..33628954ae 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
@@ -39,6 +39,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.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{IntegerType, StructType}
@@ -84,6 +85,204 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper {
assert(metricTree.sumMetricValues("memory_spilled_bytes") == 20L)
}
+ test("overlapping spill trees registered on one task count each accumulator once") {
+ val nestedDisk = new SQLMetric("nestedDisk")
+ val nestedMemory = new SQLMetric("nestedMemory")
+ val outerDisk = new SQLMetric("outerDisk")
+ val siblingDisk = new SQLMetric("siblingDisk")
+ // An unset size metric keeps its -1 initial value and must count as zero.
+ val unsetDisk = new SQLMetric("unsetDisk", -1L)
+ val nestedTree =
+ CometMetricNode(Map("spilled_bytes" -> nestedDisk, "memory_spilled_bytes" -> nestedMemory))
+ // fromCometPlan produces this shape when an outer native block consumes a nested block
+ // through a union/coalesce input boundary: the outer tree contains the nested tree.
+ val outerTree = CometMetricNode(
+ Map("spilled_bytes" -> outerDisk),
+ Seq(nestedTree, CometMetricNode(Map("spilled_bytes" -> unsetDisk))))
+ val siblingTree = CometMetricNode(Map("spilled_bytes" -> siblingDisk))
+
+ // Reusing the same trees across both attempts verifies the shared per-task seen-set is
+ // discarded once a task completes; the failure attempt verifies failed tasks still report.
+ Seq(None, Some(new IllegalStateException("failed native stage"))).foreach { failure =>
+ val ctx = TaskContext.empty()
+ // Matches execution order: the outer block registers before resolving its inputs, which
+ // then register their own trees. The nested tree registers twice, like a coalesced
+ // partition computed once per parent partition.
+ outerTree.reportSpillMetrics(ctx)
+ nestedTree.reportSpillMetrics(ctx)
+ nestedTree.reportSpillMetrics(ctx)
+ siblingTree.reportSpillMetrics(ctx)
+ // Registered last so it runs first, like native iterators publishing final metric values
+ // as they close at task completion.
+ ctx.addTaskCompletionListener[Unit] { _ =>
+ outerDisk.set(5L)
+ nestedDisk.set(7L)
+ nestedMemory.set(11L)
+ siblingDisk.set(13L)
+ }
+ ctx.markTaskCompleted(failure)
+
+ // nestedDisk is claimed once even though three registered trees contain it, while the
+ // disjoint sibling tree still reports.
+ assert(ctx.taskMetrics.diskBytesSpilled == 25L)
+ assert(ctx.taskMetrics.memoryBytesSpilled == 11L)
+ }
+ }
+
+ test("native sort in a non-shuffle stage reports task-level disk spill metrics") {
+ val expectedRecords = 20000L
+ val compressibleValue = "non-shuffle-sort-spill-metrics-" * 8
+ withTempPath { path =>
+ // A single input partition keeps all rows in one task so its sort outgrows the tiny
+ // memory pool below and must spill.
+ spark
+ .createDataFrame((0 until expectedRecords.toInt).map(index => (index, compressibleValue)))
+ .coalesce(1)
+ .write
+ .parquet(path.getAbsolutePath)
+
+ withParquetTable(path.getAbsolutePath, "tbl") {
+ withSQLConf(
+ 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") {
+ val sorted = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc)
+ val store = spark.sparkContext.statusStore
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ val stagesBefore = store.stageList(null).map(_.stageId).toSet
+
+ assert(sorted.collect().length == expectedRecords)
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ val plan = sorted.queryExecution.executedPlan
+ assert(
+ collect(plan) { case exchange: ShuffleExchangeLike => exchange }.isEmpty,
+ s"Expected a shuffle-free plan so the spill happens in a result stage:\n$plan")
+ val sorts = collect(plan) { case sort: CometSortExec => sort }
+ assert(sorts.nonEmpty, s"Expected a native sort:\n$plan")
+ val sortDiskSpilled = sorts.map(_.metrics("spilled_bytes").value).sum
+ assert(sortDiskSpilled > 0L, "Native sort did not spill")
+ assert(sorts.forall(!_.metrics.contains("memory_spilled_bytes")))
+
+ val newStages = store
+ .stageList(null)
+ .filterNot(stage => stagesBefore.contains(stage.stageId))
+ assert(newStages.nonEmpty, "No stage was recorded for the non-shuffle sort")
+ assert(newStages.map(_.diskBytesSpilled).sum == sortDiskSpilled)
+ // The native sort exposes no memory-spill counter, and no memory value may be
+ // inferred from its disk bytes.
+ assert(newStages.map(_.memoryBytesSpilled).sum == 0L)
+ }
+ }
+ }
+ }
+
+ test("failed non-shuffle native stage attempts preserve disk spill metrics") {
+ val failureRow = 8192
+ val compressibleValue = "non-shuffle-failed-spill-metrics-" * 8
+
+ withTempPath { path =>
+ spark
+ .createDataFrame((0 until failureRow + 1024).map { i =>
+ (i, if (i == failureRow) 0 else 1, compressibleValue)
+ })
+ .coalesce(1)
+ .write
+ .parquet(path.getAbsolutePath)
+
+ withParquetTable(path.getAbsolutePath, "failed_sort_tbl") {
+ val failedTaskMetrics = mutable.ArrayBuffer.empty[(Long, Long)]
+ val targetStageIds = mutable.HashSet.empty[Int]
+ val jobGroupId =
+ s"failed-native-sort-spill-metrics-${java.util.UUID.randomUUID().toString}"
+
+ val listener = new SparkListener {
+ override def onJobStart(jobStart: SparkListenerJobStart): Unit = {
+ val isTargetJob = Option(jobStart.properties)
+ .flatMap(props => Option(props.getProperty(SparkContext.SPARK_JOB_GROUP_ID)))
+ .contains(jobGroupId)
+ if (isTargetJob) {
+ targetStageIds.synchronized {
+ targetStageIds ++= jobStart.stageInfos.map(_.stageId)
+ }
+ }
+ }
+
+ override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
+ val isTargetStage = targetStageIds.synchronized {
+ targetStageIds.contains(taskEnd.stageId)
+ }
+ if (isTargetStage && taskEnd.reason != Success) {
+ val taskMetrics = taskEnd.taskMetrics
+ failedTaskMetrics.synchronized {
+ failedTaskMetrics += ((
+ taskMetrics.memoryBytesSpilled,
+ taskMetrics.diskBytesSpilled))
+ }
+ }
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+
+ try {
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ withSQLConf(
+ 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.ANSI_ENABLED.key -> "true") {
+ // The sort consumes (and spills) its whole input before the projection above it
+ // reaches the failing division, so the failed attempt has real spills to report.
+ val failing = sql("SELECT * FROM failed_sort_tbl")
+ .sortWithinPartitions($"_1".desc)
+ .selectExpr("_1", "_1 / _2 AS quotient", "_3")
+ val plan = failing.queryExecution.executedPlan
+ assert(
+ collect(plan) { case exchange: ShuffleExchangeLike => exchange }.isEmpty,
+ s"Expected a shuffle-free plan so the failure happens in a result stage:\n$plan")
+ assert(
+ collect(plan) { case sort: CometSortExec => sort }.nonEmpty,
+ s"Expected a native sort below the failing projection:\n$plan")
+ assert(
+ collect(plan) { case project: CometProjectExec => project }.nonEmpty,
+ s"Expected the failing division to execute in a native project:\n$plan")
+
+ spark.sparkContext.setJobGroup(jobGroupId, "failed native result stage spill metrics")
+ try {
+ val failure = intercept[Exception] {
+ failing.collect()
+ }
+ val messages = causeChain(failure).flatMap(error => Option(error.getMessage))
+ assert(
+ messages.exists(message =>
+ message.contains("DIVIDE_BY_ZERO") || message.contains("Division by zero")),
+ s"Expected the late-row division failure, got:\n${messages.mkString("\n")}")
+ } finally {
+ spark.sparkContext.clearJobGroup()
+ }
+ }
+
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ val metrics = failedTaskMetrics.synchronized {
+ failedTaskMetrics.toSeq
+ }
+ assert(metrics.nonEmpty, "No failed native result task was recorded")
+ assert(
+ metrics.exists { case (_, diskBytesSpilled) => diskBytesSpilled > 0L },
+ s"Failed result-stage attempts must preserve disk spill metrics: $metrics")
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+ }
+ }
+
test("JVM shuffle peak memory includes a larger earlier spill") {
withSQLConf(CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key -> Int.MaxValue.toString) {
val pageSize = 256L
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org