RevertNativeForTransitionHeavyStages strips transitions in the stage below when AQE is off

[andygrove]

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.
• fix: restore Spark write execs when reverting transition-heavy stages #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:

<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:

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.