dwsmith1983 commented on code in PR #5880:
URL: https://github.com/apache/datafusion-comet/pull/5880#discussion_r3996922344


##########
spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala:
##########
@@ -105,24 +101,34 @@ 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. 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)

Review Comment:
   > Would you add the reversed-arm query as a test that pins the current 
behaviour, and mention it under "What this does not cover" next to #5265 and 
#5879?
   
   Added both. The test now runs the union in both arm orders. Native first 
asserts bytesRead covers the sum of the two sides measured on their own. 
Fallback first asserts records still add up and bytesRead stays below that sum, 
with a comment on FileScanRDD's close setting the value it snapshotted at 
construction. The description has a bullet for it next to the other two.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to