This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5265-971380971064082215744f4fc4d86177bdc416c7 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 244bb6ce73efaf28f72c8e2c71a07215554ee836 Author: hsiang-c <[email protected]> AuthorDate: Thu Sep 24 00:44:19 2026 +0000 fix: report `Input` column when native Iceberg scan is enabled or native shuffle is enabled (#5265) * Populate executor's input metrics on fused operators * Test Shuffle Writer metrics * Style fix * CometCsvNativeScanExec doesn't report bytes_scanned * Refactor input metrics test --- .../org/apache/spark/sql/comet/operators.scala | 9 +- .../org/apache/comet/CometIcebergNativeSuite.scala | 262 +++++++++++++++------ 2 files changed, 198 insertions(+), 73 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 7bf6e84548..ed2d62a1a3 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 @@ -1050,7 +1050,14 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - hasScanInput = sparkPlans.exists(_.isInstanceOf[CometNativeScanExec])) + // 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])) } /** diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 0e88739de7..256ed22135 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.iceberg.data.IcebergGenerics @@ -34,10 +34,11 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{CometTestBase, DataFrame, Row} 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.{ArrayType, StringType, StructType, TimestampType} @@ -4073,6 +4074,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") @@ -4085,88 +4163,128 @@ 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" + 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( + icebergScanFusingParents(plan).isEmpty, + s"Expected the scan to be un-fused:\n$plan") - // Same drain-run-drain pattern as CometTaskMetricsSuite's shuffle test - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + assertScanInputMetrics(scans, bytesRead, recordsRead) + } finally { + spark.sql("DROP TABLE test_cat.db.task_metrics_test") + } + } + } + } - // 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 + /** + * 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") - // Comet native Iceberg scan - bytesReadValues.clear() - recordsReadValues.clear() - val df = spark.sql(query) + 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") { - val scanNodes = df.queryExecution.executedPlan - .collectLeaves() - .collect { case s: CometIcebergNativeScanExec => s } - assert(scanNodes.nonEmpty, "Expected CometIcebergNativeScanExec in plan") + 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 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()) + + val plan = df.queryExecution.executedPlan + val scans = collectIcebergNativeScans(plan) + assert(scans.nonEmpty, s"Expected CometIcebergNativeScanExec in plan:\n$plan") + assert( + icebergScanFusingParents(plan).nonEmpty, + s"Expected the scan to be fused under a native parent operator:\n$plan") - df.collect() - CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + assertScanInputMetrics(scans, bytesRead, recordsRead) + } finally { + spark.sql("DROP TABLE test_cat.db.fused_metrics_test") + } + } + } + } - val cometBytes = bytesReadValues.sum - val cometRecords = recordsReadValues.sum + /** + * 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" + * 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") - // 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") + 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") { + createInputMetricsTable("test_cat.db.shuffle_metrics_test") + try { + val df = spark.table("test_cat.db.shuffle_metrics_test").repartition(4, col("id")) + val (bytesRead, recordsRead) = taskInputMetrics(df.collect()) + + // 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( - cometRecords == sparkRecords, - s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords") - - // SQL-level metric should match task-level metric - val sqlBytes = scanNodes.head.metrics("bytes_scanned").value + nativeShuffles.nonEmpty, + s"Expected a CometShuffleExchangeExec with CometNativeShuffle in plan:\n$plan") + val scans = nativeShuffles.flatMap { s => + assert( + s.child.isInstanceOf[CometNativeExec], + "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( - sqlBytes == cometBytes, - s"SQL bytes_scanned ($sqlBytes) should match task bytesRead ($cometBytes)") + scans.nonEmpty, + s"Expected the Iceberg scan to be inlined under the native shuffle:\n$plan") + + assertScanInputMetrics(scans, bytesRead, recordsRead) } finally { - spark.sparkContext.removeSparkListener(listener) - spark.sql("DROP TABLE test_cat.db.task_metrics_test") + spark.sql("DROP TABLE test_cat.db.shuffle_metrics_test") } } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
