Fallback to CPU for mod only for DB11.3 [Databricks]

integration_tests/src/main/python/arithmetic_ops_test.py

@@ -272,6 +272,7 @@ def test_int_division_mixed(lhs, rhs):
                 'a DIV b'))
 
 @pytest.mark.parametrize('data_gen', _arith_data_gens, ids=idfn)
+@pytest.mark.skipif(not is_before_spark_340() or is_databricks113_or_later())
 def test_mod(data_gen):
     data_type = data_gen.data_type
     assert_gpu_and_cpu_are_equal_collect(
@@ -282,6 +283,20 @@ def test_mod(data_gen):
                 f.col('b') % f.lit(None).cast(data_type),
                 f.col('a') % f.col('b')))
 
+# This test is only added because we are skipping test_mod for spark 3.4 and databricks 11.3 because of https://github.com/NVIDIA/spark-rapids/issues/7595
+# Once that is resolved we should remove this test and not skip test_mod for spark 3.4 and db 11.3
+@pytest.mark.parametrize('data_gen', _arith_data_gens, ids=idfn)
+@pytest.mark.skipif(is_before_spark_340() and not is_databricks113_or_later())
+def test_mod_db11_3(data_gen):
+    data_type = data_gen.data_type
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : binary_op_df(spark, data_gen).select(
+            f.col('a') % f.lit(100).cast(data_type),
+            f.lit(-12).cast(data_type) % f.col('b'),
+            f.lit(None).cast(data_type) % f.col('a'),
+            f.col('b') % f.lit(None).cast(data_type),
+            f.col('a') % f.col('b')))
+
 # pmod currently falls back for Decimal(precision=38)
 # https://github.com/NVIDIA/spark-rapids/issues/6336
 # only testing numeric_gens because of https://github.com/NVIDIA/spark-rapids/issues/7553
@@ -376,10 +391,39 @@ def test_mod_pmod_by_zero_not_ansi(data_gen):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, data_gen).selectExpr(
             'pmod(a, cast(0 as {}))'.format(string_type),
-            'pmod(cast(-12 as {}), cast(0 as {}))'.format(string_type, string_type),
-            'a % (cast(0 as {}))'.format(string_type),
-            'cast(-12 as {}) % cast(0 as {})'.format(string_type, string_type)),
+            'pmod(cast(-12 as {}), cast(0 as {}))'.format(string_type, string_type)),
         {'spark.sql.ansi.enabled': 'false'})
+    # Skip decimal tests for mod on spark 3.4 and databricks 11.3, reason=https://github.com/NVIDIA/spark-rapids/issues/7595
+    if is_before_spark_340() or not is_databricks113_or_later():
+        assert_gpu_and_cpu_are_equal_collect(
+            lambda spark : unary_op_df(spark, data_gen).selectExpr(
+                'a % (cast(0 as {}))'.format(string_type),
+                'cast(-12 as {}) % cast(0 as {})'.format(string_type, string_type)),
+            {'spark.sql.ansi.enabled': 'false'})
+
+mod_mixed_decimals_lhs = [DecimalGen(6, 5), DecimalGen(6, 4), DecimalGen(5, 4), DecimalGen(5, 3), DecimalGen(4, 2),
+                           DecimalGen(3, -2), DecimalGen(16, 7), DecimalGen(19, 0), DecimalGen(30, 10)]
+mod_mixed_decimals_rhs = [DecimalGen(6, 3), DecimalGen(10, -2), DecimalGen(15, 3), DecimalGen(30, 12),
+                          DecimalGen(3, -3), DecimalGen(27, 7), DecimalGen(20, -3)]
+mod_mixed_lhs = [byte_gen, short_gen, int_gen, long_gen]
+mod_mixed_lhs.extend(pytest.param(t, marks=pytest.mark.skipif(is_databricks113_or_later() or not is_before_spark_340(),
+                                     reason='https://github.com/NVIDIA/spark-rapids/issues/7595')) for t in mod_mixed_decimals_lhs)
+mod_mixed_rhs = [byte_gen, short_gen, int_gen, long_gen]
+mod_mixed_rhs.extend(pytest.param(t, marks=pytest.mark.skipif(is_databricks113_or_later() or not is_before_spark_340(),
+                                     reason='https://github.com/NVIDIA/spark-rapids/issues/7595')) for t in mod_mixed_decimals_rhs)
+@pytest.mark.parametrize('lhs', mod_mixed_lhs, ids=idfn)
+@pytest.mark.parametrize('rhs', mod_mixed_rhs, ids=idfn)
+def test_mod_mixed(lhs, rhs):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : two_col_df(spark, lhs, rhs).selectExpr(f"a % b"))
+
+@allow_non_gpu('ProjectExec')
+@pytest.mark.skipif(not is_databricks113_or_later() or is_spark_340_or_later(), reason='https://github.com/NVIDIA/spark-rapids/issues/7595')
+@pytest.mark.parametrize('lhs', mod_mixed_decimals_lhs, ids=idfn)
+@pytest.mark.parametrize('rhs', mod_mixed_decimals_rhs, ids=idfn)
+def test_mod_fallback(lhs, rhs):
+    assert_gpu_fallback_collect(
+        lambda spark : two_col_df(spark, lhs, rhs).selectExpr(f"a % b"), 'Remainder')
 
 # Split into 4 tests to permute https://github.com/NVIDIA/spark-rapids/issues/7553 failures
 # @pytest.mark.parametrize('lhs', [byte_gen, short_gen, int_gen, long_gen, DecimalGen(6, 5),

