schenksj opened a new issue, #7481:
URL: https://github.com/apache/datafusion-comet/issues/7481
### Describe the bug
Before timing a case, `runExpressionBenchmark` in
`spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala`
checks that the
Comet plan is fully native:
```scala
// Check that the plan is fully Comet native before running the benchmark.
Unlike the check
// above this one does execute the query, because AQE can re-plan a join
after the first
// stage completes and the Comet operators are what we are checking for.
withSQLConf(cometConfigs: _*) {
val df = spark.sql(query)
df.noop()
val plan = stripAQEPlan(df.queryExecution.executedPlan)
findFirstNonCometOperator(plan).foreach { op => ...
```
`df.noop()` runs a separate write query. It does not execute `df`'s own
plan. After it,
`df.queryExecution.executedPlan` is still `AdaptiveSparkPlan
isFinalPlan=false`, and
`stripAQEPlan` returns the initial plan. The check therefore never sees what
AQE changes during
execution, which is what the comment says it is checking for. The warning it
prints shows the
initial plan, not the plan that ran.
### Steps to reproduce
Run on `main` (`0f7f8e453`) with the default profile (Spark 4.1.3), in a
`CometBenchmarkBase`
subclass, over a Parquet table of 4096 rows made with
`spark.range(4096).selectExpr("id % 10 AS k", "id AS v")`, with
`spark.comet.exec.shuffle.enabled=false`. The query is `SELECT k, sum(v)
FROM t GROUP BY k`.
1. With the three lines above, after `df.noop()`,
`df.queryExecution.executedPlan` is
`AdaptiveSparkPlan isFinalPlan=false`. The plan that the check reads, and
that the warning
prints, is
`HashAggregate` > `Exchange` > `CometHashAggregate` > `CometNativeScan`.
It has no `CometColumnarToRow`, `AQEShuffleRead` or `ShuffleQueryStage`.
2. For a second DataFrame of the same query, after `collect()`,
`executedPlan` is `AdaptiveSparkPlan isFinalPlan=true`, and the plan is
`ResultQueryStage` > `HashAggregate` > `AQEShuffleRead coalesced` >
`ShuffleQueryStage` >
`Exchange` > `CometColumnarToRow` > `CometHashAggregate` >
`CometNativeScan`.
### Expected behavior
The check reads the plan that ran, after AQE has finished, and the warning
prints that plan.
### Additional context
Executing the DataFrame's own plan, for example with
`df.queryExecution.executedPlan.execute().foreach(_ => ())`, makes the plan
final. The final plan
contains query stages. `findFirstNonCometOperator` in
`spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala`
walks the plan with
`foreach`, which does not descend into a `QueryStageExec` (a leaf node), and
its allow-list
contains neither `QueryStageExec` nor `AQEShuffleReadExec`. So a fix also
needs the walk to go
into the query stages, for example through `AdaptiveSparkPlanHelper`, and to
allow the AQE
wrapper nodes.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
