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..1259c1e237 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")