andygrove opened a new pull request, #6646:
URL: https://github.com/apache/datafusion-comet/pull/6646
## Which issue does this PR close?
Closes #6645.
## Rationale for this change
#6547 patches the stale `Scan` leaf of a native operator that AQE reuses
over a shuffle stage. It does this in `CometExecRule` while that rule runs as
AQE's query stage preparation rule, so during each re-plan. The re-planned plan
then differs from the current one, and AQE adopts it. After that, AQE cannot
link a broadcast stage it builds over the reused operator back into its logical
plan, so the join's build side is no longer pinned and a later re-plan can move
the broadcast to the other side. On Spark 3.4 the AQE DPP probe side stays in
Spark and its shuffle reports a smaller size than Comet's, so the join flips to
`BuildLeft` and the DPP subquery's broadcast is no longer reused. This fails
Spark's `SPARK-34637` test in the Spark 3.4 `sql_core-1` shard on every pull
request that runs `run-spark-3.4-tests`, for example #4565 and #5841.
## What changes are included in this PR?
- `CometExecRule` takes a `queryStagePrep` flag, which `CometRule` passes
from its `injectQueryStagePrepRule` instance. The stale `Scan` is now patched
only in the columnar pass that prepares a query stage or the final plan, so AQE
re-plans are unchanged again, as they were before #6547.
- By the time a stage is prepared, AQE may have put an `AQEShuffleReadExec`
between the stage and the reused operator, for example when it coalesces
partitions, which it does by default. The refresh looks through that read to
the stage's `ShuffleScan` sink. The `ShuffleScan` then reads the partitions
that the read specifies, as it already did for coalesced reads under #6547.
## How are these changes tested?
- New `CometExecRuleSuite` tests:
- the query stage preparation pass leaves an AQE-reused aggregate's `Scan`
alone;
- the refresh goes through an AQE shuffle read;
- an end-to-end check that AQE re-plans keep a broadcast side's native
aggregate over its shuffle stage as it was, which runs on every Spark version.
- The Spark 3.5+ test "AQE DPP broadcast roots retain temporary logical
links after an unchanged replan" now also runs with direct read on. The test
that expected a re-plan to change the broadcast is removed.
- Mutation checks:
- Patching during re-plans again fails all three re-plan tests.
- Dropping the look-through for the AQE shuffle read fails the new unit
test and #6547's coalesced case in "final aggregate reads its shuffle through
ShuffleScan only under AQE".
- `CometExecRuleSuite` and `CometNativeShuffleSuite` pass locally on Spark
3.4 and 4.1, including #6547's checks that executed plans read their shuffles
through `ShuffleScan`. Scalafix also passes on 3.4.
- Spark 3.4's `DynamicPartitionPruningV1SuiteAEOn`, run from the Spark test
jar with Comet enabled, passes in full, `SPARK-34637` included. Before this
change that test fails locally the same way it does in CI.
--
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]
