add expression test

apache · Oct 7, 2023 · 32abfcf · 32abfcf
1 parent e7f37d8
commit 32abfcf
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/...k/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/GlutenPlanManager.scala b/...k/kyuubi-extension-spark-3-3/src/main/scala/org/apache/kyuubi/sql/GlutenPlanManager.scala
@@ -19,7 +19,7 @@ package org.apache.kyuubi.sql
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.UserDefinedExpression
+import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayIntersect, ArraySort, Bin, Contains, EndsWith, LastDay, MakeDate, Overlay, Rand, Randn, Size, SortArray, StartsWith, ToUnixTimestamp, UnaryMinus, UnixTimestamp, UserDefinedExpression}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.aggregate.{ObjectHashAggregateExec, SortAggregateExec}
@@ -56,7 +56,6 @@ object GlutenPlanAnalysis extends Rule[SparkPlan] {
           if !p.relation.fileFormat.isInstanceOf[ParquetFileFormat] =>
         true
       case _: RowDataSourceScanExec |
-          _: UserDefinedExpression |
           _: CartesianProductExec |
           _: ShuffleExchangeExec |
           _: ObjectHashAggregateExec |
@@ -67,6 +66,30 @@ object GlutenPlanAnalysis extends Rule[SparkPlan] {
           _: SampleExec |
           _: BroadcastNestedLoopJoinExec =>
         true
+      case p: SparkPlan
+          if p.expressions.exists(e =>
+            e.exists {
+              case _: UserDefinedExpression |
+                  _: UnaryMinus |
+                  _: Bin |
+                  _: Contains |
+                  _: StartsWith |
+                  _: EndsWith |
+                  _: Overlay |
+                  _: Rand |
+                  _: Randn |
+                  _: ArrayContains |
+                  _: ArrayIntersect |
+                  _: ArraySort |
+                  _: SortArray |
+                  _: Size |
+                  _: LastDay |
+                  _: MakeDate |
+                  _: ToUnixTimestamp |
+                  _: UnixTimestamp => true
+              case _ => false
+            }) =>
+        true
       // TODO check whether the plan contains unsupported expressions
     }.size
     check(count)

diff --git a/...ubi-extension-spark-3-3/src/test/scala/org/apache/kyuubi/sql/GlutenPlanManagerSuite.scala b/...ubi-extension-spark-3-3/src/test/scala/org/apache/kyuubi/sql/GlutenPlanManagerSuite.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.KyuubiSparkSQLExtensionTest
 
 class GlutenPlanManagerSuite extends KyuubiSparkSQLExtensionTest {
 
-  test("Kyuubi Extension fast fail if over un-supported operator threshold") {
+  test("Kyuubi Extension fast fail with Plan if over un-supported operator threshold") {
     withSQLConf(
       KyuubiSQLConf.GLUTEN_FALLBACK_OPERATOR_THRESHOLD.key -> "1") {
       withTable("gluten_tmp_1") {
@@ -36,4 +36,17 @@ class GlutenPlanManagerSuite extends KyuubiSparkSQLExtensionTest {
       }
     }
   }
+
+  test("Kyuubi Extension fast fail with Expression if over un-supported operator threshold") {
+    withSQLConf(KyuubiSQLConf.GLUTEN_FALLBACK_OPERATOR_THRESHOLD.key -> "1") {
+      assertThrows[TooMuchGlutenUnsupportedOperationException] {
+        sql("SELECT rand(100)").collect()
+      }
+
+      spark.udf.register("str_len", (s: String) => s.length)
+      assertThrows[TooMuchGlutenUnsupportedOperationException] {
+        sql("SELECT str_len('123')").collect()
+      }
+    }
+  }
 }