Repository navigation
fix: restore Spark write execs when reverting transition-heavy stages #5957
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,7 +23,8 @@ import org.apache.spark.internal.Logging | |
| import org.apache.spark.sql.SparkSession | ||
| import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} | ||
| import org.apache.spark.sql.catalyst.rules.Rule | ||
| import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec} | ||
| import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} | ||
| import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan} | ||
| import org.apache.spark.sql.execution.adaptive.QueryStageExec | ||
| import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike} | ||
|
|
@@ -62,32 +63,41 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan | |
| plan match { | ||
| case _: BroadcastExchangeLike => plan | ||
| case exchange: ShuffleExchangeLike => | ||
| revertShuffleStageIfNeeded(exchange) | ||
| .map(reverted => exchange.withNewChildren(Seq(reverted))) | ||
| .getOrElse(plan) | ||
| case _ => | ||
| // Result stage: its output is collected as rows. | ||
| revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan) | ||
|
Member
There was a problem hiding this comment. The result stage still assumes it has to produce rows, and that isn't true for a cached query on Spark 4.0+. When the cache serializer accepts columnar input, which Comet's |
||
| } | ||
| } | ||
|
|
||
| private def applyForNonAQE(plan: SparkPlan): SparkPlan = { | ||
| val withRevertedStages = plan.transformUp { case exchange: ShuffleExchangeLike => | ||
| revertShuffleStageIfNeeded(exchange) | ||
| .map(reverted => exchange.withNewChildren(Seq(reverted))) | ||
| .getOrElse(exchange) | ||
| } | ||
| revertStageIfNeeded(withRevertedStages, outputColumnar = false) | ||
| .getOrElse(withRevertedStages) | ||
| } | ||
|
|
||
| private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = { | ||
| val outputArrow = exchange match { | ||
| case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle | ||
| case _ => false | ||
| } | ||
| revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow) | ||
| } | ||
|
|
||
| /** | ||
| * Reverts the stage if C2R count exceeds threshold, restoring the stage's output format when | ||
| * the reverted root does not satisfy it. | ||
| */ | ||
| private def revertStageIfNeeded( | ||
| stagePlan: SparkPlan, | ||
| outputColumnar: Boolean, | ||
| outputArrow: Boolean = false): Option[SparkPlan] = { | ||
| val transitionCount = countTransitions(stagePlan) | ||
| if (transitionCount <= maxTransitions) return None | ||
|
|
||
|
|
@@ -101,11 +111,27 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan | |
| val reason = | ||
| s"Stage reverted: $transitionCount C2R transitions exceed threshold $maxTransitions" | ||
|
|
||
| val reverted = | ||
| try { | ||
| revertToSpark(stagePlan) | ||
| } catch { | ||
| case e: CometExec.InvalidSparkFallbackException => | ||
| logWarning( | ||
| "Skipping transition-heavy stage reversion because a Comet operator could not " + | ||
| s"restore its Spark plan: ${e.getMessage}") | ||
| return None | ||
| } | ||
| val revertedWithReason = withFallbackReason(reverted, reason) | ||
| val result = if (outputArrow) { | ||
| // Native shuffle consumes Arrow-backed Comet vectors, not arbitrary Spark columnar | ||
| // batches. This bridge converts both row-based and vectorized Spark fallback roots. | ||
| CometSparkToColumnarExec(revertedWithReason) | ||
| } else if (outputColumnar && !reverted.supportsColumnar) { | ||
| RowToColumnarExec(revertedWithReason) | ||
| } else if (!outputColumnar && reverted.supportsColumnar) { | ||
| ColumnarToRowExec(revertedWithReason) | ||
| } else { | ||
| revertedWithReason | ||
| } | ||
| Some(result) | ||
| } | ||
|
|
@@ -146,16 +172,31 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan | |
| } | ||
|
|
||
| /** | ||
| * Like `transformDown`, never descends stage-boundary children. If the rule rewrites the | ||
| * current node, re-apply it to the result so stacked transitions such as | ||
| * `CometSparkToColumnarExec(CometNativeColumnarToRowExec(x))` are fully unwrapped before | ||
| * children are visited. Spark's `transformDown` does not do this; leaving the inner C2R in | ||
| * place later calls `CometNativeColumnarToRowExec.withNewChildren` with a reverted row-based | ||
| * child, which asserts `child.supportsColumnar`. | ||
| * | ||
| * A rewrite can itself be the stage boundary. Unwrapping a transition that sits directly on a | ||
| * shuffle yields that shuffle, and descending into it strips transitions in the next stage. | ||
| * `transformStageUp` and `insertTransitions` do not cross the exchange, so those transitions | ||
| * would not be restored (#6152). Return the boundary unchanged. | ||
| */ | ||
| private def transformStageDown(plan: SparkPlan)( | ||
| rule: PartialFunction[SparkPlan, SparkPlan]): SparkPlan = { | ||
| val transformed = rule.applyOrElse(plan, identity[SparkPlan]) | ||
| if (transformed ne plan) { | ||
| if (isStageBoundary(transformed)) transformed | ||
|
Member
There was a problem hiding this comment. This stops the strip at the exchange, but if the stage's root is itself the transition sitting on the exchange, |
||
| else transformStageDown(transformed)(rule) | ||
| } else { | ||
| val newChildren = transformed.children.map { child => | ||
| if (isStageBoundary(child)) child else transformStageDown(child)(rule) | ||
|
Comment on lines
+190
to
+195
Member
There was a problem hiding this comment. With AQE off this still descends into the stage below an exchange. The boundary check only runs on the children of the transformed node. So when the stripped transition sits directly on a shuffle, the recursion strips the transitions inside the next stage down, and nothing puts them back, because |
||
| } | ||
| if (newChildren == transformed.children) transformed | ||
| else transformed.withNewChildren(newChildren) | ||
| } | ||
| } | ||
|
|
||
| /** Like `transformUp`, never descends stage-boundary children. */ | ||
|
|
@@ -184,48 +225,40 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan | |
| count | ||
| } | ||
|
|
||
| /** | ||
| * Checks for Comet operators whose original Spark plan is also their input. This must run | ||
| * before any bottom-up rewrite replaces children. Otherwise a stable `originalPlan` reference | ||
| * can keep pointing at the old Comet child after `withNewChildren`, hiding the alias and | ||
| * causing fallback to reconstruct that Comet child instead of a Spark operator. | ||
| */ | ||
| private def validateOriginalPlanAliases(plan: SparkPlan): Unit = plan match { | ||
| case _ if isStageBoundary(plan) => () | ||
| case cometExec: CometExec => | ||
| val sparkPlan = cometExec.originalPlan | ||
| if (sparkPlan != null && cometExec.children.exists(_ eq sparkPlan)) { | ||
| throw new CometExec.InvalidSparkFallbackException( | ||
| s"${cometExec.getClass.getSimpleName} aliases its original Spark plan with a child") | ||
| } | ||
| cometExec.children.foreach(validateOriginalPlanAliases) | ||
| case _ => | ||
| plan.children.foreach(validateOriginalPlanAliases) | ||
| } | ||
|
|
||
| private[rules] def revertToSpark(plan: SparkPlan): SparkPlan = { | ||
| validateOriginalPlanAliases(plan) | ||
| val stripped = transformStageDown(plan) { | ||
| case CometNativeColumnarToRowExec(child) => child | ||
| case CometColumnarToRowExec(child) => child | ||
| case ColumnarToRowExec(child) => child | ||
| case sparkToColumnar: CometSparkToColumnarExec => sparkToColumnar.child | ||
| case RowToColumnarExec(child) => child | ||
| } | ||
| if (isStageBoundary(stripped)) { | ||
| throw new CometExec.InvalidSparkFallbackException( | ||
| "Cannot revert a stage whose stripped root is a stage boundary") | ||
| } | ||
| val reverted = transformStageUp(stripped) { case cometExec: CometExec => | ||
| cometExec.sparkFallback(cometExec.children) | ||
| } | ||
| insertTransitions(reverted) | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This says the rule rebuilds the Spark operator through
sparkFallback, butrevertToSparkhandlesCometLocalTopKExec,CometNativeScanExecandCometIcebergNativeScanExecitself before it gets there (RevertNativeForTransitionHeavyStages.scala:263-281). SosparkFallbackisn't the whole contract yet, and onCometLocalTopKExecthe default would return a secondTakeOrderedAndProjectExecthat applies the TopK twice. Could those three becomesparkFallbackoverrides, with the local TopK returning its child and the two scans carrying their live DPP filters across? Then the rule only ever callssparkFallback, and this section holds for the next operator whose live state differs from itsoriginalPlan. If you'd rather not move the code, could the section name the exceptions?