Check whether AQE is supported

oap-project · Jul 18, 2022 · 2a99d57 · 2a99d57
1 parent 0da5459
commit 2a99d57
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 20 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
@@ -562,29 +562,13 @@ case class ColumnarOverrideRules(session: SparkSession) extends ColumnarRule wit
   var originalPlan: SparkPlan = _
   var fallbacks = 0
 
-  private def supportAdaptive(plan: SparkPlan): Boolean = {
-    // TODO migrate dynamic-partition-pruning onto adaptive execution.
-    // Only QueryStage will have Exchange as Leaf Plan
-    val isLeafPlanExchange = plan match {
-      case e: Exchange => true
-      case other => false
-    }
-    isLeafPlanExchange || (SQLConf.get.adaptiveExecutionEnabled && (sanityCheck(plan) &&
-    !plan.logicalLink.exists(_.isStreaming) &&
-    !plan.expressions.exists(_.find(_.isInstanceOf[DynamicPruningSubquery]).isDefined) &&
-    plan.children.forall(supportAdaptive)))
-  }
-
-  private def sanityCheck(plan: SparkPlan): Boolean =
-    plan.logicalLink.isDefined
-
   override def preColumnarTransitions: Rule[SparkPlan] = plan => {
     if (columnarEnabled) {
       // According to Spark's Columnar.scala, the plan is tackled one by one.
       // By recording the original plan, we can easily let the whole stage
       // fallback at #postColumnarTransitions.
       originalPlan = plan
-      isSupportAdaptive = supportAdaptive(plan)
+      isSupportAdaptive = SparkShimLoader.getSparkShims.supportAdaptiveWithExchangeConsidered(plan)
       val rule = preOverrides
       rule.setAdaptiveSupport(isSupportAdaptive)
       rule(rowGuardOverrides(plan))
@@ -651,7 +635,7 @@ case class ColumnarOverrideRules(session: SparkSession) extends ColumnarRule wit
 
   override def postColumnarTransitions: Rule[SparkPlan] = plan => {
     if (columnarEnabled) {
-      if (SQLConf.get.adaptiveExecutionEnabled && fallbackWholeStage(plan)) {
+      if (isSupportAdaptive && fallbackWholeStage(plan)) {
         // BatchScan with ArrowScan initialized can still connect
         // to ColumnarToRow for transition.
         insertTransitions(originalPlan, false)

diff --git a/shims/common/src/main/scala/com/intel/oap/sql/shims/SparkShims.scala b/shims/common/src/main/scala/com/intel/oap/sql/shims/SparkShims.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec
 import org.apache.spark.sql.execution.datasources.OutputWriter
 import org.apache.spark.sql.execution.datasources.parquet.{ParquetFilters, ParquetOptions, ParquetReadSupport, VectorizedParquetRecordReader}
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleOrigin}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange, ShuffleOrigin}
 import org.apache.spark.sql.internal.SQLConf
 
 sealed abstract class ShimDescriptor
@@ -121,4 +121,20 @@ trait SparkShims {
   def getEndMapIndexOfCoalescedMapperPartitionSpec(spec: ShufflePartitionSpec): Int
 
   def getNumReducersOfCoalescedMapperPartitionSpec(spec: ShufflePartitionSpec): Int
+
+  def isLeafPlanExchange (plan: SparkPlan): Boolean = {
+    plan match {
+      case e: Exchange => true
+      case other => false
+    }
+  }
+
+  def sanityCheck(plan: SparkPlan): Boolean =
+    plan.logicalLink.isDefined
+
+  def supportAdaptive(plan: SparkPlan): Boolean
+
+  def supportAdaptiveWithExchangeConsidered(plan: SparkPlan): Boolean = {
+    isLeafPlanExchange(plan) || supportAdaptive(plan)
+  }
 }
diff --git a/shims/spark311/src/main/scala/com/intel/oap/sql/shims/spark311/Spark311Shims.scala b/shims/spark311/src/main/scala/com/intel/oap/sql/shims/spark311/Spark311Shims.scala
@@ -33,7 +33,7 @@ import org.apache.spark.shuffle.sort.SortShuffleWriter
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.{DynamicPruningSubquery, Expression}
 import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode, Partitioning}
 import org.apache.spark.sql.execution.{ShufflePartitionSpec, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, CustomShuffleReaderExec, ShuffleQueryStageExec}
@@ -205,4 +205,14 @@ class Spark311Shims extends SparkShims {
     throw new RuntimeException("This method should not be invoked in spark 3.1.")
   }
 
+  /**
+    * Ported from InsertAdaptiveSparkPlan.
+    */
+  override def supportAdaptive(plan: SparkPlan): Boolean = {
+    // TODO migrate dynamic-partition-pruning onto adaptive execution.
+    sanityCheck(plan) &&
+      !plan.logicalLink.exists(_.isStreaming) &&
+      !plan.expressions.exists(_.find(_.isInstanceOf[DynamicPruningSubquery]).isDefined) &&
+      plan.children.forall(supportAdaptive)
+  }
 }
diff --git a/shims/spark321/src/main/scala/com/intel/oap/sql/shims/spark321/Spark321Shims.scala b/shims/spark321/src/main/scala/com/intel/oap/sql/shims/spark321/Spark321Shims.scala
@@ -235,4 +235,11 @@ class Spark321Shims extends SparkShims {
     }
   }
 
+  /**
+    * Ported from InsertAdaptiveSparkPlan.
+    */
+  override def supportAdaptive(plan: SparkPlan): Boolean = {
+    sanityCheck(plan) && !plan.logicalLink.exists(_.isStreaming) &&
+      plan.children.forall(supportAdaptive)
+  }
 }