From 84f6f70f3751c0b14e985cd6c83d957593fb5f47 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 3 Jun 2024 16:02:35 -0600
Subject: [PATCH] make ANSI fallback more granular

---
 .../apache/comet/CometSparkSessionExtensions.scala  | 13 +------------
 .../org/apache/comet/serde/QueryPlanSerde.scala     |  8 ++++++++
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 168d2bb52..50a1a7868 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/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
 
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 5fe290cf6..8bb3cf8d0 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -1598,6 +1598,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           val optExpr = scalarExprToProto("pow", leftExpr, rightExpr)
           optExprWithInfo(optExpr, expr, left, right)
 
+        case r: Round if r.ansiEnabled && !CometConf.COMET_ANSI_MODE_ENABLED.get =>
+          // https://github.com/apache/datafusion-comet/issues/466
+          withInfo(
+            r,
+            "Round does not support ANSI mode. " +
+              s"Set ${CometConf.COMET_ANSI_MODE_ENABLED.key}=true to enable it anyway")
+          None
+
         // round function for Spark 3.2 does not allow negative round target scale. In addition,
         // it has different result precision/scale for decimals. Supporting only 3.3 and above.
         case r: Round if !isSpark32 =>