sql-plugin/src/main/311until340-non330db/scala/com/nvidia/spark/rapids/shims/DecimalArithmeticOverrides.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,9 +20,9 @@ import ai.rapids.cudf.DType
 import com.nvidia.spark.rapids.{BaseExprMeta, BinaryAstExprMeta, BinaryExprMeta, CastExprMeta, DecimalUtil, ExprChecks, ExprMeta, ExprRule, GpuCheckOverflow, GpuExpression, GpuPromotePrecision, LiteralExprMeta, TypeSig, UnaryExprMeta}
 import com.nvidia.spark.rapids.GpuOverrides.expr
 
-import org.apache.spark.sql.catalyst.expressions.{CastBase, CheckOverflow, Divide, Expression, IntegralDivide, Literal, Multiply, PromotePrecision}
+import org.apache.spark.sql.catalyst.expressions.{CastBase, CheckOverflow, Divide, Expression, IntegralDivide, Literal, Multiply, PromotePrecision, Remainder}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{DecimalMultiplyChecks, GpuAnsi, GpuDecimalDivide, GpuDecimalMultiply, GpuDivide, GpuIntegralDivide, GpuMultiply}
+import org.apache.spark.sql.rapids.{DecimalMultiplyChecks, GpuAnsi, GpuDecimalDivide, GpuDecimalMultiply, GpuDivide, GpuIntegralDivide, GpuMultiply, GpuRemainder}
 import org.apache.spark.sql.types.{Decimal, DecimalType}
 
 object DecimalArithmeticOverrides {
@@ -188,6 +188,16 @@ object DecimalArithmeticOverrides {
       (a, conf, p, r) => new BinaryExprMeta[IntegralDivide](a, conf, p, r) {
         override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
           GpuIntegralDivide(lhs, rhs)
+      }),
+    expr[Remainder](
+      "Remainder or modulo",
+      ExprChecks.binaryProject(
+        TypeSig.gpuNumeric, TypeSig.cpuNumeric,
+        ("lhs", TypeSig.gpuNumeric, TypeSig.cpuNumeric),
+        ("rhs", TypeSig.gpuNumeric, TypeSig.cpuNumeric)),
+      (a, conf, p, r) => new BinaryExprMeta[Remainder](a, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
+          GpuRemainder(lhs, rhs)
       })
   ).map(r => (r.getClassFor.asSubclass(classOf[Expression]), r)).toMap
 }

sql-plugin/src/main/340+-and-330db/scala/com/nvidia/spark/rapids/shims/DecimalArithmeticOverrides.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,9 +20,9 @@ import ai.rapids.cudf.DType
 import com.nvidia.spark.rapids.{BinaryAstExprMeta, BinaryExprMeta, DecimalUtil, ExprChecks, ExprRule, GpuExpression, TypeSig}
 import com.nvidia.spark.rapids.GpuOverrides.expr
 
-import org.apache.spark.sql.catalyst.expressions.{Divide, Expression, IntegralDivide, Multiply}
+import org.apache.spark.sql.catalyst.expressions.{Divide, Expression, IntegralDivide, Multiply, Remainder}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{DecimalMultiplyChecks, GpuAnsi, GpuDecimalDivide, GpuDecimalMultiply, GpuDivide, GpuIntegralDecimalDivide, GpuIntegralDivide, GpuMultiply}
+import org.apache.spark.sql.rapids.{DecimalMultiplyChecks, GpuAnsi, GpuDecimalDivide, GpuDecimalMultiply, GpuDivide, GpuIntegralDecimalDivide, GpuIntegralDivide, GpuMultiply, GpuRemainder}
 import org.apache.spark.sql.types.DecimalType
 
 object DecimalArithmeticOverrides {
@@ -93,6 +93,22 @@ object DecimalArithmeticOverrides {
             } else {
               GpuIntegralDivide(lhs, rhs)
             }
+        }),
+
+      /**
+       * Because of https://github.com/NVIDIA/spark-rapids/issues/7595 we are not supporting Decimals for spark 3.4 and
+       * db 11.3. Once we do we should revert the changes made to the following tests test_mod, test_mod_mixed and test_mod_pmod_by_zero_not_ansi
+       * or we should just revert this commit
+       */
+      expr[Remainder](
+        "Remainder or modulo",
+        ExprChecks.binaryProject(
+          TypeSig.gpuNumeric, TypeSig.cpuNumeric,
+          ("lhs", TypeSig.integral + TypeSig.fp, TypeSig.cpuNumeric),
+          ("rhs", TypeSig.integral + TypeSig.fp, TypeSig.cpuNumeric)),
+        (a, conf, p, r) => new BinaryExprMeta[Remainder](a, conf, p, r) {
+          override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
+            GpuRemainder(lhs, rhs)
         })
     ).map(r => (r.getClassFor.asSubclass(classOf[Expression]), r)).toMap
   }

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

@@ -1949,16 +1949,6 @@ object GpuOverrides extends Logging {
         override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
           GpuPow(lhs, rhs)
       }),
-    expr[Remainder](
-      "Remainder or modulo",
-      ExprChecks.binaryProject(
-        TypeSig.gpuNumeric, TypeSig.cpuNumeric,
-        ("lhs", TypeSig.gpuNumeric, TypeSig.cpuNumeric),
-        ("rhs", TypeSig.gpuNumeric, TypeSig.cpuNumeric)),
-      (a, conf, p, r) => new BinaryExprMeta[Remainder](a, conf, p, r) {
-        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression =
-          GpuRemainder(lhs, rhs)
-      }),
     expr[AggregateExpression](
       "Aggregate expression",
       ExprChecks.fullAgg(