chore: Make ANSI fallback more granular

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -383,15 +383,6 @@ object CometConf extends ShimCometConf {
       .toSequence
       .createWithDefault(Seq("Range,InMemoryTableScan"))
 
-  val COMET_ANSI_MODE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.ansi.enabled")
-    .internal()
-    .doc(
-      "Comet does not respect ANSI mode in most cases and by default will not accelerate " +
-        "queries when ansi mode is enabled. Enable this setting to test Comet's experimental " +
-        "support for ANSI mode. This should not be used in production.")
-    .booleanConf
-    .createWithDefault(COMET_ANSI_MODE_ENABLED_DEFAULT)
-
   val COMET_CAST_ALLOW_INCOMPATIBLE: ConfigEntry[Boolean] =
     conf("spark.comet.cast.allowIncompatible")
       .doc(

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -45,7 +45,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 import org.apache.comet.CometConf._
-import org.apache.comet.CometSparkSessionExtensions.{createMessage, getCometShuffleNotEnabledReason, isANSIEnabled, isCometBroadCastForceEnabled, isCometEnabled, isCometExecEnabled, isCometJVMShuffleMode, isCometNativeShuffleMode, isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled, isSchemaSupported, isSpark34Plus, shouldApplyRowToColumnar, withInfo, withInfos}
+import org.apache.comet.CometSparkSessionExtensions.{createMessage, getCometShuffleNotEnabledReason, isCometBroadCastForceEnabled, isCometEnabled, isCometExecEnabled, isCometJVMShuffleMode, isCometNativeShuffleMode, isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled, isSchemaSupported, isSpark34Plus, shouldApplyRowToColumnar, withInfo, withInfos}
 import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
 import org.apache.comet.serde.OperatorOuterClass.Operator
 import org.apache.comet.serde.QueryPlanSerde
@@ -712,17 +712,6 @@ class CometSparkSessionExtensions
     }
 
     override def apply(plan: SparkPlan): SparkPlan = {
-      // DataFusion doesn't have ANSI mode. For now we just disable CometExec if ANSI mode is
-      // enabled.
-      if (isANSIEnabled(conf)) {
-        if (COMET_ANSI_MODE_ENABLED.get()) {
-          logWarning("Using Comet's experimental support for ANSI mode.")
-        } else {
-          logInfo("Comet extension disabled for ANSI mode")
-          return plan
-        }
-      }
-
       // We shouldn't transform Spark query plan if Comet is disabled.
       if (!isCometEnabled(conf)) return plan
 

spark/src/main/scala/org/apache/comet/GenerateDocs.scala

@@ -25,7 +25,7 @@ import scala.io.Source
 
 import org.apache.spark.sql.catalyst.expressions.Cast
 
-import org.apache.comet.expressions.{CometCast, Compatible, Incompatible}
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible, Incompatible}
 
 /**
  * Utility for generating markdown documentation from the configs.
@@ -72,7 +72,7 @@ object GenerateDocs {
             if (Cast.canCast(fromType, toType) && fromType != toType) {
               val fromTypeName = fromType.typeName.replace("(10,2)", "")
               val toTypeName = toType.typeName.replace("(10,2)", "")
-              CometCast.isSupported(fromType, toType, None, "LEGACY") match {
+              CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) match {
                 case Compatible(notes) =>
                   val notesStr = notes.getOrElse("").trim
                   w.write(s"| $fromTypeName | $toTypeName | $notesStr |\n".getBytes)
@@ -89,7 +89,7 @@ object GenerateDocs {
             if (Cast.canCast(fromType, toType) && fromType != toType) {
               val fromTypeName = fromType.typeName.replace("(10,2)", "")
               val toTypeName = toType.typeName.replace("(10,2)", "")
-              CometCast.isSupported(fromType, toType, None, "LEGACY") match {
+              CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) match {
                 case Incompatible(notes) =>
                   val notesStr = notes.getOrElse("").trim
                   w.write(s"| $fromTypeName | $toTypeName  | $notesStr |\n".getBytes)

spark/src/main/scala/org/apache/comet/expressions/CometCast.scala

@@ -55,7 +55,7 @@ object CometCast {
       fromType: DataType,
       toType: DataType,
       timeZoneId: Option[String],
-      evalMode: String): SupportLevel = {
+      evalMode: CometEvalMode.Value): SupportLevel = {
 
     if (fromType == toType) {
       return Compatible()
@@ -102,7 +102,7 @@ object CometCast {
   private def canCastFromString(
       toType: DataType,
       timeZoneId: Option[String],
-      evalMode: String): SupportLevel = {
+      evalMode: CometEvalMode.Value): SupportLevel = {
     toType match {
       case DataTypes.BooleanType =>
         Compatible()

spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+package org.apache.comet.expressions
+
+/**
+ * We cannot reference Spark's EvalMode directly because the package is different between Spark
+ * versions, so we copy it here.
+ *
+ * Expression evaluation modes.
+ *   - LEGACY: the default evaluation mode, which is compliant to Hive SQL.
+ *   - ANSI: a evaluation mode which is compliant to ANSI SQL standard.
+ *   - TRY: a evaluation mode for `try_*` functions. It is identical to ANSI evaluation mode
+ *     except for returning null result on errors.
+ */
+object CometEvalMode extends Enumeration {
+  val LEGACY, ANSI, TRY = Value
+
+  def fromBoolean(ansiEnabled: Boolean): Value = if (ansiEnabled) {
+    ANSI
+  } else {
+    LEGACY
+  }
+
+  def fromString(str: String): CometEvalMode.Value = CometEvalMode.withName(str)
+}

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -19,8 +19,6 @@
 
 package org.apache.comet.serde
 
-import java.util.Locale
-
 import scala.collection.JavaConverters._
 
 import org.apache.spark.internal.Logging
@@ -45,7 +43,7 @@ import org.apache.spark.unsafe.types.UTF8String
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.{isCometOperatorEnabled, isCometScan, isSpark32, isSpark34Plus, withInfo}
-import org.apache.comet.expressions.{CometCast, Compatible, Incompatible, Unsupported}
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible, Incompatible, Unsupported}
 import org.apache.comet.serde.ExprOuterClass.{AggExpr, DataType => ProtoDataType, Expr, ScalarFunc}
 import org.apache.comet.serde.ExprOuterClass.DataType.{DataTypeInfo, DecimalInfo, ListInfo, MapInfo, StructInfo}
 import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, JoinType, Operator}
