andygrove opened a new issue, #6152:
URL: https://github.com/apache/datafusion-comet/issues/6152
### Describe the bug
`RevertNativeForTransitionHeavyStages.transformStageDown` is meant to stay
inside one stage ("never descends stage-boundary children"), but it can cross
into the stage below when it strips a transition whose child is a stage
boundary.
The boundary check only runs on a node's children before recursing into
them. When the strip rule rewrites a node, the result is not checked:
- On `main`, the rule turns `ColumnarToRowExec(boundary)` into `boundary`,
and then `transformed.children.map { ... }` recurses into the boundary's
children.
- #5957 changes this branch to `if (transformed ne plan)
transformStageDown(transformed)(rule)`, which makes the boundary the new root
and then maps its children. It crosses in the same way.
With AQE on this is harmless. The boundaries are `QueryStageExec` nodes,
which have no children in the physical plan.
With AQE off (`applyForNonAQE`), the boundary is a `ShuffleExchangeLike`
whose child is the stage below. Take an upper stage where a Spark row operator
reads a Comet shuffle through a transition:
```text
<row operator>
+- ColumnarToRowExec
+- CometShuffleExchangeExec <- stage boundary
+- <lower stage: Comet and Spark operators, with its own transitions>
```
If the upper stage exceeds the transition threshold, `revertToSpark` strips
the transitions inside the lower stage. The lower stage's Comet operators are
not reverted, because the `transformStageUp` pass stops at the exchange.
`insertTransitions` also stops at the exchange, so the stripped transitions are
never put back. A row operator in the lower stage can be left over a columnar
child with no transition between them.
This was found by reading the code while reviewing #5957, not by a
reproduction. It needs `spark.comet.exec.transitionRevert.enabled=true`
(default `false`) and AQE off.
### Steps to reproduce
A unit test on `revertToSpark` is probably the easiest route: build a plan
where a transition-heavy upper stage has a `ColumnarToRowExec` directly above a
`CometShuffleExchangeExec`, and the lower stage contains a transition of its
own. Call `revertToSpark` on the upper stage and check whether the lower
stage's transition is still present.
End to end: disable AQE, set
`spark.comet.exec.transitionRevert.enabled=true` and
`spark.comet.exec.transitionRevert.maxTransitions=0`, and run a query whose
upper stage has a JVM operator reading a Comet shuffle while the lower stage
mixes Comet and JVM operators.
### Expected behavior
Reverting a stage changes nothing below its stage boundaries. In
`transformStageDown`, a rewritten node that is itself a stage boundary should
be returned unchanged rather than recursed into, for example:
```scala
if (transformed ne plan) {
if (isStageBoundary(transformed)) transformed else
transformStageDown(transformed)(rule)
}
```
The same check is needed on `main`'s current shape if this is fixed before
#5957 merges.
### Additional context
Pre-existing on `main`; #5957 keeps the behavior while fixing #5719 in the
same function. Related: #5719, #5957.
--
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]
