andygrove commented on code in PR #2840:
URL: https://github.com/apache/datafusion-comet/pull/2840#discussion_r2583177929
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -296,57 +281,10 @@ case class CometExecRule(session: SparkSession) extends
Rule[SparkPlan] {
case s @ ShuffleQueryStageExec(_, ReusedExchangeExec(_, _:
CometShuffleExchangeExec), _) =>
newPlanWithProto(s, CometSinkPlaceHolder(_, s, s))
- // Native shuffle for Comet operators
case s: ShuffleExchangeExec =>
- val nativeShuffle: Option[SparkPlan] =
- if (nativeShuffleSupported(s)) {
- val newOp = operator2ProtoIfAllChildrenAreNative(s)
- newOp match {
- case Some(nativeOp) =>
- // Switch to use Decimal128 regardless of precision, since
Arrow native execution
- // doesn't support Decimal32 and Decimal64 yet.
- conf.setConfString(CometConf.COMET_USE_DECIMAL_128.key, "true")
- val cometOp = CometShuffleExchangeExec(s, shuffleType =
CometNativeShuffle)
- Some(CometSinkPlaceHolder(nativeOp, s, cometOp))
- case None =>
- None
- }
- } else {
- None
- }
-
- val nativeOrColumnarShuffle = if (nativeShuffle.isDefined) {
- nativeShuffle
- } else {
- // Columnar shuffle for regular Spark operators (not Comet) and
Comet operators
- // (if configured).
- // If the child of ShuffleExchangeExec is also a
ShuffleExchangeExec, we should not
- // convert it to CometColumnarShuffle,
- if (columnarShuffleSupported(s)) {
- val newOp = operator2Proto(s)
- newOp match {
- case Some(nativeOp) =>
- s.child match {
- case n if n.isInstanceOf[CometNativeExec] ||
!n.supportsColumnar =>
- val cometOp =
- CometShuffleExchangeExec(s, shuffleType =
CometColumnarShuffle)
- Some(CometSinkPlaceHolder(nativeOp, s, cometOp))
- case _ =>
- None
- }
- case None =>
- None
- }
- } else {
- None
- }
- }
-
- if (nativeOrColumnarShuffle.isDefined) {
- nativeOrColumnarShuffle.get
- } else {
- s
- }
+ // try native shuffle first, then columnar shuffle, then fall back to
Spark
+ // if neither are supported
+ tryNativeShuffle(s).orElse(tryColumnarShuffle(s)).getOrElse(s)
Review Comment:
This is much easier to comprehend, IMO
--
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]
