fix: report task input metrics after the native iterator closes and add to Spark's counters

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)
       }
     }
   }

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,38 @@ 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. That holds only when the fallback scan registers its
+   * completion listener after this one, which a `CometSparkToColumnarExec` input always does and
+   * a coalesced Spark-scan partition computed first does not: `FileScanRDD`'s close then runs
+   * last and sets bytesRead from the value it snapshotted at construction. 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 +145,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 +175,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 +239,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`.
    */

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)
       }
     }
   }

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)
       }
     }
   }

spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala

@@ -3738,6 +3738,82 @@ 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 {
+          // This Iceberg scan is its own native block with no JVM input, so its metrics publish
+          // per batch and this covers the registration site rather than the listener order,
+          // which stays unguarded for a fused Iceberg scan until the scan-input gate widens.
+          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")