[PR] fix: refresh stale shuffle scans when AQE prepares a stage, not while it re-plans [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: refresh stale shuffle scans when AQE prepares a stage, not while it re-plans [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove merged PR #6646:
URL: https://github.com/apache/datafusion-comet/pull/6646


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org