ANSI check for aggregates

integration_tests/src/main/python/hash_aggregate_test.py

@@ -182,6 +182,14 @@
     _grpkey_floats_with_nulls_and_nans]
 
 
+# Used to test ANSI-mode fallback
+_no_overflow_ansi_gens = [
+    ByteGen(min_val = 1, max_val = 10, special_cases=[]),
+    ShortGen(min_val = 1, max_val = 100, special_cases=[]),
+    IntegerGen(min_val = 1, max_val = 1000, special_cases=[]),
+    LongGen(min_val = 1, max_val = 3000, special_cases=[])]
+
+
 def get_params(init_list, marked_params=[]):
     """
     A method to build the test inputs along with their passed in markers to allow testing
@@ -259,29 +267,6 @@ def test_hash_grpby_avg(data_gen, conf):
         conf=conf
     )
 
-@ignore_order
-@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
-@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
-@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
-def test_hash_grpby_avg_nulls(data_gen, conf, ansi_enabled):
-    local_conf = copy_and_update(conf, {'spark.sql.ansi.enabled': ansi_enabled})
-    assert_gpu_and_cpu_are_equal_collect(
-        lambda spark: gen_df(spark, data_gen, length=100).groupby('a')
-          .agg(f.avg('c')),
-        conf=local_conf
-    )
-
-@ignore_order
-@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
-@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
-@pytest.mark.parametrize('ansi_enabled', ['true', 'false'])
-def test_hash_reduction_avg_nulls(data_gen, conf, ansi_enabled):
-    local_conf = copy_and_update(conf, {'spark.sql.ansi.enabled': ansi_enabled})
-    assert_gpu_and_cpu_are_equal_collect(
-        lambda spark: gen_df(spark, data_gen, length=100)
-          .agg(f.avg('c')),
-        conf=local_conf
-    )
 
 # tracks https://github.com/NVIDIA/spark-rapids/issues/154
 @approximate_float
@@ -1088,3 +1073,99 @@ def do_it(spark):
         df = two_col_df(spark, StringGen('k{1,5}'), ArrayGen(MapGen(StringGen('a{1,5}', nullable=False), StringGen('[ab]{1,5}'))))
         return df.groupBy('a').agg(f.min(df.b[1]["a"]))
     assert_gpu_and_cpu_are_equal_collect(do_it)
+
+@ignore_order
+@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
+@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
+def test_hash_grpby_avg_nulls(data_gen, conf):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, data_gen, length=100).groupby('a')
+          .agg(f.avg('c')),
+        conf=conf
+    )
+
+@ignore_order
+@allow_non_gpu('HashAggregateExec', 'Alias', 'AggregateExpression', 'Cast',
+  'HashPartitioning', 'ShuffleExchangeExec', 'Average')
+@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
+@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
+def test_hash_grpby_avg_nulls_ansi(data_gen, conf):
+    local_conf = copy_and_update(conf, {'spark.sql.ansi.enabled': 'true'})
+    assert_gpu_fallback_collect(
+        lambda spark: gen_df(spark, data_gen, length=100).groupby('a')
+          .agg(f.avg('c')),
+        'Average',
+        conf=local_conf
+    )
+
+@ignore_order
+@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
+@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
+def test_hash_reduction_avg_nulls(data_gen, conf):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, data_gen, length=100)
+          .agg(f.avg('c')),
+        conf=conf
+    )
+
+@ignore_order
+@allow_non_gpu('HashAggregateExec', 'Alias', 'AggregateExpression', 'Cast',
+  'HashPartitioning', 'ShuffleExchangeExec', 'Average')
+@pytest.mark.parametrize('data_gen', [_grpkey_strings_with_extra_nulls], ids=idfn)
+@pytest.mark.parametrize('conf', get_params(_confs, params_markers_for_confs), ids=idfn)
+def test_hash_reduction_avg_nulls_ansi(data_gen, conf):
+    local_conf = copy_and_update(conf, {'spark.sql.ansi.enabled': 'true'})
+    assert_gpu_fallback_collect(
+        lambda spark: gen_df(spark, data_gen, length=100)
+          .agg(f.avg('c')),
+        'Average',
+        conf=local_conf
+    )
+
+
+@allow_non_gpu('HashAggregateExec', 'Alias', 'AggregateExpression', 'Cast',
+  'HashPartitioning', 'ShuffleExchangeExec', 'Sum')
+@pytest.mark.parametrize('data_gen', _no_overflow_ansi_gens, ids=idfn)
+def test_sum_fallback_when_ansi_enabled(data_gen):
+    def do_it(spark):
+        df = gen_df(spark, [('a', data_gen), ('b', data_gen)], length=100)
+        return df.groupBy('a').agg(f.sum("b"))
+
+    assert_gpu_fallback_collect(do_it, 'Sum',
+        conf={'spark.sql.ansi.enabled': 'true'})
+
+
+@allow_non_gpu('HashAggregateExec', 'Alias', 'AggregateExpression', 'Cast',
+  'HashPartitioning', 'ShuffleExchangeExec', 'Average')
+@pytest.mark.parametrize('data_gen', _no_overflow_ansi_gens, ids=idfn)
+def test_avg_fallback_when_ansi_enabled(data_gen):
+    def do_it(spark):
+        df = gen_df(spark, [('a', data_gen), ('b', data_gen)], length=100)
+        return df.groupBy('a').agg(f.avg("b"))
+
+    assert_gpu_fallback_collect(do_it, 'Average',
+        conf={'spark.sql.ansi.enabled': 'true'})
+
+
+@allow_non_gpu('HashAggregateExec', 'Alias', 'AggregateExpression',
+  'HashPartitioning', 'ShuffleExchangeExec', 'Count', 'Literal')
+@pytest.mark.parametrize('data_gen', _no_overflow_ansi_gens, ids=idfn)
+def test_count_fallback_when_ansi_enabled(data_gen):
+    def do_it(spark):
+        df = gen_df(spark, [('a', data_gen), ('b', data_gen)], length=100)
+        return df.groupBy('a').agg(f.count("b"), f.count("*"))
+
+    assert_gpu_fallback_collect(do_it, 'Count',
+        conf={'spark.sql.ansi.enabled': 'true'})
+
+
+@pytest.mark.parametrize('data_gen', _no_overflow_ansi_gens, ids=idfn)
+def test_no_fallback_when_ansi_enabled(data_gen):
+    def do_it(spark):
+        df = gen_df(spark, [('a', data_gen), ('b', data_gen)], length=100)
+        # coalescing because of first/last are not deterministic
+        df = df.coalesce(1).orderBy("a", "b")
+        return df.groupBy('a').agg(f.first("b"), f.last("b"), f.min("b"), f.max("b"))
+
+    assert_gpu_and_cpu_are_equal_collect(do_it,
+        conf={'spark.sql.ansi.enabled': 'true'})

shims/spark320/src/main/scala/com/nvidia/spark/rapids/shims/spark320/Spark320Shims.scala

@@ -228,12 +228,16 @@ class Spark320Shims extends Spark32XShims {
         Seq(ParamCheck("input", TypeSig.integral + TypeSig.fp + TypeSig.NULL,
           TypeSig.numericAndInterval + TypeSig.NULL))),
       (a, conf, p, r) => new AggExprMeta[Average](a, conf, p, r) {
-        override def tagExprForGpu(): Unit = {
+        override def tagAggForGpu(): Unit = {
           val dataType = a.child.dataType
           GpuOverrides.checkAndTagFloatAgg(dataType, conf, this)
         }
 
-        override def convertToGpu(child: Expression): GpuExpression = GpuAverage(child)
+        override def convertToGpu(childExprs: Seq[Expression]): GpuExpression =
+          GpuAverage(childExprs.head)
+
+        // Average is not supported in ANSI mode right now, no matter the type
+        override val ansiTypeToCheck: Option[DataType] = None
       }),
     GpuOverrides.expr[Abs](
       "Absolute value",

sql-plugin/src/main/301db/scala/com/nvidia/spark/rapids/shims/v2/SparkBaseShims.scala

@@ -310,12 +310,16 @@ abstract class SparkBaseShims extends Spark30XShims {
           TypeSig.DOUBLE, TypeSig.DOUBLE + TypeSig.DECIMAL_128_FULL,
           Seq(ParamCheck("input", TypeSig.integral + TypeSig.fp, TypeSig.numeric))),
         (a, conf, p, r) => new AggExprMeta[Average](a, conf, p, r) {
-          override def tagExprForGpu(): Unit = {
+          override def tagAggForGpu(): Unit = {
             val dataType = a.child.dataType
             GpuOverrides.checkAndTagFloatAgg(dataType, conf, this)
           }
 
-          override def convertToGpu(child: Expression): GpuExpression = GpuAverage(child)
+          override def convertToGpu(childExprs: Seq[Expression]): GpuExpression = 
+            GpuAverage(childExprs.head)
+
+          // Average is not supported in ANSI mode right now, no matter the type
+          override val ansiTypeToCheck: Option[DataType] = None
         }),
       GpuOverrides.expr[Abs](
         "Absolute value",

sql-plugin/src/main/301until310-nondb/scala/com/nvidia/spark/rapids/shims/v2/SparkBaseShims.scala

@@ -273,12 +273,16 @@ abstract class SparkBaseShims extends Spark30XShims {
           TypeSig.DOUBLE, TypeSig.DOUBLE + TypeSig.DECIMAL_128_FULL,
           Seq(ParamCheck("input", TypeSig.integral + TypeSig.fp, TypeSig.numeric))),
         (a, conf, p, r) => new AggExprMeta[Average](a, conf, p, r) {
-          override def tagExprForGpu(): Unit = {
+          override def tagAggForGpu(): Unit = {
             val dataType = a.child.dataType
             GpuOverrides.checkAndTagFloatAgg(dataType, conf, this)
           }
 
-          override def convertToGpu(child: Expression): GpuExpression = GpuAverage(child)
+          override def convertToGpu(childExprs: Seq[Expression]): GpuExpression = 
+            GpuAverage(childExprs.head)
+
+          // Average is not supported in ANSI mode right now, no matter the type
+          override val ansiTypeToCheck: Option[DataType] = None
         }),
       GpuOverrides.expr[Abs](
         "Absolute value",

sql-plugin/src/main/311db/scala/com/nvidia/spark/rapids/shims/v2/SparkBaseShims.scala

@@ -210,12 +210,16 @@ abstract class SparkBaseShims extends Spark30XShims {
         TypeSig.DOUBLE, TypeSig.DOUBLE + TypeSig.DECIMAL_128_FULL,
         Seq(ParamCheck("input", TypeSig.integral + TypeSig.fp, TypeSig.numeric))),
       (a, conf, p, r) => new AggExprMeta[Average](a, conf, p, r) {
-        override def tagExprForGpu(): Unit = {
+        override def tagAggForGpu(): Unit = {
           val dataType = a.child.dataType
           GpuOverrides.checkAndTagFloatAgg(dataType, conf, this)
         }
 
-        override def convertToGpu(child: Expression): GpuExpression = GpuAverage(child)
+        override def convertToGpu(childExprs: Seq[Expression]): GpuExpression = 
+          GpuAverage(childExprs.head)
+
+        // Average is not supported in ANSI mode right now, no matter the type
+        override val ansiTypeToCheck: Option[DataType] = None
       }),
     GpuOverrides.expr[Abs](
       "Absolute value",

sql-plugin/src/main/311until320-nondb/scala/com/nvidia/spark/rapids/shims/v2/SparkBaseShims.scala

@@ -208,12 +208,16 @@ abstract class SparkBaseShims extends Spark30XShims {
         TypeSig.DOUBLE, TypeSig.DOUBLE + TypeSig.DECIMAL_128_FULL,
         Seq(ParamCheck("input", TypeSig.integral + TypeSig.fp, TypeSig.numeric))),
       (a, conf, p, r) => new AggExprMeta[Average](a, conf, p, r) {
-        override def tagExprForGpu(): Unit = {
+        override def tagAggForGpu(): Unit = {
           val dataType = a.child.dataType
           GpuOverrides.checkAndTagFloatAgg(dataType, conf, this)
         }
 
-        override def convertToGpu(child: Expression): GpuExpression = GpuAverage(child)
+        override def convertToGpu(childExprs: Seq[Expression]): GpuExpression = 
+          GpuAverage(childExprs.head)
+
+        // Average is not supported in ANSI mode right now, no matter the type
+        override val ansiTypeToCheck: Option[DataType] = None
       }),
     GpuOverrides.expr[Abs](
       "Absolute value",