feat: Only allow incompatible cast expressions to run in comet if a config is enabled

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

@@ -1042,10 +1042,10 @@ object CometSparkSessionExtensions extends Logging {
    *   The node with information (if any) attached
    */
   def withInfo[T <: TreeNode[_]](node: T, info: String, exprs: T*): T = {
-    val exprInfo = exprs
-      .flatMap { e => Seq(e.getTagValue(CometExplainInfo.EXTENSION_INFO)) }
-      .flatten
-      .mkString("\n")
+    val exprInfo =
+      (Seq(node.getTagValue(CometExplainInfo.EXTENSION_INFO)) ++ exprs
+        .flatMap(e => Seq(e.getTagValue(CometExplainInfo.EXTENSION_INFO)))).flatten
+        .mkString("\n")
     if (info != null && info.nonEmpty && exprInfo.nonEmpty) {
       node.setTagValue(CometExplainInfo.EXTENSION_INFO, Seq(exprInfo, info).mkString("\n"))
     } else if (exprInfo.nonEmpty) {

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

@@ -0,0 +1,136 @@
+/*
+ * 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
+
+import org.apache.spark.sql.catalyst.expressions.Cast
+import org.apache.spark.sql.types.{DataType, DataTypes, DecimalType}
+
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.withInfo
+
+object CometCast {
+
+  def isSupported(
+      cast: Cast,
+      fromType: DataType,
+      toType: DataType,
+      timeZoneId: Option[String],
+      evalMode: String): Boolean = {
+    (fromType, toType) match {
+      case (DataTypes.StringType, _) =>
+        castFromStringSupported(cast, toType, timeZoneId, evalMode)
+      case (_, DataTypes.StringType) =>
+        castToStringSupported(cast, fromType, timeZoneId, evalMode)
+      case (DataTypes.TimestampType, _) =>
+        castFromTimestampSupported(cast, toType, timeZoneId, evalMode)
+      case (_: DecimalType, DataTypes.FloatType | DataTypes.DoubleType) => true
+      case (
+            DataTypes.BooleanType,
+            DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType |
+            DataTypes.LongType | DataTypes.FloatType | DataTypes.DoubleType) =>
+        true
+      case (
+            DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType | DataTypes.LongType,
+            DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
+            DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
+            DataTypes.DoubleType) =>
+        true
+      case (
+            DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType | DataTypes.LongType,
+            _: DecimalType) =>
+        true
+      case (DataTypes.FloatType, DataTypes.BooleanType | DataTypes.DoubleType) => true
+      case (DataTypes.DoubleType, DataTypes.BooleanType | DataTypes.FloatType) => true
+      case _ => false
+    }
+  }
+
+  private def castFromStringSupported(
+      cast: Cast,
+      toType: DataType,
+      timeZoneId: Option[String],
+      evalMode: String): Boolean = {
+    toType match {
+      case DataTypes.BooleanType =>
+        true
+      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType |
+          DataTypes.LongType =>
+        true
+      case DataTypes.FloatType | DataTypes.DoubleType =>
+        // https://github.com/apache/datafusion-comet/issues/326
+        false
+      case _: DecimalType =>
+        // https://github.com/apache/datafusion-comet/issues/325
+        false
+      case DataTypes.DateType =>
+        // https://github.com/apache/datafusion-comet/issues/327
+        false
+      case DataTypes.TimestampType =>
+        val enabled = CometConf.COMET_CAST_STRING_TO_TIMESTAMP.get()
+        if (!enabled) {
+          // https://github.com/apache/datafusion-comet/issues/328
+          withInfo(cast, s"${CometConf.COMET_CAST_STRING_TO_TIMESTAMP.key} is disabled")
+        }
+        enabled
+      case _ =>
+        false
+    }
+  }
+
+  private def castToStringSupported(
+      cast: Cast,
+      fromType: DataType,
+      timeZoneId: Option[String],
+      evalMode: String): Boolean = {
+    fromType match {
+      case DataTypes.BooleanType => true
+      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType |
+          DataTypes.LongType =>
+        true
+      case DataTypes.DateType => true
+      case DataTypes.TimestampType => true
+      case DataTypes.FloatType | DataTypes.DoubleType =>
+        // https://github.com/apache/datafusion-comet/issues/326
+        false
+      case _ => false
+    }
+  }
+
+  private def castFromTimestampSupported(
+      cast: Cast,
+      toType: DataType,
+      timeZoneId: Option[String],
+      evalMode: String): Boolean = {
+    toType match {
+      case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
+          DataTypes.IntegerType =>
+        // https://github.com/apache/datafusion-comet/issues/352
+        // this seems like an edge case that isn't important for us to support
+        false
+      case DataTypes.LongType =>
+        // https://github.com/apache/datafusion-comet/issues/352
+        false
+      case DataTypes.StringType => true
+      case DataTypes.DateType => true
+      case _ => false
+    }
+  }
+
+}

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

@@ -43,6 +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
 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}
@@ -575,7 +576,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
           val value = cast.eval()
           exprToProtoInternal(Literal(value, dataType), inputs)
 
-        case Cast(child, dt, timeZoneId, evalMode) =>
+        case cast @ Cast(child, dt, timeZoneId, evalMode) =>
           val childExpr = exprToProtoInternal(child, inputs)
           if (childExpr.isDefined) {
             val evalModeStr = if (evalMode.isInstanceOf[Boolean]) {
@@ -585,19 +586,12 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
               // Spark 3.4+ has EvalMode enum with values LEGACY, ANSI, and TRY
               evalMode.toString
             }
-            val supportedCast = (child.dataType, dt) match {
-              case (DataTypes.StringType, DataTypes.TimestampType)
-                  if !CometConf.COMET_CAST_STRING_TO_TIMESTAMP.get() =>
-                // https://github.com/apache/datafusion-comet/issues/328
-                withInfo(expr, s"${CometConf.COMET_CAST_STRING_TO_TIMESTAMP.key} is disabled")
-                false
-              case _ => true
-            }
-            if (supportedCast) {
+            if (CometCast.isSupported(cast, child.dataType, dt, timeZoneId, evalModeStr)) {
               castToProto(timeZoneId, dt, childExpr, evalModeStr)
             } else {
-              // no need to call withInfo here since it was called when determining
-              // the value for `supportedCast`
+              withInfo(
+                expr,
+                s"Unsupported cast from ${child.dataType} to $dt with timezone $timeZoneId and evalMode $evalModeStr")
               None
             }
           } else {

spark/src/test/scala/org/apache/comet/CometCastSuite.scala

@@ -771,12 +771,12 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       withSQLConf((SQLConf.ANSI_ENABLED.key, "false")) {
         // cast() should return null for invalid inputs when ansi mode is disabled
         val df = spark.sql(s"select a, cast(a as ${toType.sql}) from t order by a")
-        checkSparkAnswer(df)
+        checkSparkAnswerAndOperator(df)
 
         // try_cast() should always return null for invalid inputs
         val df2 =
           spark.sql(s"select a, try_cast(a as ${toType.sql}) from t order by a")
-        checkSparkAnswer(df2)
+        checkSparkAnswerAndOperator(df2)
       }
 
       // with ANSI enabled, we should produce the same exception as Spark
@@ -818,7 +818,7 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         // try_cast() should always return null for invalid inputs
         val df2 =
           spark.sql(s"select a, try_cast(a as ${toType.sql}) from t order by a")
-        checkSparkAnswer(df2)
+        checkSparkAnswerAndOperator(df2)
       }
     }
   }