Fall back to CPU for try_cast in Spark 3.4.0 [databricks]

integration_tests/src/main/python/cast_test.py

@@ -16,7 +16,7 @@
 
 from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_and_cpu_are_equal_sql, assert_gpu_and_cpu_error, assert_gpu_fallback_collect, assert_py4j_exception
 from data_gen import *
-from spark_session import is_before_spark_320, is_before_spark_330, with_gpu_session
+from spark_session import is_before_spark_320, is_before_spark_330, is_spark_32X, is_spark_33X, with_gpu_session
 from marks import allow_non_gpu, approximate_float
 from pyspark.sql.types import *
 from spark_init_internal import spark_version
@@ -115,6 +115,19 @@ def test_cast_string_date_invalid_ansi(invalid):
               'spark.sql.ansi.enabled': 'true'},
         error_message="DateTimeException")
 
+
+# test try_cast in Spark versions >= 340
+test_try_cast_fallback_non_gpu = ['ProjectExec', 'Cast'] if is_spark_340_or_later() or is_databricks113_or_later() else ['ProjectExec','TryCast']
+@pytest.mark.skipif(is_before_spark_320(), reason="try_cast only in Spark 3.2+")
+@allow_non_gpu(test_try_cast_fallback_non_gpu)
+@pytest.mark.parametrize('invalid', invalid_values_string_to_date)
+def test_try_cast_fallback_340(invalid):
+    assert_gpu_fallback_collect(
+        lambda spark: spark.createDataFrame([(invalid,)], "a string").selectExpr("try_cast(a as date)"),
+        'Cast',
+        conf={'spark.rapids.sql.hasExtendedYearValues': 'false',
+              'spark.sql.ansi.enabled': 'true'})
+
 # test all Spark versions, non ANSI mode, invalid value will be converted to NULL
 def test_cast_string_date_non_ansi():
     data_rows = [(v,) for v in values_string_to_data]

integration_tests/src/main/python/spark_session.py

@@ -163,6 +163,9 @@ def is_spark_340_or_later():
 def is_spark_330():
     return spark_version() == "3.3.0"
 
+def is_spark_32X():
+    return "3.2.0" <= spark_version() < "3.3.0"
+
 def is_spark_33X():
     return "3.3.0" <= spark_version() < "3.4.0"
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCast.scala

