fix add columnar to row

Author: wForget

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -20,7 +20,6 @@
 package org.apache.comet.rules
 
 import scala.collection.mutable.ListBuffer
-
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
@@ -30,16 +29,15 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec, CometShuffleManager}
 import org.apache.spark.sql.execution._
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AdaptiveSparkPlanExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.command.ExecutedCommandExec
 import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange, ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
-
 import org.apache.comet.{CometConf, ExtendedExplainInfo}
 import org.apache.comet.CometConf.COMET_EXEC_SHUFFLE_ENABLED
 import org.apache.comet.CometSparkSessionExtensions._
@@ -608,18 +606,24 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
         }
     }
 
-    plan.transformUp { case op =>
-      val newOp = convertNode(op)
-      // if newOp is not columnar and newOp.children has columnar, we need to add columnar to row
-      if (!newOp.supportsColumnar && !newOp.isInstanceOf[ColumnarToRowTransition]) {
-        val newChildren = newOp.children.map {
+    val newPlan = plan.transformUp { case op =>
+      convertNode(op)
+    }
+
+    // insert CometColumnarToRowExec if necessary
+    newPlan.transformUp {
+      case c2r: ColumnarToRowTransition => c2r
+      case op if !op.supportsColumnar =>
+        val newChildren = op.children.map {
+          // CometExec already handles columnar to row conversion internally
+          // Don't explicitly add CometColumnarToRowExec helps broadcast reuse,
+          // for plan like: BroadcastExchangeExec(CometExec)
+          case cometExec: CometExec => cometExec
           case c if c.supportsColumnar => CometColumnarToRowExec(c)
           case other => other
         }
-        newOp.withNewChildren(newChildren)
-      } else {
-        newOp
-      }
+        op.withNewChildren(newChildren)
+      case o => o
     }
   }