Fix ras and only fallback bhj

Author: liujiayi771

gluten-core/src/main/scala/org/apache/gluten/extension/columnar/heuristic/RewriteSparkPlanRulesManager.scala

@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.{LeafExecNode, ProjectExec, SparkPlan}
-import org.apache.spark.sql.execution.joins.BaseJoinExec
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
 
 case class RewrittenNodeWall(originalChild: SparkPlan) extends LeafExecNode {
   override protected def doExecute(): RDD[InternalRow] = throw new UnsupportedOperationException()
@@ -109,7 +109,9 @@ class RewriteSparkPlanRulesManager private (
             FallbackTags.add(origin, FallbackTags.getOption(rewriteNode).get)
             origin
           } else if (
-            rewriteNode.isInstanceOf[BaseJoinExec] && allFallbackTags.exists(_.isDefined)
+            (rewriteNode.isInstanceOf[BroadcastHashJoinExec] ||
+              rewriteNode.isInstanceOf[BroadcastNestedLoopJoinExec]) &&
+            allFallbackTags.exists(_.isDefined)
           ) {
             // If the inserted projects for join is not transformable, return the original plan.
             val reason = allFallbackTags.collect { case Some(s) => s.reason() }.mkString(", ")

gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/enumerated/RasOffload.scala

@@ -18,7 +18,7 @@ package org.apache.gluten.extension.columnar.enumerated
 
 import org.apache.gluten.execution.{GlutenPlan, ValidatablePlan}
 import org.apache.gluten.extension.columnar.FallbackTags
-import org.apache.gluten.extension.columnar.offload.OffloadSingleNode
+import org.apache.gluten.extension.columnar.offload.{OffloadOthers, OffloadSingleNode}
 import org.apache.gluten.extension.columnar.rewrite.RewriteSingleNode
 import org.apache.gluten.extension.columnar.validator.Validator
 import org.apache.gluten.ras.path.Pattern
@@ -27,7 +27,8 @@ import org.apache.gluten.ras.rule.{RasRule, Shape}
 import org.apache.gluten.ras.rule.Shapes.pattern
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
 
 import scala.reflect.{classTag, ClassTag}
 
@@ -129,7 +130,21 @@ object RasOffload {
                   case t: ValidatablePlan => t
                 }
                 val outComes = offloadedNodes.map(_.doValidate()).filter(!_.ok())
-                if (outComes.nonEmpty) {
+                // 4.1 Validate pre project of broadcast join
+                val notOffload = from match {
+                  case _: BroadcastHashJoinExec | _: BroadcastNestedLoopJoinExec =>
+                    val projectOffload = RasOffload.from[ProjectExec](OffloadOthers())
+                    from
+                      .collect {
+                        case preProject: ProjectExec => projectOffload.offload(preProject)
+                      }
+                      .exists {
+                        case t: ValidatablePlan => !t.doValidate().ok()
+                        case plan if !plan.isInstanceOf[GlutenPlan] => true
+                      }
+                  case _ => false
+                }
+                if (outComes.nonEmpty || notOffload) {
                   // 5. If native validation fails on at least one of the offloaded nodes, return
                   // the original one.
                   //

gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/PullOutPreProject.scala

@@ -17,7 +17,6 @@
 package org.apache.gluten.extension.columnar.rewrite
 
 import org.apache.gluten.backendsapi.BackendsApiManager
-import org.apache.gluten.config.GlutenConfig
 import org.apache.gluten.extension.columnar.heuristic.RewrittenNodeWall
 import org.apache.gluten.sql.shims.SparkShimLoader
 import org.apache.gluten.utils.PullOutProjectHelper
@@ -26,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, Partial}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, TypedAggregateExpression}
-import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin}
+import org.apache.spark.sql.execution.joins.{BaseJoinExec, HashJoin}
 import org.apache.spark.sql.execution.python.ArrowEvalPythonExec
 import org.apache.spark.sql.execution.window.WindowExec
 
@@ -295,17 +294,6 @@ object PullOutPreProject extends RewriteSingleNode with PullOutProjectHelper {
         arrowEvalPythonExec)
 
     case join: BaseJoinExec if needsPreProject(join) =>
-      join match {
-        case _: BroadcastHashJoinExec | _: BroadcastNestedLoopJoinExec
-            if !GlutenConfig.get.enableColumnarProject =>
-          // If columnar project is disabled, we cannot pull out project for join, since ProjectExec
-          // not override doExecuteBroadcast methods, we cannot add project between broadcast join
-          // and broadcast exchange.
-          throw new UnsupportedOperationException("columnar project is disabled, " +
-            "broadcast join operator does not support pull out pre-project, and it will fallback.")
-        case _ =>
-      }
-
       // Spark has an improvement which would patch integer joins keys to a Long value.
       // But this improvement would cause adding extra project before hash join in velox,
       // disabling this improvement as below would help reduce the project.

gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala

@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete, Partial}
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.aggregate._
-import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.types.{ByteType, DateType, IntegerType, LongType, ShortType}
 
@@ -175,6 +175,22 @@ trait PullOutProjectHelper {
         condition = newCondition)
       newSmj.copyTagsFrom(smj)
       newSmj
+    case nestedLoopJoin: BroadcastNestedLoopJoinExec =>
+      val newNestedLoopJoin = nestedLoopJoin.copy(
+        left = newLeft,
+        right = newRight,
+        condition = newCondition
+      )
+      newNestedLoopJoin.copyTagsFrom(nestedLoopJoin)
+      newNestedLoopJoin
+    case cartesianProduct: CartesianProductExec =>
+      val newCartesianProduct = cartesianProduct.copy(
+        left = newLeft,
+        right = newRight,
+        condition = newCondition
+      )
+      newCartesianProduct.copyTagsFrom(cartesianProduct)
+      newCartesianProduct
     case _ =>
       throw new UnsupportedOperationException(s"Unsupported join $join")
   }