Delta Lake and AQE on Databricks 10.4 workaround [databricks] (#6587)

* Handle Delta and AQE workaround around on Databricks 10.4 due to bug Signed-off-by: Thomas Graves <[email protected]> * Update to use SQLConf Signed-off-by: Thomas Graves <[email protected]> * cleanup switch to SQLConf * fix line length Signed-off-by: Thomas Graves <[email protected]>
NVIDIA · Sep 20, 2022 · bc0aa04 · bc0aa04
1 parent 998c76e
commit bc0aa04
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 5 deletions.
diff --git a/sql-plugin/src/main/311+-nondb/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala b/sql-plugin/src/main/311+-nondb/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala
@@ -18,6 +18,7 @@ package com.nvidia.spark.rapids.shims
 
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.adaptive.{QueryStageExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.internal.SQLConf
 
 /** Utility methods for manipulating Catalyst classes involved in Adaptive Query Execution */
 object AQEUtils {
@@ -26,5 +27,5 @@ object AQEUtils {
     sqse.newReuseInstance(sqse.id, newOutput)
   }
 
-  def isAdaptiveExecutionSupportedInSparkVersion: Boolean = true
+  def isAdaptiveExecutionSupportedInSparkVersion(conf: SQLConf): Boolean = true
 }
diff --git a/sql-plugin/src/main/312db/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala b/sql-plugin/src/main/312db/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala
@@ -19,6 +19,7 @@ package com.nvidia.spark.rapids.shims
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.adaptive.{QueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.internal.SQLConf
 
 /** Utility methods for manipulating Catalyst classes involved in Adaptive Query Execution */
 object AQEUtils {
@@ -28,5 +29,5 @@ object AQEUtils {
     ShuffleQueryStageExec(sqse.id, reusedExchange, sqse.originalPlan)
   }
 
-  def isAdaptiveExecutionSupportedInSparkVersion: Boolean = true
+  def isAdaptiveExecutionSupportedInSparkVersion(conf: SQLConf): Boolean = true
 }
diff --git a/sql-plugin/src/main/321db/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala b/sql-plugin/src/main/321db/scala/com/nvidia/spark/rapids/shims/AQEUtils.scala
@@ -19,6 +19,7 @@ package com.nvidia.spark.rapids.shims
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.adaptive.{QueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.internal.SQLConf
 
 /** Utility methods for manipulating Catalyst classes involved in Adaptive Query Execution */
 object AQEUtils {
@@ -28,5 +29,10 @@ object AQEUtils {
     ShuffleQueryStageExec(sqse.id, reusedExchange, sqse.originalPlan, sqse.isSparkExchange)
   }
 
-  def isAdaptiveExecutionSupportedInSparkVersion: Boolean = true
+  // Databricks 10.4 has an issue where if you turn off AQE it can still use it for
+  // certain operations. This causes issues with the plugin so this is to work around
+  // that.
+  def isAdaptiveExecutionSupportedInSparkVersion(conf: SQLConf): Boolean = {
+    conf.getConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED)
+  }
 }
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
@@ -4421,10 +4421,12 @@ case class GpuOverrides() extends Rule[SparkPlan] with Logging {
           logDebug(s"Fallback for RDDScanExec delta log: $rdd")
         }
         found
-      case aqe: AdaptiveSparkPlanExec if !AQEUtils.isAdaptiveExecutionSupportedInSparkVersion =>
+      case aqe: AdaptiveSparkPlanExec if 
+        !AQEUtils.isAdaptiveExecutionSupportedInSparkVersion(plan.conf) =>
         logDebug(s"AdaptiveSparkPlanExec found on unsupported Spark Version: $aqe")
         true
-      case project: ProjectExec if !AQEUtils.isAdaptiveExecutionSupportedInSparkVersion =>
+      case project: ProjectExec if
+        !AQEUtils.isAdaptiveExecutionSupportedInSparkVersion(plan.conf) =>
         val foundExprs = project.expressions.flatMap { e =>
           PlanUtils.findExpressions(e, {
             case udf: ScalaUDF =>