DecimalType support for Aggregate Count (NVIDIA#1476)

Signed-off-by: Raza Jafri <rjafri@nvidia.com>
nartal1 · Jan 8, 2021 · 1958465 · 1958465
1 parent 8b4d931
commit 1958465
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/docs/supported_ops.md b/docs/supported_ops.md
@@ -15483,7 +15483,7 @@ Accelerator support is described below.
 <td>S</td>
 <td>S*</td>
 <td>S</td>
-<td><b>NS</b></td>
+<td>S*</td>
 <td>S</td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>
@@ -15526,7 +15526,7 @@ Accelerator support is described below.
 <td>S</td>
 <td>S*</td>
 <td>S</td>
-<td><b>NS</b></td>
+<td>S*</td>
 <td>S</td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>
@@ -15569,7 +15569,7 @@ Accelerator support is described below.
 <td>S</td>
 <td>S*</td>
 <td>S</td>
-<td><b>NS</b></td>
+<td>S*</td>
 <td>S</td>
 <td><b>NS</b></td>
 <td><b>NS</b></td>

diff --git a/integration_tests/src/main/python/hash_aggregate_test.py b/integration_tests/src/main/python/hash_aggregate_test.py
@@ -418,3 +418,10 @@ def test_arithmetic_reductions(data_gen):
                 'avg(a)'),
             conf = _no_nans_float_conf)
 
+@ignore_order
+@pytest.mark.parametrize('data_gen', all_gen, ids=idfn)
+@pytest.mark.parametrize('count_func', [f.count, f.countDistinct])
+def test_agg_count(data_gen, count_func):
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark : gen_df(spark, [('a', data_gen), ('b', data_gen)],
+                              length=1024).groupBy('a').agg(count_func("b")))
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
@@ -1648,7 +1648,7 @@ object GpuOverrides {
       ExprChecks.fullAgg(
         TypeSig.LONG, TypeSig.LONG,
         repeatingParamCheck = Some(RepeatingParamCheck(
-          "input", TypeSig.commonCudfTypes + TypeSig.NULL, TypeSig.all))),
+          "input", TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL, TypeSig.all))),
       (count, conf, p, r) => new ExprMeta[Count](count, conf, p, r) {
         override def tagExprForGpu(): Unit = {
           if (count.children.size > 1) {