viirya commented on code in PR #6852:
URL: https://github.com/apache/datafusion-comet/pull/6852#discussion_r4238845482
##########
spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala:
##########
@@ -65,19 +65,27 @@ case class RevertNativeForTransitionHeavyStages(session:
SparkSession, wholePlan
case exchange: ShuffleExchangeLike =>
revertShuffleStageIfNeeded(exchange).getOrElse(plan)
case _ =>
- // Result stage: its output is collected as rows.
- revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan)
+ revertResultStageIfNeeded(plan).getOrElse(plan)
}
}
private def applyForNonAQE(plan: SparkPlan): SparkPlan = {
val withRevertedStages = plan.transformUp { case exchange:
ShuffleExchangeLike =>
revertShuffleStageIfNeeded(exchange).getOrElse(exchange)
}
- revertStageIfNeeded(withRevertedStages, outputColumnar = false)
- .getOrElse(withRevertedStages)
+ revertResultStageIfNeeded(withRevertedStages).getOrElse(withRevertedStages)
}
+ /**
+ * Reverts the result stage if needed, keeping the output format its
consumer reads. That is
+ * usually rows, but on Spark 4.0+, when the cache serializer accepts
columnar input, as Comet's
+ * does, Spark marks a cached `AdaptiveSparkPlanExec` columnar and AQE plans
its final stage to
+ * output columnar batches. Only then does Spark's transition insertion
leave a columnar-only
+ * root; when rows are wanted it puts a columnar-to-row transition on top.
+ */
+ private def revertResultStageIfNeeded(plan: SparkPlan): Option[SparkPlan] =
+ revertStageIfNeeded(plan, outputColumnar = plan.supportsColumnar &&
!plan.supportsRowBased)
Review Comment:
What should happen when the root supports both formats? With
`supportsColumnar && supportsRowBased` this infers rows. But when columnar
output was requested, Spark's `ensureOutputsColumnar` keeps such a root as is,
so the root alone doesn't say which format the consumer wants.
Up to 4.1 I couldn't find a way to reach this. The only dual-format
operators there are `UnionExec` and `InMemoryTableScanExec`, and neither can
have a C2R below it in the same stage. On Spark 4.2, though,
`ArrowEvalPythonExec` reports `supportsRowBased = true` and `supportsColumnar =
child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled`, and that
conf defaults to true. Take a cached query whose result stage is rooted at a
pandas UDF over a columnar child with C2Rs below it, for example with
`spark.comet.sparkToColumnar.enabled`. The reverted `ArrowEvalPythonExec`
becomes row-only, and the cache builder's `executeColumnar()` would hit the
same mismatch this PR fixes.
Would it make sense to skip the revert when the original root supports both
formats and the reverted root no longer supports columnar? The doc comment
above could also note that a dual-format root leaves the requested format
unknown.
--
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]
