Changed extractExecutedPlan to consider ResultQueryStageExec for Da…

…tabricks 13.3 [databricks] (#9690) * Changed extractExecutedPlan to consider ResultQueryStageExec * Signing off Signed-off-by: Raza Jafri <[email protected]> --------- Signed-off-by: Raza Jafri <[email protected]> Co-authored-by: raza <[email protected]>
NVIDIA · Nov 14, 2023 · 0f0a0ed · 0f0a0ed
1 parent f120c48
commit 0f0a0ed
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 1 deletion.
diff --git a/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/PlanShimsImpl.scala b/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/PlanShimsImpl.scala
@@ -32,7 +32,6 @@
 {"spark": "333"}
 {"spark": "340"}
 {"spark": "341"}
-{"spark": "341db"}
 {"spark": "350"}
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims

diff --git a/sql-plugin/src/main/spark341db/scala/com/nvidia/spark/rapids/shims/PlanShimsImpl.scala b/sql-plugin/src/main/spark341db/scala/com/nvidia/spark/rapids/shims/PlanShimsImpl.scala
@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "341db"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+import com.nvidia.spark.rapids.{GpuAlias, PlanShims}
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Expression}
+import org.apache.spark.sql.execution.{CommandResultExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.ResultQueryStageExec
+
+class PlanShimsImpl extends PlanShims {
+  def extractExecutedPlan(plan: SparkPlan): SparkPlan = plan match {
+    case p: CommandResultExec => p.commandPhysicalPlan
+    case q: ResultQueryStageExec => q.plan
+    case _ => plan
+  }
+
+  def isAnsiCast(e: Expression): Boolean = AnsiCastShim.isAnsiCast(e)
+
+  def isAnsiCastOptionallyAliased(e: Expression): Boolean = e match {
+    case Alias(e, _) => isAnsiCast(e)
+    case GpuAlias(e, _) => isAnsiCast(e)
+    case e => isAnsiCast(e)
+  }
+}