make ANSI fallback more granular

apache · Jun 3, 2024 · 84f6f70 · 84f6f70
1 parent edd63ef
commit 84f6f70
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git 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
@@ -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 =>