From ed23e2d092a05c7a108c6892269794aea4980af6 Mon Sep 17 00:00:00 2001 From: huanghsiang_cheng Date: Tue, 4 Aug 2026 12:50:42 -0700 Subject: [PATCH 1/5] Populate executor's input metrics on fused operators --- .../apache/spark/sql/comet/operators.scala | 10 +- .../comet/CometIcebergNativeSuite.scala | 107 ++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) 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 8e1efa00872..1ce8d508a99 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 @@ -822,7 +822,15 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec])) + // Any leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`, + // `CometCsvNativeScanExec`, contrib scans) can contribute `bytes_scanned` / + // `output_rows` to Spark's task-level input metrics, which drive the Input column + // on the UI's Stages and Executors tabs. Matching on `CometLeafExec` rather than + // `CometNativeScanExec` keeps every scan reported once the scan is fused into a + // larger native block, where only the block root's `compute` runs. + // `reportScanInputMetrics` self-filters on the `bytes_scanned` metric, so leaves + // that don't track it are a no-op. + hasScanInput = sparkPlans.exists(_.isInstanceOf[CometLeafExec])) } /** diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index c8162559248..1259c1e237e 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -3659,6 +3659,113 @@ class CometIcebergNativeSuite } } + /** + * Companion to "task-level inputMetrics.bytesRead is populated for Iceberg native scan", which + * runs `SELECT *` and therefore leaves the scan as the root of its native block. In that shape + * Spark calls `CometIcebergNativeScanExec.doExecuteColumnar` directly, and that method reports + * input metrics itself. + * + * This test covers the other reporting site. With an operator above the scan, the two fuse into + * one native block and only the block root's `compute` runs -- a parent `CometNativeExec` reads + * its scan children via `PlanDataInjector.findAllPlanData` instead of executing them. Reporting + * then comes from `CometNativeExec.executeColumnarWithContext`, whose `hasScanInput` gate used + * to match only `CometNativeScanExec` and so skipped Iceberg, leaving the Input column on the + * UI's Stages and Executors tabs blank. + */ + test("task-level inputMetrics is populated when Iceberg native scan is fused into a block") { + 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.fused_metrics_test ( + id INT, + value DOUBLE + ) USING iceberg + """) + + spark + .range(10000) + .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") + .repartition(5) + .write + .format("iceberg") + .mode("append") + .saveAsTable("test_cat.db.fused_metrics_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 + if (im.bytesRead > 0) { + bytesReadValues.synchronized { + bytesReadValues += im.bytesRead + recordsReadValues += im.recordsRead + } + } + } + } + spark.sparkContext.addSparkListener(listener) + + try { + // Arithmetic in the projection keeps it from being collapsed into the scan, so a + // CometProjectExec sits above the scan and the two fuse into one native block. + val query = + "SELECT id + 1 AS id2, value * 2 AS value2 FROM test_cat.db.fused_metrics_test" + + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + bytesReadValues.clear() + recordsReadValues.clear() + + val df = spark.sql(query) + val plan = df.queryExecution.executedPlan + val scanNodes = collectIcebergNativeScans(plan) + assert(scanNodes.nonEmpty, s"Expected CometIcebergNativeScanExec in plan:\n$plan") + // Assert the fusion actually happened, so this test cannot silently degrade into a + // duplicate of the SELECT * one if a future planner change collapses the projection + // into the scan. Checking the executedPlan root is not enough: with AQE on, the root + // is an AdaptiveSparkPlanExec either way. + val fusingParents = collect(plan) { + case p: CometNativeExec + if p.children.exists(_.isInstanceOf[CometIcebergNativeScanExec]) => + p + } + assert( + fusingParents.nonEmpty, + s"Expected the Iceberg scan to be fused under a native parent operator:\n$plan") + + df.collect() + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + + val cometBytes = bytesReadValues.sum + val cometRecords = recordsReadValues.sum + + assert(cometBytes > 0, s"bytesRead should be > 0 for a fused scan, got $cometBytes") + assert( + cometRecords == 10000, + s"recordsRead should equal the scanned row count, got $cometRecords") + + val sqlBytes = scanNodes.map(_.metrics("bytes_scanned").value).sum + assert( + sqlBytes == cometBytes, + s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + } finally { + spark.sparkContext.removeSparkListener(listener) + spark.sql("DROP TABLE test_cat.db.fused_metrics_test") + } + } + } + } + test("exchange reuse must not collapse scans with different pushed filters (#4774)") { assume(icebergAvailable, "Iceberg not available") From e103ba1efe7e8eec4f34378adad36170261f9301 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 6 Aug 2026 14:50:55 -0700 Subject: [PATCH 2/5] Test Shuffle Writer metrics --- .../comet/CometIcebergNativeSuite.scala | 127 +++++++++++++++++- 1 file changed, 126 insertions(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 1259c1e237e..43d96670aef 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -31,10 +31,11 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression import org.apache.spark.sql.comet._ -import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{InSubqueryExec, ReusedSubqueryExec, SparkPlan, SubqueryExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, BroadcastQueryStageExec} import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{StringType, TimestampType} @@ -3766,6 +3767,130 @@ class CometIcebergNativeSuite } } + /** + * `CometNativeExec.buildNativeContext` is also consumed by the native shuffle path: + * `CometShuffleExchangeExec.nativeChildContext` builds the context from its `CometNativeExec` + * child, and the child's whole native subtree -- scan included -- is inlined under the + * `ShuffleWriter` protobuf operator and executed by `CometNativeShuffleWriter` in the + * ShuffleMapTask. No `CometExecRDD` runs for that subtree, so neither + * `CometIcebergNativeScanExec.doExecuteColumnar` nor + * `CometNativeExec.executeColumnarWithContext` reports anything; the writer has its own + * `ctx.hasScanInput` check instead. + * + * Before the fix, an Iceberg scan feeding a native shuffle left the map stage's Input column + * blank. The Parquet equivalent ("native shuffle reports task input metrics for its scan child" + * in `CometTaskMetricsSuite`) passed all along because the old gate matched + * `CometNativeScanExec`. + */ + test("task-level inputMetrics is populated when Iceberg native scan feeds a native shuffle") { + 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", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + // "auto" would also pick native here, but pin it so a future change to the auto + // heuristic turns this into a skip-with-assertion-failure rather than a silent + // switch to columnar shuffle (which reports input metrics through a different path). + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + + spark.sql(""" + CREATE TABLE test_cat.db.shuffle_metrics_test ( + id INT, + value DOUBLE + ) USING iceberg + """) + + spark + .range(10000) + .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") + .repartition(5) + .write + .format("iceberg") + .mode("append") + .saveAsTable("test_cat.db.shuffle_metrics_test") + + val mapInputBytes = mutable.ArrayBuffer.empty[Long] + val mapInputRecords = mutable.ArrayBuffer.empty[Long] + + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + // Only the map stage runs the scan; the reduce stage reads shuffle blocks and must + // not contribute to inputMetrics. + if (taskEnd.taskType.contains("ShuffleMapTask")) { + val im = taskEnd.taskMetrics.inputMetrics + mapInputBytes.synchronized { + mapInputBytes += im.bytesRead + mapInputRecords += im.recordsRead + } + } + } + } + spark.sparkContext.addSparkListener(listener) + + try { + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + mapInputBytes.clear() + mapInputRecords.clear() + + val df = spark + .table("test_cat.db.shuffle_metrics_test") + .repartition(4, col("id")) + + df.collect() + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + + // Assert the plan shape after execution so we inspect what AQE actually ran. All three + // conditions are required for the writer to be the reporting site: a native (not + // columnar) shuffle, a CometNativeExec child so `nativeChildContext` is `Some`, and an + // Iceberg scan inside that child's subtree so `hasScanInput` must be true. + val plan = df.queryExecution.executedPlan + val nativeShuffles = collect(plan) { + case s: CometShuffleExchangeExec if s.shuffleType == CometNativeShuffle => s + } + assert( + nativeShuffles.nonEmpty, + s"Expected a CometShuffleExchangeExec with CometNativeShuffle in plan:\n$plan") + val scanNodes = nativeShuffles.flatMap { s => + assert( + s.child.isInstanceOf[CometNativeExec], + s"Expected the shuffle's child to be a CometNativeExec so its subtree is " + + s"inlined into the writer plan, got ${s.child.getClass.getSimpleName}:\n$plan") + collectIcebergNativeScans(s.child) + } + assert( + scanNodes.nonEmpty, + s"Expected the Iceberg scan to be inlined under the native shuffle:\n$plan") + + assert(mapInputRecords.nonEmpty, "no ShuffleMapTask metrics captured") + + val cometBytes = mapInputBytes.sum + val cometRecords = mapInputRecords.sum + + assert( + cometBytes > 0, + s"bytesRead across map tasks should be > 0 for a shuffled scan, got $cometBytes") + assert( + cometRecords == 10000, + s"recordsRead across map tasks should equal the scanned row count, got $cometRecords") + + val sqlBytes = scanNodes.map(_.metrics("bytes_scanned").value).sum + assert( + sqlBytes == cometBytes, + s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + } finally { + spark.sparkContext.removeSparkListener(listener) + spark.sql("DROP TABLE test_cat.db.shuffle_metrics_test") + } + } + } + } + test("exchange reuse must not collapse scans with different pushed filters (#4774)") { assume(icebergAvailable, "Iceberg not available") From 8030ff3ea4de1427c69b310d0c4cf0053248624e Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Mon, 10 Aug 2026 22:05:18 -0700 Subject: [PATCH 3/5] Style fix --- .../test/scala/org/apache/comet/CometIcebergNativeSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index e31cf7ee1a6..eb69bf0032b 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -3859,7 +3859,7 @@ class CometIcebergNativeSuite val scanNodes = nativeShuffles.flatMap { s => assert( s.child.isInstanceOf[CometNativeExec], - s"Expected the shuffle's child to be a CometNativeExec so its subtree is " + + "Expected the shuffle's child to be a CometNativeExec so its subtree is " + s"inlined into the writer plan, got ${s.child.getClass.getSimpleName}:\n$plan") collectIcebergNativeScans(s.child) } From 534ad3732a079cf4a9f9fd825d90d7a246801a96 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 27 Aug 2026 15:52:17 -0700 Subject: [PATCH 4/5] CometCsvNativeScanExec doesn't report bytes_scanned --- .../org/apache/spark/sql/comet/operators.scala | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) 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 3d451e7b94f..9a723560ead 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 @@ -822,14 +822,13 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - // Any leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`, - // `CometCsvNativeScanExec`, contrib scans) can contribute `bytes_scanned` / - // `output_rows` to Spark's task-level input metrics, which drive the Input column - // on the UI's Stages and Executors tabs. Matching on `CometLeafExec` rather than - // `CometNativeScanExec` keeps every scan reported once the scan is fused into a - // larger native block, where only the block root's `compute` runs. - // `reportScanInputMetrics` self-filters on the `bytes_scanned` metric, so leaves - // that don't track it are a no-op. + // A leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`) can + // contribute `bytes_scanned` / `output_rows` to Spark's task-level input metrics, + // which drive the Input column on the UI's Stages and Executors tabs. + // Matching on `CometLeafExec` rather than `CometNativeScanExec` keeps every scan + // reported once the scan is fused into a larger native block, where only the block + // root's `compute` runs. `reportScanInputMetrics` self-filters on the `bytes_scanned` + // metric, so leaves that don't track it are a no-op. hasScanInput = sparkPlans.exists(_.isInstanceOf[CometLeafExec])) } From 51994ad492cf225c5e161661b358bd121e45064e Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 27 Aug 2026 22:27:55 -0700 Subject: [PATCH 5/5] Refactor input metrics test --- .../comet/CometIcebergNativeSuite.scala | 348 ++++++------------ 1 file changed, 117 insertions(+), 231 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index eb69bf0032b..5ae4044378e 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -22,8 +22,8 @@ package org.apache.comet import java.io.File import java.net.URI import java.nio.charset.StandardCharsets.UTF_8 +import java.util.concurrent.atomic.AtomicLong -import scala.collection.mutable import scala.jdk.CollectionConverters._ import org.apache.spark.CometListenerBusUtils @@ -3561,6 +3561,83 @@ class CometIcebergNativeSuite } } + /** Row count written by [[createInputMetricsTable]], and the expected `recordsRead`. */ + private val inputMetricsRows = 10000L + + /** + * Creates `table` as (id INT, value DOUBLE) holding [[inputMetricsRows]] rows spread over + * several files, so scanning it runs multiple tasks that each read a non-zero number of bytes. + */ + private def createInputMetricsTable(table: String): Unit = { + spark.sql(s"CREATE TABLE $table (id INT, value DOUBLE) USING iceberg") + spark + .range(inputMetricsRows) + .selectExpr("CAST(id AS INT) AS id", "CAST(id * 1.5 AS DOUBLE) AS value") + .repartition(5) + .write + .format("iceberg") + .mode("append") + .saveAsTable(table) + } + + /** + * Native operators holding an Iceberg scan as a direct child, i.e. the scan is fused with them. + */ + private def icebergScanFusingParents(plan: SparkPlan): Seq[CometNativeExec] = + collect(plan) { + case p: CometNativeExec if p.children.exists(_.isInstanceOf[CometIcebergNativeScanExec]) => + p + } + + /** + * Runs `body` and returns the (bytesRead, recordsRead) totals Spark reported across the tasks + * it launched. Events from the table setup are drained first so they cannot leak into the + * totals. Reduce-stage tasks read shuffle blocks rather than files, so they add nothing and + * need no filtering. + */ + private def taskInputMetrics(body: => Unit): (Long, Long) = { + val bytesRead = new AtomicLong() + val recordsRead = new AtomicLong() + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + bytesRead.addAndGet(taskEnd.taskMetrics.inputMetrics.bytesRead) + recordsRead.addAndGet(taskEnd.taskMetrics.inputMetrics.recordsRead) + } + } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + spark.sparkContext.addSparkListener(listener) + try { + body + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + (bytesRead.get(), recordsRead.get()) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + /** + * Asserts the task-level input metrics account for every row scanned and every byte the scans + * report, which is what drives the Input column on the UI's Stages and Executors tabs. + */ + private def assertScanInputMetrics( + scans: Seq[CometIcebergNativeScanExec], + bytesRead: Long, + recordsRead: Long): Unit = { + assert(bytesRead > 0, s"bytesRead should be > 0, got $bytesRead") + assert( + recordsRead == inputMetricsRows, + s"recordsRead should equal the scanned row count $inputMetricsRows, got $recordsRead") + val sqlBytes = scans.map(_.metrics("bytes_scanned").value).sum + assert( + sqlBytes == bytesRead, + s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($bytesRead)") + } + + /** + * `SELECT *` leaves the scan as the root of its own native block, so Spark calls + * `CometIcebergNativeScanExec.doExecuteColumnar` directly and that method registers the + * input-metric reporting listener itself. + */ test("task-level inputMetrics.bytesRead is populated for Iceberg native scan") { assume(icebergAvailable, "Iceberg not available in classpath") @@ -3573,87 +3650,23 @@ class CometIcebergNativeSuite CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { - spark.sql(""" - CREATE TABLE test_cat.db.task_metrics_test ( - id INT, - value DOUBLE - ) USING iceberg - """) - - spark - .range(10000) - .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") - .repartition(5) - .write - .format("iceberg") - .mode("append") - .saveAsTable("test_cat.db.task_metrics_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 - if (im.bytesRead > 0) { - bytesReadValues.synchronized { - bytesReadValues += im.bytesRead - recordsReadValues += im.recordsRead - } - } - } - } - spark.sparkContext.addSparkListener(listener) - + createInputMetricsTable("test_cat.db.task_metrics_test") try { - val query = "SELECT * FROM test_cat.db.task_metrics_test" - - // Same drain-run-drain pattern as CometTaskMetricsSuite's shuffle test - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - - // Baseline: iceberg-Java scan (Comet native disabled) - withSQLConf(CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "false") { - bytesReadValues.clear() - recordsReadValues.clear() - spark.sql(query).collect() - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - } - val sparkBytes = bytesReadValues.sum - val sparkRecords = recordsReadValues.sum - - // Comet native Iceberg scan - bytesReadValues.clear() - recordsReadValues.clear() - val df = spark.sql(query) - - val scanNodes = df.queryExecution.executedPlan - .collectLeaves() - .collect { case s: CometIcebergNativeScanExec => s } - assert(scanNodes.nonEmpty, "Expected CometIcebergNativeScanExec in plan") - - df.collect() - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - - val cometBytes = bytesReadValues.sum - val cometRecords = recordsReadValues.sum - - // Both paths should report metrics - assert(sparkBytes > 0, s"Spark bytesRead should be > 0, got $sparkBytes") - assert(sparkRecords > 0, s"Spark recordsRead should be > 0, got $sparkRecords") - assert(cometBytes > 0, s"Comet bytesRead should be > 0, got $cometBytes") - assert(cometRecords > 0, s"Comet recordsRead should be > 0, got $cometRecords") + val df = spark.sql("SELECT * FROM test_cat.db.task_metrics_test") + val (bytesRead, recordsRead) = taskInputMetrics(df.collect()) + // Inspect the plan after execution so we assert on what AQE actually ran. + val plan = df.queryExecution.executedPlan + val scans = collectIcebergNativeScans(plan) + assert(scans.nonEmpty, s"Expected CometIcebergNativeScanExec in plan:\n$plan") + // No native parent, so the scan reports for itself. Pinning the shape keeps this test + // from silently becoming a duplicate of the fused one below. assert( - cometRecords == sparkRecords, - s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords") + icebergScanFusingParents(plan).isEmpty, + s"Expected the scan to be un-fused:\n$plan") - // SQL-level metric should match task-level metric - val sqlBytes = scanNodes.head.metrics("bytes_scanned").value - assert( - sqlBytes == cometBytes, - s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + assertScanInputMetrics(scans, bytesRead, recordsRead) } finally { - spark.sparkContext.removeSparkListener(listener) spark.sql("DROP TABLE test_cat.db.task_metrics_test") } } @@ -3661,17 +3674,11 @@ class CometIcebergNativeSuite } /** - * Companion to "task-level inputMetrics.bytesRead is populated for Iceberg native scan", which - * runs `SELECT *` and therefore leaves the scan as the root of its native block. In that shape - * Spark calls `CometIcebergNativeScanExec.doExecuteColumnar` directly, and that method reports - * input metrics itself. - * - * This test covers the other reporting site. With an operator above the scan, the two fuse into - * one native block and only the block root's `compute` runs -- a parent `CometNativeExec` reads - * its scan children via `PlanDataInjector.findAllPlanData` instead of executing them. Reporting - * then comes from `CometNativeExec.executeColumnarWithContext`, whose `hasScanInput` gate used - * to match only `CometNativeScanExec` and so skipped Iceberg, leaving the Input column on the - * UI's Stages and Executors tabs blank. + * With an operator above it the scan fuses into one native block, so + * `CometIcebergNativeScanExec.doExecuteColumnar` never runs -- the parent reads its scan child + * via `PlanDataInjector.findAllPlanData` instead of executing it -- and reporting comes from + * `CometNativeExec.executeColumnarWithContext`, whose `hasScanInput` gate used to match only + * `CometNativeScanExec` and so skipped Iceberg, leaving the Input column blank. */ test("task-level inputMetrics is populated when Iceberg native scan is fused into a block") { assume(icebergAvailable, "Iceberg not available in classpath") @@ -3685,82 +3692,23 @@ class CometIcebergNativeSuite CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { - spark.sql(""" - CREATE TABLE test_cat.db.fused_metrics_test ( - id INT, - value DOUBLE - ) USING iceberg - """) - - spark - .range(10000) - .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") - .repartition(5) - .write - .format("iceberg") - .mode("append") - .saveAsTable("test_cat.db.fused_metrics_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 - if (im.bytesRead > 0) { - bytesReadValues.synchronized { - bytesReadValues += im.bytesRead - recordsReadValues += im.recordsRead - } - } - } - } - spark.sparkContext.addSparkListener(listener) - + createInputMetricsTable("test_cat.db.fused_metrics_test") try { // Arithmetic in the projection keeps it from being collapsed into the scan, so a // CometProjectExec sits above the scan and the two fuse into one native block. - val query = - "SELECT id + 1 AS id2, value * 2 AS value2 FROM test_cat.db.fused_metrics_test" + val df = spark.sql( + "SELECT id + 1 AS id2, value * 2 AS value2 FROM test_cat.db.fused_metrics_test") + val (bytesRead, recordsRead) = taskInputMetrics(df.collect()) - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - bytesReadValues.clear() - recordsReadValues.clear() - - val df = spark.sql(query) val plan = df.queryExecution.executedPlan - val scanNodes = collectIcebergNativeScans(plan) - assert(scanNodes.nonEmpty, s"Expected CometIcebergNativeScanExec in plan:\n$plan") - // Assert the fusion actually happened, so this test cannot silently degrade into a - // duplicate of the SELECT * one if a future planner change collapses the projection - // into the scan. Checking the executedPlan root is not enough: with AQE on, the root - // is an AdaptiveSparkPlanExec either way. - val fusingParents = collect(plan) { - case p: CometNativeExec - if p.children.exists(_.isInstanceOf[CometIcebergNativeScanExec]) => - p - } - assert( - fusingParents.nonEmpty, - s"Expected the Iceberg scan to be fused under a native parent operator:\n$plan") - - df.collect() - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - - val cometBytes = bytesReadValues.sum - val cometRecords = recordsReadValues.sum - - assert(cometBytes > 0, s"bytesRead should be > 0 for a fused scan, got $cometBytes") + val scans = collectIcebergNativeScans(plan) + assert(scans.nonEmpty, s"Expected CometIcebergNativeScanExec in plan:\n$plan") assert( - cometRecords == 10000, - s"recordsRead should equal the scanned row count, got $cometRecords") + icebergScanFusingParents(plan).nonEmpty, + s"Expected the scan to be fused under a native parent operator:\n$plan") - val sqlBytes = scanNodes.map(_.metrics("bytes_scanned").value).sum - assert( - sqlBytes == cometBytes, - s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + assertScanInputMetrics(scans, bytesRead, recordsRead) } finally { - spark.sparkContext.removeSparkListener(listener) spark.sql("DROP TABLE test_cat.db.fused_metrics_test") } } @@ -3768,14 +3716,10 @@ class CometIcebergNativeSuite } /** - * `CometNativeExec.buildNativeContext` is also consumed by the native shuffle path: - * `CometShuffleExchangeExec.nativeChildContext` builds the context from its `CometNativeExec` - * child, and the child's whole native subtree -- scan included -- is inlined under the - * `ShuffleWriter` protobuf operator and executed by `CometNativeShuffleWriter` in the - * ShuffleMapTask. No `CometExecRDD` runs for that subtree, so neither - * `CometIcebergNativeScanExec.doExecuteColumnar` nor - * `CometNativeExec.executeColumnarWithContext` reports anything; the writer has its own - * `ctx.hasScanInput` check instead. + * The native shuffle path inlines the child's whole native subtree -- scan included -- under + * the `ShuffleWriter` protobuf operator and executes it in the ShuffleMapTask. No + * `CometExecRDD` runs for that subtree, so neither site above reports anything and + * `CometNativeShuffleWriter` has its own `ctx.hasScanInput` check instead. * * Before the fix, an Iceberg scan feeding a native shuffle left the map stage's Input column * blank. The Parquet equivalent ("native shuffle reports task input metrics for its scan child" @@ -3799,56 +3743,14 @@ class CometIcebergNativeSuite // switch to columnar shuffle (which reports input metrics through a different path). CometConf.COMET_SHUFFLE_MODE.key -> "native") { - spark.sql(""" - CREATE TABLE test_cat.db.shuffle_metrics_test ( - id INT, - value DOUBLE - ) USING iceberg - """) - - spark - .range(10000) - .selectExpr("CAST(id AS INT)", "CAST(id * 1.5 AS DOUBLE) as value") - .repartition(5) - .write - .format("iceberg") - .mode("append") - .saveAsTable("test_cat.db.shuffle_metrics_test") - - val mapInputBytes = mutable.ArrayBuffer.empty[Long] - val mapInputRecords = mutable.ArrayBuffer.empty[Long] - - val listener = new SparkListener { - override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { - // Only the map stage runs the scan; the reduce stage reads shuffle blocks and must - // not contribute to inputMetrics. - if (taskEnd.taskType.contains("ShuffleMapTask")) { - val im = taskEnd.taskMetrics.inputMetrics - mapInputBytes.synchronized { - mapInputBytes += im.bytesRead - mapInputRecords += im.recordsRead - } - } - } - } - spark.sparkContext.addSparkListener(listener) - + createInputMetricsTable("test_cat.db.shuffle_metrics_test") try { - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - mapInputBytes.clear() - mapInputRecords.clear() - - val df = spark - .table("test_cat.db.shuffle_metrics_test") - .repartition(4, col("id")) - - df.collect() - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + val df = spark.table("test_cat.db.shuffle_metrics_test").repartition(4, col("id")) + val (bytesRead, recordsRead) = taskInputMetrics(df.collect()) - // Assert the plan shape after execution so we inspect what AQE actually ran. All three - // conditions are required for the writer to be the reporting site: a native (not - // columnar) shuffle, a CometNativeExec child so `nativeChildContext` is `Some`, and an - // Iceberg scan inside that child's subtree so `hasScanInput` must be true. + // All three conditions are required for the writer to be the reporting site: a native + // (not columnar) shuffle, a CometNativeExec child so `nativeChildContext` is `Some`, and + // an Iceberg scan inside that child's subtree so `hasScanInput` must be true. val plan = df.queryExecution.executedPlan val nativeShuffles = collect(plan) { case s: CometShuffleExchangeExec if s.shuffleType == CometNativeShuffle => s @@ -3856,7 +3758,7 @@ class CometIcebergNativeSuite assert( nativeShuffles.nonEmpty, s"Expected a CometShuffleExchangeExec with CometNativeShuffle in plan:\n$plan") - val scanNodes = nativeShuffles.flatMap { s => + val scans = nativeShuffles.flatMap { s => assert( s.child.isInstanceOf[CometNativeExec], "Expected the shuffle's child to be a CometNativeExec so its subtree is " + @@ -3864,27 +3766,11 @@ class CometIcebergNativeSuite collectIcebergNativeScans(s.child) } assert( - scanNodes.nonEmpty, + scans.nonEmpty, s"Expected the Iceberg scan to be inlined under the native shuffle:\n$plan") - assert(mapInputRecords.nonEmpty, "no ShuffleMapTask metrics captured") - - val cometBytes = mapInputBytes.sum - val cometRecords = mapInputRecords.sum - - assert( - cometBytes > 0, - s"bytesRead across map tasks should be > 0 for a shuffled scan, got $cometBytes") - assert( - cometRecords == 10000, - s"recordsRead across map tasks should equal the scanned row count, got $cometRecords") - - val sqlBytes = scanNodes.map(_.metrics("bytes_scanned").value).sum - assert( - sqlBytes == cometBytes, - s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + assertScanInputMetrics(scans, bytesRead, recordsRead) } finally { - spark.sparkContext.removeSparkListener(listener) spark.sql("DROP TABLE test_cat.db.shuffle_metrics_test") } }