From a34f2e6e2a940c06e6d30c90dfa54fc4dd92b787 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sat, 12 Sep 2026 20:29:05 +0700 Subject: [PATCH] fix: report task input metrics after the native iterator closes and add to Spark's counters Register the scan input metrics listener before super.compute creates the CometExecIterator at the native scan, Iceberg scan and CometExecRDD sites. Spark runs completion listeners in reverse registration order, so the report now runs after the iterator's close publishes the final native metrics. A block with a JVM input only publishes on the metrics update interval, and a consumer that stops early, such as a limit, left the report reading zero. Add to the task's bytesRead and recordsRead instead of replacing them, so bytes a fallback Spark scan or a cached input accumulated in the same task survive, and claim each accumulator once per task through the registry the spill reporter already uses. --- .../comet/CometIcebergNativeScanExec.scala | 5 +- .../spark/sql/comet/CometMetricNode.scala | 85 ++-- .../spark/sql/comet/CometNativeScanExec.scala | 8 +- .../apache/spark/sql/comet/operators.scala | 5 +- .../comet/CometIcebergNativeSuite.scala | 75 ++++ .../sql/comet/CometTaskMetricsSuite.scala | 404 ++++++++++++++++++ 6 files changed, 540 insertions(+), 42 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index 3182a9daeb4..bdd39ef76a3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -235,9 +235,10 @@ case class CometIcebergNativeScanExec( nativeMetrics = nativeMetrics, subqueries = Seq.empty) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { - val res = super.compute(split, context) + // Register before super.compute creates the CometExecIterator, so this listener runs + // after the iterator's close has published the final scan metrics. Option(context).foreach(nativeMetrics.reportScanInputMetrics) - res + super.compute(split, context) } } } 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 0308779ce07..cc345d2a921 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 @@ -77,11 +77,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM 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) { - math.max(metric.value, 0L) - } else { - 0L - } + CometMetricNode.claimMetricValue(metric, seenMetrics) } nodeValue + metricNode.children.iterator.map(sumFromNode).sum } @@ -105,24 +101,34 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM }) /** - * 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 - * be called in a TaskCompletionListener after the iterator is fully consumed. + * Reports the scan leaves' bytes and rows (summed across joins and unions) to Spark's task + * input metrics, which drive the Input column on the UI's Stages and Executors tabs. + * + * Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so its + * completion listener publishes final SQL metrics before this listener runs. A block with a JVM + * input only publishes on the metrics update interval, and a consumer that stops early, such as + * a limit, leaves the final publish to that close. + * + * Adds to the task's counters instead of replacing them, so bytes that a fallback Spark scan + * accumulated in the same task survive. Trees registered on one task may share accumulators + * (see [[reportSpillMetrics]]), so each accumulator is counted once per task. */ def reportScanInputMetrics(ctx: TaskContext): Unit = { + val seenMetrics = CometMetricNode.taskSeenMetrics(ctx).scanInput ctx.addTaskCompletionListener[Unit] { _ => val scanLeaves = leafNodes.filter(_.metrics.contains("bytes_scanned")) - if (scanLeaves.nonEmpty) { - val totalBytes = scanLeaves.map(_.metrics("bytes_scanned").value).sum - val totalRows = scanLeaves.map { leaf => - val outputRows = - leaf.metrics.get("output_rows").map(_.value).getOrElse(0L) - val prunedRows = - leaf.metrics.get("pushdown_rows_pruned").map(_.value).getOrElse(0L) - outputRows + prunedRows - }.sum - ctx.taskMetrics().inputMetrics.setBytesRead(totalBytes) - ctx.taskMetrics().inputMetrics.setRecordsRead(totalRows) + def claimed(leaf: CometMetricNode, metricName: String): Long = + leaf.metrics.get(metricName).fold(0L)(CometMetricNode.claimMetricValue(_, seenMetrics)) + + val totalBytes = scanLeaves.map(claimed(_, "bytes_scanned")).sum + val totalRows = scanLeaves.map { leaf => + claimed(leaf, "output_rows") + claimed(leaf, "pushdown_rows_pruned") + }.sum + if (totalBytes > 0L) { + ctx.taskMetrics().inputMetrics.incBytesRead(totalBytes) + } + if (totalRows > 0L) { + ctx.taskMetrics().inputMetrics.incRecordsRead(totalRows) } } } @@ -135,6 +141,10 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM * Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so * Spark's completion listener stack invokes the iterator `close` (final SQL metric update) * before this listener runs. + * + * The native writer is the root of its task's plan and runs once per task, so its values + * replace the task's output counters without the per-task registry that the scan input and + * spill reports share. */ def reportNativeWriteOutputMetrics(ctx: TaskContext): Unit = { ctx.addTaskCompletionListener[Unit] { _ => @@ -161,7 +171,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM * 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) + val seenMetrics = CometMetricNode.taskSeenMetrics(ctx) ctx.addTaskCompletionListener[Unit] { _ => val diskBytesSpilled = sumMetricValues("spilled_bytes", seenMetrics.disk) if (diskBytesSpilled > 0L) { @@ -225,30 +235,39 @@ 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]) + private type SeenMetricSet = IdentityHashMap[SQLMetric, java.lang.Boolean] + + private case class SeenMetrics( + disk: SeenMetricSet, + memory: SeenMetricSet, + scanInput: SeenMetricSet) - // 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]() + // Per running task attempt: the accumulators already claimed by a reporting listener, one + // identity set per reported task metric (see reportSpillMetrics and reportScanInputMetrics). + // 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 seenMetricsByTask = new ConcurrentHashMap[Long, SeenMetrics]() - private def taskSeenSpillMetrics(ctx: TaskContext): SeenSpillMetrics = { + private def taskSeenMetrics(ctx: TaskContext): SeenMetrics = { val attemptId = ctx.taskAttemptId() - val existing = seenSpillMetricsByTask.get(attemptId) + val existing = seenMetricsByTask.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)) + val created = + SeenMetrics(new IdentityHashMap(), new IdentityHashMap(), new IdentityHashMap()) + seenMetricsByTask.put(attemptId, created) + ctx.addTaskCompletionListener[Unit](_ => seenMetricsByTask.remove(attemptId)) created } } + /** The metric's value the first time it is claimed for a task, zero afterwards. */ + private def claimMetricValue(metric: SQLMetric, seenMetrics: SeenMetricSet): Long = + if (seenMetrics.put(metric, java.lang.Boolean.TRUE) == null) math.max(metric.value, 0L) + else 0L + /** * The baseline SQL metrics for DataFusion `BaselineMetrics`. */ diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 303e8279c3e..53b5e1e8373 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -277,12 +277,10 @@ case class CometNativeScanExec( encryptedFilePaths, perPartitionFilePaths = perPartitionFilePaths) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { - val res = super.compute(split, context) - - // Report scan input metrics after the iterator is fully consumed. + // Register before super.compute creates the CometExecIterator, so this listener runs + // after the iterator's close has published the final scan metrics. Option(context).foreach(nativeMetrics.reportScanInputMetrics) - - res + super.compute(split, context) } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3aa6b7220d3..149e444ca5d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -626,11 +626,12 @@ abstract class CometNativeExec extends CometExec { ctx.encryptedFilePaths, ctx.shuffleScanIndices) { override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { - val res = super.compute(split, context) + // Register before super.compute creates the CometExecIterator, so this listener runs + // after the iterator's close has published the final scan metrics. if (ctx.hasScanInput) { Option(context).foreach(nativeMetrics.reportScanInputMetrics) } - res + super.compute(split, context) } } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index d92723ecf83..b4eef71a0ca 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -3738,6 +3738,81 @@ class CometIcebergNativeSuite } } + test("Iceberg native scan left unconsumed by a limit still reports task input metrics") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + spark.sql(""" + CREATE TABLE test_cat.db.task_metrics_limit_test ( + id INT, + value DOUBLE + ) USING iceberg + """) + spark + .range(20000) + .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") + .repartition(4) + .write + .format("iceberg") + .mode("append") + .saveAsTable("test_cat.db.task_metrics_limit_test") + + val bytesReadValues = mutable.ArrayBuffer.empty[Long] + val recordsReadValues = mutable.ArrayBuffer.empty[Long] + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + val im = taskEnd.taskMetrics.inputMetrics + bytesReadValues.synchronized { + bytesReadValues += im.bytesRead + recordsReadValues += im.recordsRead + } + } + } + spark.sparkContext.addSparkListener(listener) + + try { + // The limit stops pulling before the scan is exhausted, so the final metric publish + // happens in the iterator's completion-time close. + val query = "SELECT * FROM test_cat.db.task_metrics_limit_test LIMIT 3" + Seq("-1", CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach { + interval => + withSQLConf(CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval) { + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + bytesReadValues.clear() + recordsReadValues.clear() + val df = spark.sql(query) + assert( + collectIcebergNativeScans(df.queryExecution.executedPlan).nonEmpty, + "Expected CometIcebergNativeScanExec in plan") + df.collect() + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + + val cometBytes = bytesReadValues.sum + val cometRecords = recordsReadValues.sum + assert( + cometBytes > 0, + s"bytesRead should be > 0 at interval $interval, got $cometBytes") + assert( + cometRecords >= 3 && cometRecords <= 20000, + s"recordsRead should cover at least the limit at interval $interval, got $cometRecords") + } + } + } finally { + spark.sparkContext.removeSparkListener(listener) + spark.sql("DROP TABLE test_cat.db.task_metrics_limit_test") + } + } + } + } + test("task-level inputMetrics.bytesRead is populated for Iceberg native scan") { assume(icebergAvailable, "Iceberg not available in classpath") 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 1c3500426c3..01513ec37f5 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 @@ -128,6 +128,87 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + // The unit tests below build their contexts with TaskContext.empty(), which always carries + // task attempt id 0, so they share one per-task registry entry. Each test must reach + // markTaskCompleted so the cleanup listener discards that entry before the next test claims. + test("scan input metrics add to the task once per accumulator and keep existing values") { + val nestedBytes = new SQLMetric("nestedBytes", -1L) + val nestedRows = new SQLMetric("nestedRows") + val nestedPruned = new SQLMetric("nestedPruned") + val siblingBytes = new SQLMetric("siblingBytes", -1L) + val siblingRows = new SQLMetric("siblingRows") + val nestedScan = CometMetricNode( + Map( + "bytes_scanned" -> nestedBytes, + "output_rows" -> nestedRows, + "pushdown_rows_pruned" -> nestedPruned)) + // An operator above the scans contributes output_rows of its own, which is not input. + val outerRows = new SQLMetric("outerRows") + outerRows.set(99L) + val outerTree = CometMetricNode( + Map("output_rows" -> outerRows), + Seq( + nestedScan, + CometMetricNode(Map("bytes_scanned" -> siblingBytes, "output_rows" -> siblingRows)))) + + Seq(None, Some(new IllegalStateException("failed native stage"))).foreach { failure => + val ctx = TaskContext.empty() + // A fallback Spark scan in the same task accumulated its own input before Comet reports. + ctx.taskMetrics.inputMetrics.incBytesRead(1000L) + ctx.taskMetrics.inputMetrics.incRecordsRead(3L) + // The outer block registers before its inputs, and the nested scan registers twice like + // a coalesced partition computed once per parent partition. + outerTree.reportScanInputMetrics(ctx) + nestedScan.reportScanInputMetrics(ctx) + nestedScan.reportScanInputMetrics(ctx) + // Registered last so it runs first, like native iterators publishing final metric values + // as they close at task completion. + ctx.addTaskCompletionListener[Unit] { _ => + nestedBytes.set(100L) + nestedRows.set(10L) + nestedPruned.set(5L) + siblingBytes.set(50L) + siblingRows.set(7L) + } + ctx.markTaskCompleted(failure) + + assert(ctx.taskMetrics.inputMetrics.bytesRead == 1150L) + assert(ctx.taskMetrics.inputMetrics.recordsRead == 25L) + } + } + + test("scan input metrics registered before the iterator see its final values") { + val bytes = new SQLMetric("bytes", -1L) + val rows = new SQLMetric("rows") + val scan = CometMetricNode(Map("bytes_scanned" -> bytes, "output_rows" -> rows)) + + val ctx = TaskContext.empty() + scan.reportScanInputMetrics(ctx) + // The iterator's close publishes the final values at task completion, after a polling + // block with a JVM input has left them stale, and runs before the report registered ahead. + ctx.addTaskCompletionListener[Unit] { _ => + bytes.set(64L) + rows.set(8L) + } + ctx.markTaskCompleted(None) + + assert(ctx.taskMetrics.inputMetrics.bytesRead == 64L) + assert(ctx.taskMetrics.inputMetrics.recordsRead == 8L) + } + + test("scan input metrics report nothing for scans that never ran") { + val scan = CometMetricNode( + Map("bytes_scanned" -> new SQLMetric("bytes", -1L), "output_rows" -> new SQLMetric("rows"))) + val ctx = TaskContext.empty() + ctx.taskMetrics.inputMetrics.incBytesRead(42L) + scan.reportScanInputMetrics(ctx) + ctx.markTaskCompleted(None) + + // An unset size metric keeps its -1 initial value and must not subtract from the task. + assert(ctx.taskMetrics.inputMetrics.bytesRead == 42L) + assert(ctx.taskMetrics.inputMetrics.recordsRead == 0L) + } + 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 @@ -1012,6 +1093,329 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("native scan left unconsumed by a limit still reports task input metrics") { + withTempPath { dir => + spark + .createDataFrame((0 until 20000).map(i => (i, s"big_$i"))) + .repartition(4) + .write + .parquet(dir.getAbsolutePath) + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("limit_big") + spark + .createDataFrame((0 until 100).map(i => (i * 7, s"local_$i"))) + .createOrReplaceTempView("limit_local") + + // The broadcast side is a JVM input to the join block, so native execution polls and only + // publishes scan metrics on the update interval. The limit stops pulling before the scan + // is exhausted, leaving the final publish to the iterator's completion-time close. + val query = "SELECT /*+ BROADCAST(limit_local) */ * FROM limit_big JOIN limit_local " + + "ON limit_big._1 = limit_local._1 LIMIT 3" + val localBroadcast = Seq( + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "LocalTableScan") + + Seq("-1", CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach { interval => + val confs = localBroadcast :+ (CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval) + val (cometBytes, cometRecords, cometPlan) = + collectInputMetrics(query, (CometConf.COMET_ENABLED.key -> "true") +: confs: _*) + + val join = find(cometPlan)(_.isInstanceOf[CometBroadcastHashJoinExec]) + assert( + join.isDefined, + s"Expected CometBroadcastHashJoinExec in plan:\n${cometPlan.treeString}") + val streamedNativeScan = join.get.children.exists { child => + find(child)(_.isInstanceOf[CometBroadcastExchangeExec]).isEmpty && + find(child)(_.isInstanceOf[CometNativeScanExec]).isDefined + } + assert( + streamedNativeScan, + s"Expected a native scan on the streamed side of the join:\n${cometPlan.treeString}") + + assert(cometBytes > 0, s"bytesRead should be > 0 at interval $interval, got $cometBytes") + assert( + cometRecords >= 3 && cometRecords <= 20000, + s"recordsRead should cover at least the limit at interval $interval, got $cometRecords") + } + } + } + + test("task input metrics keep bytes read by a fallback Spark scan in the same task") { + withTempPath { parquetDir => + withTempPath { jsonDir => + spark + .createDataFrame((0 until 5000).map(i => (i, s"parquet_$i"))) + .repartition(1) + .write + .parquet(parquetDir.getAbsolutePath) + spark + .createDataFrame((5000 until 10000).map(i => (i, s"json_$i"))) + .repartition(1) + .write + .json(jsonDir.getAbsolutePath) + spark.read.parquet(parquetDir.getAbsolutePath).createOrReplaceTempView("mixed_parquet") + spark.read.json(jsonDir.getAbsolutePath).createOrReplaceTempView("mixed_json") + val convertJson = CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "true" + + val nativeArm = "SELECT _1 FROM mixed_parquet" + val fallbackArm = "SELECT CAST(_1 AS INT) FROM mixed_json" + val nativeBytes = cometBytesRead(nativeArm, convertJson) + val fallbackBytes = cometBytesRead(fallbackArm, convertJson) + assert(nativeBytes > 0 && fallbackBytes > 0, s"sides: $nativeBytes, $fallbackBytes") + + // Coalescing the union to one partition computes the native scan and the fallback JSON + // scan inside the same task, so Spark's own input metrics and Comet's must add up. + def coalescedUnion(first: String, second: String): String = + s"SELECT /*+ COALESCE(1) */ * FROM ($first UNION ALL $second)" + + val nativeFirst = coalescedUnion(nativeArm, fallbackArm) + val (sparkBytes, sparkRecords, _) = + collectInputMetrics(nativeFirst, CometConf.COMET_ENABLED.key -> "false", convertJson) + val (cometBytes, cometRecords, cometPlan) = + collectInputMetrics(nativeFirst, CometConf.COMET_ENABLED.key -> "true", convertJson) + assert( + find(cometPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined, + s"Expected CometNativeScanExec in plan:\n${cometPlan.treeString}") + assert( + find(cometPlan)(_.isInstanceOf[CometSparkToColumnarExec]).isDefined, + s"Expected CometSparkToColumnarExec in plan:\n${cometPlan.treeString}") + assert(sparkRecords > 0, s"Spark recordsRead should be > 0, got $sparkRecords") + assert( + cometRecords == sparkRecords, + s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords") + assert(sparkBytes > 0, s"Spark bytesRead should be > 0, got $sparkBytes") + assert( + cometBytes >= nativeBytes + fallbackBytes, + s"bytesRead should cover both sides: comet=$cometBytes, " + + s"native=$nativeBytes, fallback=$fallbackBytes") + + // With the fallback partition first, FileScanRDD registers its close listener before + // Comet's report, so it runs last and sets bytesRead to the bytes it snapshotted at + // construction plus its own reads, dropping the native side. Records still add up, + // since FileScanRDD increments those. This pins the current behaviour. + val fallbackFirst = coalescedUnion(fallbackArm, nativeArm) + val (_, sparkRecordsReversed, _) = + collectInputMetrics(fallbackFirst, CometConf.COMET_ENABLED.key -> "false", convertJson) + val (cometBytesReversed, cometRecordsReversed, _) = + collectInputMetrics(fallbackFirst, CometConf.COMET_ENABLED.key -> "true", convertJson) + assert( + cometRecordsReversed == sparkRecordsReversed, + s"recordsRead mismatch: comet=$cometRecordsReversed, spark=$sparkRecordsReversed") + assert( + cometBytesReversed >= fallbackBytes && cometBytesReversed < nativeBytes + fallbackBytes, + s"bytesRead with the fallback side first keeps only its own bytes today: " + + s"comet=$cometBytesReversed, native=$nativeBytes, fallback=$fallbackBytes") + } + } + } + + test("native scan block left unconsumed by a limit still reports task input metrics") { + withTempPath { dir => + spark + .createDataFrame((0 until 20000).map(i => (i, s"elem_$i"))) + .repartition(4) + .write + .parquet(dir.getAbsolutePath) + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("scan_limit_tbl") + + // The scan is its own native block here, so CometNativeScanExec registers the report. + val query = "SELECT * FROM scan_limit_tbl LIMIT 3" + Seq("-1", CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach { interval => + val (cometBytes, cometRecords, cometPlan) = collectInputMetrics( + query, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval) + assert( + find(cometPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined, + s"Expected CometNativeScanExec in plan:\n${cometPlan.treeString}") + assert(cometBytes > 0, s"bytesRead should be > 0 at interval $interval, got $cometBytes") + assert( + cometRecords >= 3 && cometRecords <= 20000, + s"recordsRead should cover at least the limit at interval $interval, got $cometRecords") + } + } + } + + test("coalesced union of two native scans reports task input metrics for both") { + withTempPath { dir1 => + withTempPath { dir2 => + spark + .createDataFrame((0 until 5000).map(i => (i, s"left_$i"))) + .repartition(1) + .write + .parquet(dir1.getAbsolutePath) + spark + .createDataFrame((5000 until 10000).map(i => (i, s"right_$i"))) + .repartition(1) + .write + .parquet(dir2.getAbsolutePath) + spark.read.parquet(dir1.getAbsolutePath).createOrReplaceTempView("co_union_left") + spark.read.parquet(dir2.getAbsolutePath).createOrReplaceTempView("co_union_right") + + // Both scan blocks run inside one task and register their own report, so each block's + // bytes and rows must add up rather than the last report replacing the first. + val query = "SELECT /*+ COALESCE(1) */ * FROM (SELECT * FROM co_union_left " + + "UNION ALL SELECT * FROM co_union_right)" + val leftBytes = cometBytesRead("SELECT * FROM co_union_left") + val rightBytes = cometBytesRead("SELECT * FROM co_union_right") + + val (sparkBytes, sparkRecords, _) = + collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "false") + val (cometBytes, cometRecords, cometPlan) = + collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "true") + + val scanCount = collect(cometPlan) { case s: CometNativeScanExec => s }.size + assert( + scanCount == 2, + s"Expected 2 CometNativeScanExec in plan:\n${cometPlan.treeString}") + assert( + find(cometPlan)(_.isInstanceOf[CometCoalesceExec]).isDefined, + s"Expected CometCoalesceExec in plan:\n${cometPlan.treeString}") + assert( + cometRecords == sparkRecords, + s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords") + assert(sparkBytes > 0, s"Spark bytesRead should be > 0, got $sparkBytes") + assertCometBytesReadInRange(cometBytes, sparkBytes) + assert( + cometBytes >= leftBytes + rightBytes, + s"bytesRead should cover both blocks: comet=$cometBytes, " + + s"left=$leftBytes, right=$rightBytes") + } + } + } + + test("task input metrics keep rows read from a cached JVM input in the same task") { + withTempPath { dir => + spark + .createDataFrame((0 until 5000).map(i => (i, s"parquet_$i"))) + .repartition(1) + .write + .parquet(dir.getAbsolutePath) + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("cache_union_parquet") + val cached = spark.createDataFrame((5000 until 10000).map(i => (i, s"cached_$i"))).cache() + try { + cached.count() + cached.createOrReplaceTempView("cache_union_cached") + + // Reading a cached block increments the task's input metrics on the JVM side, in the + // same task as the native scan once the union is coalesced. + val query = "SELECT /*+ COALESCE(1) */ * FROM (SELECT * FROM cache_union_parquet " + + "UNION ALL SELECT * FROM cache_union_cached)" + val parquetBytes = cometBytesRead("SELECT * FROM cache_union_parquet") + val cachedBytes = cometBytesRead("SELECT * FROM cache_union_cached") + assert(parquetBytes > 0 && cachedBytes > 0, s"sides: $parquetBytes, $cachedBytes") + + val (_, sparkRecords, _) = + collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "false") + val (cometBytes, cometRecords, cometPlan) = + collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "true") + + assert( + find(cometPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined, + s"Expected CometNativeScanExec in plan:\n${cometPlan.treeString}") + // Spark counts a cached block as one record, so the baseline is the parquet rows plus + // the cached blocks rather than a row count. + assert( + sparkRecords > 5000, + s"Spark recordsRead should exceed the parquet rows, got $sparkRecords") + assert( + cometRecords == sparkRecords, + s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords") + assert( + cometBytes >= parquetBytes + cachedBytes, + s"bytesRead should cover both sides: comet=$cometBytes, " + + s"parquet=$parquetBytes, cached=$cachedBytes") + } finally { + cached.unpersist() + } + } + } + + test("failed native stage attempts still report task input metrics") { + val failureRow = 8192 + withTempPath { path => + spark + .createDataFrame((0 until failureRow + 1024).map(i => (i, if (i == failureRow) 0 else 1))) + .coalesce(1) + .write + .parquet(path.getAbsolutePath) + + withParquetTable(path.getAbsolutePath, "failed_scan_tbl") { + val failedTaskMetrics = mutable.ArrayBuffer.empty[(Long, Long)] + val targetStageIds = mutable.HashSet.empty[Int] + val jobGroupId = s"failed-native-scan-input-metrics-${java.util.UUID.randomUUID()}" + + 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 inputMetrics = taskEnd.taskMetrics.inputMetrics + failedTaskMetrics.synchronized { + failedTaskMetrics += ((inputMetrics.bytesRead, inputMetrics.recordsRead)) + } + } + } + } + spark.sparkContext.addSparkListener(listener) + + try { + spark.sparkContext.listenerBus.waitUntilEmpty() + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "1024", + SQLConf.ANSI_ENABLED.key -> "true") { + // The scan reads several batches before the projection reaches the failing row. + val failing = sql("SELECT _1 / _2 AS quotient FROM failed_scan_tbl") + val plan = failing.queryExecution.executedPlan + assert( + collect(plan) { case scan: CometNativeScanExec => scan }.nonEmpty, + s"Expected a native scan below the failing projection:\n$plan") + + spark.sparkContext.setJobGroup(jobGroupId, "failed native scan input 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 task attempt was observed") + assert( + metrics.forall { case (bytesRead, recordsRead) => + bytesRead > 0L && recordsRead > 0L + }, + s"Failed attempts should report the bytes and rows scanned before failing: $metrics") + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + } + } + + /** The task input bytes a query reports on its own with Comet enabled. */ + private def cometBytesRead(query: String, confs: (String, String)*): Long = + collectInputMetrics(query, (CometConf.COMET_ENABLED.key -> "true") +: confs: _*)._1 + /** * Runs the given query with the given SQL config overrides and returns the aggregated * (bytesRead, recordsRead) across all tasks, along with the executed plan.