fallback on reused broadcast exchange

Signed-off-by: Yuan Zhou <yuan.zhou@intel.com>
PHILO-HE · Mar 20, 2022 · ddc8631 · ddc8631
1 parent cf38196
commit ddc8631
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/extension/ColumnarOverrides.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/extension/ColumnarOverrides.scala
@@ -209,8 +209,15 @@ case class ColumnarPreOverrides() extends Rule[SparkPlan] {
         right)
     case plan: BroadcastQueryStageExec =>
       logDebug(
-        s"Columnar Processing for ${plan.getClass} is currently supported, actual plan is ${plan.plan.getClass}.")
-      plan
+        s"Columnar Processing for ${plan.getClass} is currently supported, actual plan is ${plan.plan}.")
+      plan.plan match {
+        case ReusedExchangeExec(_, originalBroadcastPlan: ColumnarBroadcastExchangeAdaptor) =>
+          val newBroadcast = BroadcastExchangeExec(
+            originalBroadcastPlan.mode,
+            DataToArrowColumnarExec(plan.plan, 1))
+          SparkShimLoader.getSparkShims.newBroadcastQueryStageExec(plan.id, newBroadcast)
+        case other => plan
+      }
     case plan: BroadcastExchangeExec =>
       val child = replaceWithColumnarPlan(plan.child)
       logDebug(s"Columnar Processing for ${plan.getClass} is currently supported.")