@@ -578,6 +576,15 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     }
   }
 
+  def evalModeToProto(evalMode: CometEvalMode.Value): ExprOuterClass.EvalMode = {
+    evalMode match {
+      case CometEvalMode.LEGACY => ExprOuterClass.EvalMode.LEGACY
+      case CometEvalMode.TRY => ExprOuterClass.EvalMode.TRY
+      case CometEvalMode.ANSI => ExprOuterClass.EvalMode.ANSI
+      case _ => throw new IllegalStateException(s"Invalid evalMode $evalMode")
+    }
+  }
+
   /**
    * Convert a Spark expression to protobuf.
    *
@@ -590,18 +597,6 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
    * @return
    *   The protobuf representation of the expression, or None if the expression is not supported
    */
-
-  def stringToEvalMode(evalModeStr: String): ExprOuterClass.EvalMode =
-    evalModeStr.toUpperCase(Locale.ROOT) match {
-      case "LEGACY" => ExprOuterClass.EvalMode.LEGACY
-      case "TRY" => ExprOuterClass.EvalMode.TRY
-      case "ANSI" => ExprOuterClass.EvalMode.ANSI
-      case invalid =>
-        throw new IllegalArgumentException(
-          s"Invalid eval mode '$invalid' "
-        ) // Assuming we want to catch errors strictly
-    }
-
   def exprToProto(
       expr: Expression,
       input: Seq[Attribute],
@@ -610,15 +605,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         timeZoneId: Option[String],
         dt: DataType,
         childExpr: Option[Expr],
-        evalMode: String): Option[Expr] = {
+        evalMode: CometEvalMode.Value): Option[Expr] = {
       val dataType = serializeDataType(dt)
-      val evalModeEnum = stringToEvalMode(evalMode) // Convert string to enum
 
       if (childExpr.isDefined && dataType.isDefined) {
         val castBuilder = ExprOuterClass.Cast.newBuilder()
         castBuilder.setChild(childExpr.get)
         castBuilder.setDatatype(dataType.get)
-        castBuilder.setEvalMode(evalModeEnum) // Set the enum in protobuf
+        castBuilder.setEvalMode(evalModeToProto(evalMode))
 
         val timeZone = timeZoneId.getOrElse("UTC")
         castBuilder.setTimezone(timeZone)
@@ -646,26 +640,26 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           inputs: Seq[Attribute],
           dt: DataType,
           timeZoneId: Option[String],
-          actualEvalModeStr: String): Option[Expr] = {
+          evalMode: CometEvalMode.Value): Option[Expr] = {
 
         val childExpr = exprToProtoInternal(child, inputs)
         if (childExpr.isDefined) {
           val castSupport =
-            CometCast.isSupported(child.dataType, dt, timeZoneId, actualEvalModeStr)
+            CometCast.isSupported(child.dataType, dt, timeZoneId, evalMode)
 
           def getIncompatMessage(reason: Option[String]): String =
             "Comet does not guarantee correct results for cast " +
               s"from ${child.dataType} to $dt " +
-              s"with timezone $timeZoneId and evalMode $actualEvalModeStr" +
+              s"with timezone $timeZoneId and evalMode $evalMode" +
               reason.map(str => s" ($str)").getOrElse("")
 
           castSupport match {
             case Compatible(_) =>
-              castToProto(timeZoneId, dt, childExpr, actualEvalModeStr)
+              castToProto(timeZoneId, dt, childExpr, evalMode)
             case Incompatible(reason) =>
               if (CometConf.COMET_CAST_ALLOW_INCOMPATIBLE.get()) {
                 logWarning(getIncompatMessage(reason))
-                castToProto(timeZoneId, dt, childExpr, actualEvalModeStr)
+                castToProto(timeZoneId, dt, childExpr, evalMode)
               } else {
                 withInfo(
                   expr,
@@ -677,7 +671,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
               withInfo(
                 expr,
                 s"Unsupported cast from ${child.dataType} to $dt " +
-                  s"with timezone $timeZoneId and evalMode $actualEvalModeStr")
+                  s"with timezone $timeZoneId and evalMode $evalMode")
               None
           }
         } else {
@@ -701,17 +695,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
 
         case UnaryExpression(child) if expr.prettyName == "trycast" =>
           val timeZoneId = SQLConf.get.sessionLocalTimeZone
-          handleCast(child, inputs, expr.dataType, Some(timeZoneId), "TRY")
+          handleCast(child, inputs, expr.dataType, Some(timeZoneId), CometEvalMode.TRY)
 
-        case Cast(child, dt, timeZoneId, evalMode) =>
-          val evalModeStr = if (evalMode.isInstanceOf[Boolean]) {
-            // Spark 3.2 & 3.3 has ansiEnabled boolean
-            if (evalMode.asInstanceOf[Boolean]) "ANSI" else "LEGACY"
-          } else {
-            // Spark 3.4+ has EvalMode enum with values LEGACY, ANSI, and TRY
-            evalMode.toString
-          }
-          handleCast(child, inputs, dt, timeZoneId, evalModeStr)
+        case c @ Cast(child, dt, timeZoneId, _) =>
+          handleCast(child, inputs, dt, timeZoneId, evalMode(c))
+
+        case expr if isUnsupportedAnsiExpr(expr) =>
+          withInfo(expr, "ANSI mode not supported")
+          None
 
         case add @ Add(left, right, _) if supportedDataType(left.dataType) =>
           val leftExpr = exprToProtoInternal(left, inputs)
@@ -2006,7 +1997,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           // TODO: Remove this once we have new DataFusion release which includes
           // the fix: https://github.com/apache/arrow-datafusion/pull/9459
           if (childExpr.isDefined) {
-            castToProto(None, a.dataType, childExpr, "LEGACY")
+            castToProto(None, a.dataType, childExpr, CometEvalMode.LEGACY)
           } else {
             withInfo(expr, a.children: _*)
             None
@@ -2207,6 +2198,22 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     exprToProtoInternal(newExpr, input)
   }
 
+  /** These expressions are not compatible with Spark in ANSI mode yet */
+  private def isUnsupportedAnsiExpr(expr: Expression): Boolean = expr match {
+    case expr: Add => evalMode(expr) == CometEvalMode.ANSI
+    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
+  }
+
   def scalarExprToProtoWithReturnType(
       funcName: String,
       returnType: DataType,

spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala

@@ -18,7 +18,9 @@
  */
 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.
@@ -30,4 +32,37 @@ trait CometExprShim {
     def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    def evalMode(expr: Add): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Subtract): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Multiply): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    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
 }

spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala

@@ -18,7 +18,9 @@
  */
 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.
@@ -30,4 +32,37 @@ trait CometExprShim {
     def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    def evalMode(expr: Add): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Subtract): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    def evalMode(expr: Multiply): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.failOnError)
+
+    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.useAnsiAdd)
+
+    def evalMode(expr: Average): CometEvalMode.Value =
+        CometEvalMode.fromBoolean(expr.useAnsiAdd)
+
+    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
 }