Add config for cast float to integral types

docs/compatibility.md

@@ -241,6 +241,19 @@ To enable all formats on GPU, set
 
 In general, performing `cast` and `ansi_cast` operations on the GPU is compatible with the same operations on the CPU. However, there are some exceptions. For this reason, certain casts are disabled on the GPU by default and require configuration options to be specified to enable them. 
 
+### Float to Integral Types
+
+With both `cast` and `ansi_cast`, Spark uses the expression
+`Math.floor(x) <= MAX && Math.ceil(x) >= MIN` to determine whether a floating-point value can be
+converted to an integral type. Prior to Spark 3.1.0 the MIN and MAX values were floating-point
+values such as `Int.MaxValue.toFloat` but starting with 3.1.0 these are now integral types such as
+`Int.MaxValue` so this has slightly affected the valid range of values and now differs slightly
+from the behavior on GPU in some cases.
+
+To enable this operation on the GPU, set
+[`spark.rapids.sql.castFloatToIntegralTypes.enabled`](configs.md#sql.castFloatToIntegralTypes.enabled)
+to `true`.
+
 ### Float to String
 
 The GPU will use different precision than Java's toString method when converting floating-point data types to strings and this can produce results that differ from the default behavior in Spark. 
@@ -286,7 +299,7 @@ The following formats/patterns are supported on the GPU. Timezone of UTC is assu
 | `"tomorrow"`          | Yes               |
 | `"yesterday"`         | Yes               |
 
-## String to Timestamp
+### String to Timestamp
 
 To allow casts from string to timestamp on the GPU, enable the configuration property 
 [`spark.rapids.sql.castStringToTimestamp.enabled`](configs.md#sql.castStringToTimestamp.enabled).

docs/configs.md

@@ -48,6 +48,7 @@ Name | Description | Default Value
 <a name="shuffle.ucx.managementServerHost"></a>spark.rapids.shuffle.ucx.managementServerHost|The host to be used to start the management server|null
 <a name="shuffle.ucx.useWakeup"></a>spark.rapids.shuffle.ucx.useWakeup|When set to true, use UCX's event-based progress (epoll) in order to wake up the progress thread when needed, instead of a hot loop.|true
 <a name="sql.batchSizeBytes"></a>spark.rapids.sql.batchSizeBytes|Set the target number of bytes for a GPU batch. Splits sizes for input data is covered by separate configs. The maximum setting is 2 GB to avoid exceeding the cudf row count limit of a column.|2147483647
+<a name="sql.castFloatToIntegralTypes.enabled"></a>spark.rapids.sql.castFloatToIntegralTypes.enabled|Casting from floating point types to integral types on the GPU supports a slightly different range of values than Spark 3.1.0, which uses the following logic to determine valid inputs (example given is for Int): `Math.floor(x) <= Int.MaxValue && Math.ceil(x) >= Int.MinValue`|false
 <a name="sql.castFloatToString.enabled"></a>spark.rapids.sql.castFloatToString.enabled|Casting from floating point types to string on the GPU returns results that have a different precision than the default Java toString behavior.|false
 <a name="sql.castStringToFloat.enabled"></a>spark.rapids.sql.castStringToFloat.enabled|When set to true, enables casting from strings to float types (float, double) on the GPU. Currently hex values aren't supported on the GPU. Also note that casting from string to float types on the GPU returns incorrect results when the string represents any number "1.7976931348623158E308" <= x < "1.7976931348623159E308" and "-1.7976931348623158E308" >= x > "-1.7976931348623159E308" in both these cases the GPU returns Double.MaxValue while CPU returns "+Infinity" and "-Infinity" respectively|false
 <a name="sql.castStringToInteger.enabled"></a>spark.rapids.sql.castStringToInteger.enabled|When set to true, enables casting from strings to integer types (byte, short, int, long) on the GPU. Casting from string to integer types on the GPU returns incorrect results when the string represents a number larger than Long.MaxValue or smaller than Long.MinValue.|false

integration_tests/src/main/python/qa_nightly_select_test.py

@@ -145,6 +145,7 @@ def idfn(val):
         'spark.rapids.sql.variableFloatAgg.enabled': 'true',
         'spark.rapids.sql.hasNans': 'false',
         'spark.rapids.sql.castStringToFloat.enabled': 'true',
+        'spark.rapids.sql.castFloatToIntegralTypes.enabled': 'true',
         'spark.rapids.sql.castFloatToString.enabled': 'true'
         }
 

shims/spark310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/Spark310Shims.scala

@@ -21,6 +21,7 @@ import com.nvidia.spark.rapids.shims.spark301.Spark301Shims
 import com.nvidia.spark.rapids.spark310.RapidsShuffleManager
 
 import org.apache.spark.SparkEnv
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.analysis.Resolver
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
@@ -93,6 +94,21 @@ class Spark310Shims extends Spark301Shims {
   }
 
   def exprs310: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = Seq(
+    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) {
+          override def tagExprForGpu(): Unit = {
+            if (!conf.isCastFloatToIntegralTypesEnabled &&
+                (fromType == DataTypes.FloatType || fromType == DataTypes.DoubleType) &&
+                (toType == DataTypes.ByteType || toType == DataTypes.ShortType ||
+                    toType == DataTypes.IntegerType || toType == DataTypes.LongType)) {
+              willNotWorkOnGpu(buildTagMessage(RapidsConf.ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES))
+            }
+            super.tagExprForGpu()
+          }
+        }),
     GpuOverrides.expr[AnsiCast](
       "Convert a column of one type of data into another type",
       new CastChecks() {
@@ -134,7 +150,17 @@ class Spark310Shims extends Spark301Shims {
         override val udtChecks: TypeSig = none
         override val sparkUdtSig: TypeSig = UDT
       },
-      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, true, conf, p, r)),
+      (cast, conf, p, r) => new CastExprMeta[AnsiCast](cast, true, conf, p, r) {
+        override def tagExprForGpu(): Unit = {
+          if (!conf.isCastFloatToIntegralTypesEnabled &&
+              (fromType == DataTypes.FloatType || fromType == DataTypes.DoubleType) &&
+              (toType == DataTypes.ByteType || toType == DataTypes.ShortType ||
+                  toType == DataTypes.IntegerType || toType == DataTypes.LongType)) {
+            willNotWorkOnGpu(buildTagMessage(RapidsConf.ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES))
+          }
+          super.tagExprForGpu()
+        }
+      }),
     GpuOverrides.expr[RegExpReplace](
       "RegExpReplace support for string literal input patterns",
       ExprChecks.projectNotLambda(TypeSig.STRING, TypeSig.STRING,

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

@@ -34,8 +34,8 @@ class CastExprMeta[INPUT <: CastBase](
   extends UnaryExprMeta[INPUT](cast, conf, parent, rule) {
 
   private val castExpr = if (ansiEnabled) "ansi_cast" else "cast"
-  private val fromType = cast.child.dataType
-  private val toType = cast.dataType
+  val fromType = cast.child.dataType
+  val toType = cast.dataType
 
   override def tagExprForGpu(): Unit = {
     if (!conf.isCastFloatToStringEnabled && toType == DataTypes.StringType &&
@@ -72,6 +72,10 @@ class CastExprMeta[INPUT <: CastBase](
     }
   }
 
+  def buildTagMessage(entry: ConfEntry[_]): String = {
+    s"${entry.doc}. To enable this operation on the GPU, set ${entry.key} to true."
+  }
+
   override def convertToGpu(child: Expression): GpuExpression =
     GpuCast(child, toType, ansiEnabled, cast.timeZoneId)
 }

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

@@ -482,6 +482,14 @@ object RapidsConf {
     .booleanConf
     .createWithDefault(false)
 
+  val ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES =
+    conf("spark.rapids.sql.castFloatToIntegralTypes.enabled")
+      .doc("Casting from floating point types to integral types on the GPU supports a " +
+          "slightly different range of values when using Spark 3.1.0 or later. Refer to the CAST " +
+          "documentation for more details.")
+      .booleanConf
+      .createWithDefault(false)
+
   val ENABLE_CAST_STRING_TO_FLOAT = conf("spark.rapids.sql.castStringToFloat.enabled")
     .doc("When set to true, enables casting from strings to float types (float, double) " +
       "on the GPU. Currently hex values aren't supported on the GPU. Also note that casting from " +
@@ -1012,6 +1020,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val isCastStringToFloatEnabled: Boolean = get(ENABLE_CAST_STRING_TO_FLOAT)
 
+  lazy val isCastFloatToIntegralTypesEnabled: Boolean = get(ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES)
+
   lazy val isCsvTimestampEnabled: Boolean = get(ENABLE_CSV_TIMESTAMPS)
 
   lazy val isParquetEnabled: Boolean = get(ENABLE_PARQUET)

tests/src/test/scala/com/nvidia/spark/rapids/AnsiCastOpSuite.scala

@@ -34,6 +34,7 @@ class AnsiCastOpSuite extends GpuExpressionTestSuite {
   private val sparkConf = new SparkConf()
     .set("spark.sql.ansi.enabled", "true")
     .set("spark.sql.storeAssignmentPolicy", "ANSI") // note this is the default in 3.0.0
+    .set(RapidsConf.ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES.key, "true")
     .set(RapidsConf.ENABLE_CAST_FLOAT_TO_STRING.key, "true")
     .set(RapidsConf.ENABLE_CAST_STRING_TO_INTEGER.key, "true")
     .set(RapidsConf.ENABLE_CAST_STRING_TO_FLOAT.key, "true")

tests/src/test/scala/com/nvidia/spark/rapids/CastOpSuite.scala

@@ -30,6 +30,7 @@ class CastOpSuite extends GpuExpressionTestSuite {
   import CastOpSuite._
 
   private val sparkConf = new SparkConf()
+    .set(RapidsConf.ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES.key, "true")
     .set(RapidsConf.ENABLE_CAST_STRING_TO_FLOAT.key, "true")
 
   private val timestampDatesMsecParquet = frameFromParquet("timestamp-date-test-msec.parquet")
@@ -55,6 +56,7 @@ class CastOpSuite extends GpuExpressionTestSuite {
     Seq(false, true).foreach { ansiEnabled =>
 
       val conf = new SparkConf()
+        .set(RapidsConf.ENABLE_CAST_FLOAT_TO_INTEGRAL_TYPES.key, "true")
         .set(RapidsConf.ENABLE_CAST_FLOAT_TO_STRING.key, "true")
         .set(RapidsConf.ENABLE_CAST_STRING_TO_TIMESTAMP.key, "true")
         .set(RapidsConf.ENABLE_CAST_STRING_TO_INTEGER.key, "true")
@@ -264,7 +266,8 @@ class CastOpSuite extends GpuExpressionTestSuite {
       col("longs").cast(TimestampType))
   }
 
-  testSparkResultsAreEqual("Test cast from float", mixedFloatDf) {
+  testSparkResultsAreEqual("Test cast from float", mixedFloatDf,
+      conf = sparkConf) {
     frame => frame.select(
       col("floats").cast(IntegerType),
       col("floats").cast(LongType),
@@ -277,7 +280,8 @@ class CastOpSuite extends GpuExpressionTestSuite {
       col("floats").cast(TimestampType))
   }
 
-  testSparkResultsAreEqual("Test cast from double", doubleWithNansDf) {
+  testSparkResultsAreEqual("Test cast from double", doubleWithNansDf,
+      conf = sparkConf) {
     frame => frame.select(
       col("doubles").cast(IntegerType),
       col("doubles").cast(LongType),