Fall back to CPU for RoundCeil and RoundFloor expressions (NVIDIA#5798)

* Fall back to CPU for RoundCeil and RoundFloor expressions

Signed-off-by: Niranjan Artal <nartal@nvidia.com>

integration_tests/src/main/python/arithmetic_ops_test.py

@@ -404,17 +404,18 @@ def test_hypot(data_gen):
             'hypot(a, b)',
         ))
 
-@pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale, ids=idfn)
+@pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale + [DecimalGen(30, 15)], ids=idfn)
 def test_floor(data_gen):
     assert_gpu_and_cpu_are_equal_collect(
             lambda spark : unary_op_df(spark, data_gen).selectExpr('floor(a)'))
 
+# This test should be enabled to run on GPU as part of https://github.com/NVIDIA/spark-rapids/issues/5797
 @pytest.mark.skipif(is_before_spark_330(), reason='scale parameter in Floor function is not supported before Spark 3.3.0')
+@allow_non_gpu('ProjectExec')
 @pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale, ids=idfn)
 def test_floor_scale_zero(data_gen):
-    assert_gpu_and_cpu_are_equal_collect(
-            lambda spark : unary_op_df(spark, data_gen).selectExpr('floor(a, 0)'),
-            conf={'spark.rapids.sql.castFloatToDecimal.enabled':'true'})
+    assert_gpu_fallback_collect(
+            lambda spark : unary_op_df(spark, data_gen).selectExpr('floor(a, 0)'), 'RoundFloor')
 
 @pytest.mark.skipif(is_before_spark_330(), reason='scale parameter in Floor function is not supported before Spark 3.3.0')
 @allow_non_gpu('ProjectExec')
@@ -423,17 +424,18 @@ def test_floor_scale_nonzero(data_gen):
     assert_gpu_fallback_collect(
             lambda spark : unary_op_df(spark, data_gen).selectExpr('floor(a, -1)'), 'RoundFloor')
 
-@pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale, ids=idfn)
+@pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale + [DecimalGen(30, 15)], ids=idfn)
 def test_ceil(data_gen):
     assert_gpu_and_cpu_are_equal_collect(
             lambda spark : unary_op_df(spark, data_gen).selectExpr('ceil(a)'))
 
+# This test should be enabled to run on GPU as part of https://github.com/NVIDIA/spark-rapids/issues/5797
 @pytest.mark.skipif(is_before_spark_330(), reason='scale parameter in Ceil function is not supported before Spark 3.3.0')
+@allow_non_gpu('ProjectExec')
 @pytest.mark.parametrize('data_gen', double_n_long_gens + _arith_decimal_gens_no_neg_scale, ids=idfn)
 def test_ceil_scale_zero(data_gen):
-    assert_gpu_and_cpu_are_equal_collect(
-            lambda spark : unary_op_df(spark, data_gen).selectExpr('ceil(a, 0)'),
-            conf={'spark.rapids.sql.castFloatToDecimal.enabled':'true'})
+    assert_gpu_fallback_collect(
+            lambda spark : unary_op_df(spark, data_gen).selectExpr('ceil(a, 0)'), 'RoundCeil')
 
 @pytest.mark.parametrize('data_gen', [_decimal_gen_36_neg5, _decimal_gen_38_neg10], ids=idfn)
 def test_floor_ceil_overflow(data_gen):

sql-plugin/src/main/330+/scala/com/nvidia/spark/rapids/shims/Spark33XShims.scala

@@ -87,6 +87,7 @@ trait Spark33XShims extends Spark321PlusShims with Spark320PlusNonDBShims {
           ("scale", TypeSig.lit(TypeEnum.INT), TypeSig.lit(TypeEnum.INT))),
         (ceil, conf, p, r) => new BinaryExprMeta[RoundCeil](ceil, conf, p, r) {
           override def tagExprForGpu(): Unit = {
+            willNotWorkOnGpu("RoundCeil is currently not supported on GPU")
             ceil.child.dataType match {
               case dt: DecimalType =>
                 val precision = GpuFloorCeil.unboundedOutputPrecision(dt)
@@ -118,6 +119,7 @@ trait Spark33XShims extends Spark321PlusShims with Spark320PlusNonDBShims {
           ("scale", TypeSig.lit(TypeEnum.INT), TypeSig.lit(TypeEnum.INT))),
         (floor, conf, p, r) => new BinaryExprMeta[RoundFloor](floor, conf, p, r) {
           override def tagExprForGpu(): Unit = {
+            willNotWorkOnGpu("RoundFloor is currently not supported on GPU")
             floor.child.dataType match {
               case dt: DecimalType =>
                 val precision = GpuFloorCeil.unboundedOutputPrecision(dt)

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/mathExpressions.scala

@@ -23,7 +23,6 @@ import ai.rapids.cudf.ast.BinaryOperator
 import com.nvidia.spark.rapids._
 
 import org.apache.spark.sql.catalyst.expressions.{Expression, ImplicitCastInputTypes}
-import org.apache.spark.sql.rapids.shims.RapidsFloorCeilUtils
 import org.apache.spark.sql.types._
 
 abstract class CudfUnaryMathExpression(name: String) extends GpuUnaryMathExpression(name)
@@ -166,7 +165,11 @@ object GpuFloorCeil {
 }
 
 case class GpuCeil(child: Expression) extends CudfUnaryMathExpression("CEIL") {
-  override def dataType: DataType = RapidsFloorCeilUtils.outputDataType(child.dataType)
+  override def dataType: DataType = child.dataType match {
+    case dt: DecimalType =>
+      DecimalType.bounded(GpuFloorCeil.unboundedOutputPrecision(dt), 0)
+    case _ => LongType
+  }
 
   override def hasSideEffects: Boolean = true
 
@@ -238,7 +241,11 @@ case class GpuExpm1(child: Expression) extends CudfUnaryMathExpression("EXPM1")
 }
 
 case class GpuFloor(child: Expression) extends CudfUnaryMathExpression("FLOOR") {
-  override def dataType: DataType = RapidsFloorCeilUtils.outputDataType(child.dataType)
+  override def dataType: DataType = child.dataType match {
+    case dt: DecimalType =>
+      DecimalType.bounded(GpuFloorCeil.unboundedOutputPrecision(dt), 0)
+    case _ => LongType
+  }
 
   override def hasSideEffects: Boolean = true