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 74ef3af91c..3d451e7b94 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 c816255924..43d96670ae 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} @@ -3659,6 +3660,237 @@ 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") + } + } + } + } + + /** + * `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")