andygrove commented on code in PR #5880:
URL: https://github.com/apache/datafusion-comet/pull/5880#discussion_r3996849011
##########
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.
Review Comment:
The Iceberg site gets the same ordering fix, but I do not think any existing
test can fail without it. `"task-level inputMetrics.bytesRead is populated for
Iceberg native scan"` is a plain `SELECT *` with no JVM input and no early
stop, which is the batch-receiver shape where `update_metrics` fires per batch,
so it passes either way. An Iceberg scan with a `LIMIT` at the default update
interval would mirror your parquet test and give this line a guard.
##########
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:
I do not think the `inc` change holds in the other arm order. `FileScanRDD`
snapshots `existingBytesRead` when its iterator is constructed in `compute()`
and its completion-time `close()` does `setBytesRead(existingBytesRead +
getBytesReadCallback())`, so it replaces rather than adds. Comet's `inc` only
survives when Spark's close listener runs first, which requires Spark's scan to
have registered its listener after Comet's.
In `"task input metrics keep bytes read by a fallback Spark scan in the same
task"` the parquet arm comes first, so Comet registers first and Spark's close
runs first. If I swap the arms so the JSON scan is the leading one, reading
exactly the same two files, Comet's `bytesRead` drops by the parquet side's
entire `bytes_scanned`:
```
native-first: sparkBytes=147770 cometBytes=166181
fallback-first: sparkBytes=147770 cometBytes=146148
```
`recordsRead` is correct in both orders, since `FileScanRDD` uses
`incRecordsRead` there. The scope is narrow: a fallback scan reaching a native
block through `CometSparkToColumnarExec` is always safe, because
`CometExecRDD.compute` registers this listener before `resolveInputObjects`
pulls the input RDDs. It takes a coalesce that merges a Spark-scan partition
and a Comet-scan partition into one task with the Spark side first, which is
the shape the new test uses.
This is not a regression, both orders are wrong on main today, and I do not
see a clean fix given `TaskContext` has no listener ordering control. 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? As written
the description reads as though the fallback bytes always survive.
##########
spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala:
##########
@@ -1012,6 +1090,288 @@ class CometTaskMetricsSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("native scan left unconsumed by a limit still reports task input
metrics") {
+ withTempPath { dir =>
+ spark
+ .createDataFrame((0 until 20000).map(i => (i, s"big_$i")))
+ .repartition(4)
+ .write
+ .parquet(dir.getAbsolutePath)
+
spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("limit_big")
+ spark
+ .createDataFrame((0 until 100).map(i => (i * 7, s"local_$i")))
+ .createOrReplaceTempView("limit_local")
+
+ // The broadcast side is a JVM input to the join block, so native
execution polls and only
+ // publishes scan metrics on the update interval. The limit stops
pulling before the scan
+ // is exhausted, leaving the final publish to the iterator's
completion-time close.
+ val query = "SELECT /*+ BROADCAST(limit_local) */ * FROM limit_big JOIN
limit_local " +
+ "ON limit_big._1 = limit_local._1 LIMIT 3"
+ val localBroadcast = Seq(
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key ->
"LocalTableScan")
+
+ Seq("-1",
CometConf.COMET_METRICS_UPDATE_INTERVAL.defaultValueString).foreach { interval
=>
+ val confs = localBroadcast :+
(CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval)
+ val (cometBytes, cometRecords, cometPlan) =
+ collectInputMetrics(query, (CometConf.COMET_ENABLED.key -> "true")
+: confs: _*)
+
+ val join = find(cometPlan)(_.isInstanceOf[CometBroadcastHashJoinExec])
+ assert(
+ join.isDefined,
+ s"Expected CometBroadcastHashJoinExec in
plan:\n${cometPlan.treeString}")
+ val streamedNativeScan = join.get.children.exists { child =>
+ find(child)(_.isInstanceOf[CometBroadcastExchangeExec]).isEmpty &&
+ find(child)(_.isInstanceOf[CometNativeScanExec]).isDefined
+ }
+ assert(
+ streamedNativeScan,
+ s"Expected a native scan on the streamed side of the
join:\n${cometPlan.treeString}")
+
+ 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")
+ }
+ }
+ }
+
+ test("task input metrics keep bytes read by a fallback Spark scan in the
same task") {
+ withTempPath { parquetDir =>
+ withTempPath { jsonDir =>
+ spark
+ .createDataFrame((0 until 5000).map(i => (i, s"parquet_$i")))
+ .repartition(1)
+ .write
+ .parquet(parquetDir.getAbsolutePath)
+ spark
+ .createDataFrame((5000 until 10000).map(i => (i, s"json_$i")))
+ .repartition(1)
+ .write
+ .json(jsonDir.getAbsolutePath)
+
spark.read.parquet(parquetDir.getAbsolutePath).createOrReplaceTempView("mixed_parquet")
+
spark.read.json(jsonDir.getAbsolutePath).createOrReplaceTempView("mixed_json")
+
+ // Coalescing the union to one partition computes the native scan and
the fallback JSON
+ // scan inside the same task, so Spark's own input metrics and Comet's
must add up.
+ val query = "SELECT /*+ COALESCE(1) */ * FROM (SELECT _1 FROM
mixed_parquet " +
+ "UNION ALL SELECT CAST(_1 AS INT) FROM mixed_json)"
+ val convertJson = CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key ->
"true"
+
+ val (sparkBytes, sparkRecords, _) =
+ collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "false",
convertJson)
+ val (cometBytes, cometRecords, cometPlan) =
+ collectInputMetrics(query, CometConf.COMET_ENABLED.key -> "true",
convertJson)
+
+ assert(
+ find(cometPlan)(_.isInstanceOf[CometNativeScanExec]).isDefined,
+ s"Expected CometNativeScanExec in plan:\n${cometPlan.treeString}")
+ assert(
+ find(cometPlan)(_.isInstanceOf[CometSparkToColumnarExec]).isDefined,
+ s"Expected CometSparkToColumnarExec in
plan:\n${cometPlan.treeString}")
+
+ assert(sparkRecords > 0, s"Spark recordsRead should be > 0, got
$sparkRecords")
+ assert(
+ cometRecords == sparkRecords,
+ s"recordsRead mismatch: comet=$cometRecords, spark=$sparkRecords")
+ assert(sparkBytes > 0, s"Spark bytesRead should be > 0, got
$sparkBytes")
+ assertCometBytesReadInRange(cometBytes, sparkBytes)
Review Comment:
When I reverted `inc` back to `set` to check what the new tests catch, all
three integration failures were on `recordsRead` (`5000 did not equal 10000`),
none on `bytesRead`. `assertCometBytesReadInRange` allows a 0.7 to 1.3 ratio on
3.5 and 4.0, and merely `cometBytes >= sparkBytes` on 4.1+. In this fixture the
parquet side is small enough relative to the JSON side that losing it entirely
still lands around 0.88, inside the band. So the bytes half of the change is
not pinned by anything here, or in the coalesced-union and cached-input tests.
Could this measure the two sides separately first and then assert the
coalesced run is at least their sum? That would make the bytes assertion fail
on a `set` revert the same way the records one does.
##########
spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala:
##########
@@ -128,6 +128,84 @@ class CometTaskMetricsSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("scan input metrics add to the task once per accumulator and keep
existing values") {
+ val nestedBytes = new SQLMetric("nestedBytes", -1L)
+ val nestedRows = new SQLMetric("nestedRows")
+ val nestedPruned = new SQLMetric("nestedPruned")
+ val siblingBytes = new SQLMetric("siblingBytes", -1L)
+ val siblingRows = new SQLMetric("siblingRows")
+ val nestedScan = CometMetricNode(
+ Map(
+ "bytes_scanned" -> nestedBytes,
+ "output_rows" -> nestedRows,
+ "pushdown_rows_pruned" -> nestedPruned))
+ // An operator above the scans contributes output_rows of its own, which
is not input.
+ val outerRows = new SQLMetric("outerRows")
+ outerRows.set(99L)
+ val outerTree = CometMetricNode(
+ Map("output_rows" -> outerRows),
+ Seq(
+ nestedScan,
+ CometMetricNode(Map("bytes_scanned" -> siblingBytes, "output_rows" ->
siblingRows))))
+
+ Seq(None, Some(new IllegalStateException("failed native stage"))).foreach
{ failure =>
+ val ctx = TaskContext.empty()
Review Comment:
Worth a line in the suite noting that these tests depend on
`markTaskCompleted` running. Every `TaskContext.empty()` is constructed with
`taskAttemptId == 0`, so they all share one `seenMetricsByTask` entry, and the
cleanup listener is the only thing keeping them isolated from each other. A
future test that registers a report and forgets `markTaskCompleted` would
silently zero the next test's claims rather than failing.
--
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]