@@ -40,7 +40,7 @@ import org.apache.spark.sql.types._
 /** Meta-data for cast and ansi_cast. */
 final class CastExprMeta[INPUT <: UnaryExpression with TimeZoneAwareExpression with NullIntolerant](
     cast: INPUT,
-    val ansiEnabled: Boolean,
+    val evalMode: GpuEvalMode.Value,
     conf: RapidsConf,
     parent: Option[RapidsMeta[_, _, _]],
     rule: DataFromReplacementRule,
@@ -52,15 +52,28 @@ final class CastExprMeta[INPUT <: UnaryExpression with TimeZoneAwareExpression w
     toTypeOverride: Option[DataType] = None)
   extends UnaryExprMeta[INPUT](cast, conf, parent, rule) {
 
-  def withToTypeOverride(newToType: DecimalType): CastExprMeta[INPUT] =
-    new CastExprMeta[INPUT](cast, ansiEnabled, conf, parent, rule,
+  val ansiEnabled = evalMode == GpuEvalMode.ANSI
+
+  def withToTypeOverride(newToType: DecimalType): CastExprMeta[INPUT] = {
+    val evalMode = if (ansiEnabled) {
+      GpuEvalMode.ANSI
+    } else {
+      GpuEvalMode.LEGACY
+    }
+    new CastExprMeta[INPUT](cast, evalMode, conf, parent, rule,
       doFloatToIntCheck, stringToAnsiDate, Some(newToType))
+  }
 
   val fromType: DataType = cast.child.dataType
   val toType: DataType = toTypeOverride.getOrElse(cast.dataType)
   val legacyCastToString: Boolean = SQLConf.get.getConf(SQLConf.LEGACY_COMPLEX_TYPES_TO_STRING)
 
-  override def tagExprForGpu(): Unit = recursiveTagExprForGpuCheck()
+  override def tagExprForGpu(): Unit = {
+    if (evalMode == GpuEvalMode.TRY) {
+      willNotWorkOnGpu("try_cast is not supported on the GPU")
+    }
+    recursiveTagExprForGpuCheck()
+  }
 
   private def recursiveTagExprForGpuCheck(
       fromDataType: DataType = fromType,

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuEvalMode.scala

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * 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 com.nvidia.spark.rapids
+
+/**
+ * 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 GpuEvalMode extends Enumeration {
+  val LEGACY, ANSI, TRY = Value
+}

sql-plugin/src/main/spark311/scala/com/nvidia/spark/rapids/shims/AnsiCastShim.scala

@@ -32,8 +32,9 @@
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims
 
-import com.nvidia.spark.rapids.GpuCast
+import com.nvidia.spark.rapids.{GpuCast, GpuEvalMode}
 
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{AnsiCast, Cast, Expression}
 
 object AnsiCastShim {
@@ -46,4 +47,12 @@ object AnsiCastShim {
       m.getBoolean(e)
     case _ => false
   }
+
+  def getEvalMode(c: Cast): GpuEvalMode.Value = {
+    if (SparkSession.active.sessionState.conf.ansiEnabled) {
+      GpuEvalMode.ANSI
+    } else {
+      GpuEvalMode.LEGACY
+    }
+  }
 }

sql-plugin/src/main/spark311/scala/com/nvidia/spark/rapids/shims/Spark31XShims.scala

@@ -235,7 +235,7 @@ abstract class Spark31XShims extends Spark31Xuntil33XShims with Logging {
         override val udtChecks: TypeSig = none
         override val sparkUdtSig: TypeSig = UDT
       },
-      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, ansiEnabled = true, conf = conf,
+      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, GpuEvalMode.ANSI, conf = conf,
         parent = p, rule = r, doFloatToIntCheck = true, stringToAnsiDate = false))
   }
 
@@ -244,7 +244,7 @@ abstract class Spark31XShims extends Spark31Xuntil33XShims with Logging {
         "Convert a column of one type of data into another type",
         new CastChecks(),
         (cast, conf, p, r) => new CastExprMeta[Cast](cast,
-          SparkSession.active.sessionState.conf.ansiEnabled, conf, p, r,
+          AnsiCastShim.getEvalMode(cast), conf, p, r,
           doFloatToIntCheck = true, stringToAnsiDate = false)),
     GpuOverrides.expr[Average](
       "Average aggregate operator",

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/AnsiCastRuleShims.scala

@@ -92,7 +92,7 @@ trait AnsiCastRuleShims extends SparkShims {
         override val udtChecks: TypeSig = none
         override val sparkUdtSig: TypeSig = UDT
       },
-      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, ansiEnabled = true, conf = conf,
+      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, GpuEvalMode.ANSI, conf = conf,
         parent = p, rule = r, doFloatToIntCheck = true, stringToAnsiDate = true))
   }
 }

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/Spark320PlusShims.scala

@@ -139,9 +139,11 @@ trait Spark320PlusShims extends SparkShims with RebaseShims with Logging {
     GpuOverrides.expr[Cast](
       "Convert a column of one type of data into another type",
       new CastChecks(),
-      (cast, conf, p, r) => new CastExprMeta[Cast](cast,
-        SparkSession.active.sessionState.conf.ansiEnabled, conf, p, r,
-        doFloatToIntCheck = true, stringToAnsiDate = true)),
+      (cast, conf, p, r) => {
+        new CastExprMeta[Cast](cast,
+          AnsiCastShim.getEvalMode(cast), conf, p, r,
+          doFloatToIntCheck = true, stringToAnsiDate = true)
+      }),
     GpuOverrides.expr[Average](
       "Average aggregate operator",
       ExprChecks.fullAgg(

sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/AnsiCastShim.scala

@@ -19,7 +19,7 @@
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims
 
-import com.nvidia.spark.rapids.GpuCast
+import com.nvidia.spark.rapids.{GpuCast, GpuEvalMode}
 
 import org.apache.spark.sql.catalyst.expressions.{Cast, EvalMode, Expression}
 
@@ -29,4 +29,12 @@ object AnsiCastShim {
     case c: Cast => c.evalMode == EvalMode.ANSI
     case _ => false
   }
+
+  def getEvalMode(c: Cast): GpuEvalMode.Value = {
+    c.evalMode match {
+      case EvalMode.LEGACY => GpuEvalMode.LEGACY
+      case EvalMode.ANSI => GpuEvalMode.ANSI
+      case EvalMode.TRY => GpuEvalMode.TRY
+    }
+  }
 }

sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/Spark340PlusShims.scala

@@ -22,8 +22,7 @@ package com.nvidia.spark.rapids.shims
 import com.nvidia.spark.rapids._
 
 import org.apache.spark.rapids.shims.GpuShuffleExchangeExec
-import org.apache.spark.sql.catalyst.expressions.{Expression, KnownNullable}
-import org.apache.spark.sql.catalyst.expressions.Empty2Null
+import org.apache.spark.sql.catalyst.expressions.{Empty2Null, Expression, KnownNullable}
 import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
 import org.apache.spark.sql.execution.{CollectLimitExec, GlobalLimitExec, SparkPlan}
 import org.apache.spark.sql.execution.command.{CreateDataSourceTableAsSelectCommand, DataWritingCommand, RunnableCommand}
@@ -79,6 +78,7 @@ trait Spark340PlusShims extends Spark331PlusShims {
   // AnsiCast is removed from Spark3.4.0
   override def ansiCastRule: ExprRule[_ <: Expression] = null
 
+
   override def getExprs: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = {
     val shimExprs: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = Seq(
       // Empty2Null is pulled out of FileFormatWriter by default since Spark 3.4.0,