fall back for more expressions in ANSI mode

apache · Jun 6, 2024 · 901a622 · 901a622
1 parent 9688cc6
commit 901a622
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 0 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -2207,9 +2207,13 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     case expr: Subtract => evalMode(expr) == CometEvalMode.ANSI
     case expr: Multiply => evalMode(expr) == CometEvalMode.ANSI
     case expr: Divide => evalMode(expr) == CometEvalMode.ANSI
+    case expr: IntegralDivide => evalMode(expr) == CometEvalMode.ANSI
     case expr: Remainder => evalMode(expr) == CometEvalMode.ANSI
     case expr: Pmod => evalMode(expr) == CometEvalMode.ANSI
     case expr: Round => evalMode(expr) == CometEvalMode.ANSI
+    case expr: BRound => evalMode(expr) == CometEvalMode.ANSI
+    case expr: Sum => evalMode(expr) == CometEvalMode.ANSI
+    case expr: Average => evalMode(expr) == CometEvalMode.ANSI
     case _ => false
   }
 

diff --git a/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala
@@ -20,6 +20,7 @@ package org.apache.comet.shims
 
 import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Sum}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -44,13 +45,24 @@ trait CometExprShim {
     def evalMode(expr: Divide): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
+    def evalMode(expr: IntegralDivide): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
     def evalMode(expr: Remainder): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
     def evalMode(expr: Pmod): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
+    def evalMode(expr: Sum): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Average): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
     def evalMode(c: Cast): CometEvalMode.Value = CometEvalMode.fromBoolean(c.ansiEnabled)
 
     def evalMode(r: Round): CometEvalMode.Value = CometEvalMode.LEGACY
+
+    def evalMode(r: BRound): CometEvalMode.Value = CometEvalMode.LEGACY
 }
diff --git a/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala
@@ -20,6 +20,7 @@ package org.apache.comet.shims
 
 import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Sum}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -44,13 +45,24 @@ trait CometExprShim {
     def evalMode(expr: Divide): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
+    def evalMode(expr: IntegralDivide): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
     def evalMode(expr: Remainder): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
     def evalMode(expr: Pmod): CometEvalMode.Value =
         CometEvalMode.fromBoolean(expr.failOnError)
 
+    def evalMode(expr: Sum): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Average): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
     def evalMode(c: Cast): CometEvalMode.Value = CometEvalMode.fromBoolean(c.ansiEnabled)
 
     def evalMode(r: Round): CometEvalMode.Value = CometEvalMode.LEGACY
+
+    def evalMode(r: BRound): CometEvalMode.Value = CometEvalMode.LEGACY
 }
diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
@@ -20,6 +20,7 @@ package org.apache.comet.shims
 
 import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Sum}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -44,16 +45,27 @@ trait CometExprShim {
     def evalMode(expr: Divide): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
+    def evalMode(expr: IntegralDivide): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
     def evalMode(expr: Remainder): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
     def evalMode(expr: Pmod): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
+    def evalMode(expr: Sum): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
+    def evalMode(expr: Average): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
     def evalMode(c: Cast): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 
     def evalMode(r: Round): CometEvalMode.Value = CometEvalMode.fromBoolean(r.ansiEnabled)
+
+    def evalMode(r: BRound): CometEvalMode.Value = CometEvalMode.fromBoolean(r.ansiEnabled)
 }
 
 object CometEvalModeUtil {

diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
@@ -20,6 +20,7 @@ package org.apache.comet.shims
 
 import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Sum}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -44,17 +45,29 @@ trait CometExprShim {
     def evalMode(expr: Divide): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
+    def evalMode(expr: IntegralDivide): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
     def evalMode(expr: Remainder): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
     def evalMode(expr: Pmod): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
 
+    def evalMode(expr: Sum): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
+    def evalMode(expr: Average): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(expr.evalMode)
+
+
     def evalMode(c: Cast): CometEvalMode.Value =
         CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 
     def evalMode(r: Round): CometEvalMode.Value = CometEvalMode.fromBoolean(r.ansiEnabled)
 
+    def evalMode(r: BRound): CometEvalMode.Value = CometEvalMode.fromBoolean(r.ansiEnabled)
+
 }
 
 object CometEvalModeUtil {