Support DecimalType for approx_count_distinct

apache · May 25, 2024 · f99168c · f99168c
1 parent 9f424a1
commit f99168c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/HLLRewriteRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/HLLRewriteRule.scala
@@ -73,6 +73,7 @@ case class HLLRewriteRule(spark: SparkSession) extends Rule[LogicalPlan] {
       case LongType => true
       case ShortType => true
       case StringType => true
+      case _: DecimalType => true
       case _ => false
     }
   }

diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
@@ -571,6 +571,26 @@ abstract class VeloxAggregateFunctionsSuite extends VeloxWholeStageTransformerSu
     }
   }
 
+  test("approx_count_distinct decimal") {
+    // The data type of l_discount is decimal.
+    runQueryAndCompare("""
+                         |select approx_count_distinct(l_discount) from lineitem;
+                         |""".stripMargin) {
+      checkGlutenOperatorMatch[HashAggregateExecTransformer]
+    }
+    runQueryAndCompare(
+      "select approx_count_distinct(l_discount), count(distinct l_orderkey) from lineitem") {
+      df =>
+        {
+          assert(
+            getExecutedPlan(df).count(
+              plan => {
+                plan.isInstanceOf[HashAggregateExecTransformer]
+              }) == 0)
+        }
+    }
+  }
+
   test("max_by") {
     runQueryAndCompare(s"""
                           |select max_by(l_linenumber, l_comment) from lineitem;