feat: do not fallback to Spark for distincts

Author: comphead

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -564,7 +564,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
     // COUNT(DISTINCT x, y) - not supported
     if (aggExpr.isDistinct
       &&
-      (aggExpr.aggregateFunction.prettyName == "count" &&
+      !(aggExpr.aggregateFunction.prettyName == "count" &&
         aggExpr.aggregateFunction.children.length == 1)) {
       withInfo(aggExpr, s"Distinct aggregate not supported for: $aggExpr")
       return None

spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala

@@ -19,12 +19,14 @@
 
 package org.apache.comet
 
+import org.apache.comet.DataTypeSupport.isComplexType
+
 class CometFuzzAggregateSuite extends CometFuzzTestBase {
 
-  test("count distinct") {
+  test("count distinct - simple columns") {
     val df = spark.read.parquet(filename)
     df.createOrReplaceTempView("t1")
-    for (col <- df.columns) {
+    for (col <- df.schema.fields.filterNot(f => isComplexType(f.dataType)).map(_.name)) {
       val sql = s"SELECT count(distinct $col) FROM t1"
       val (_, cometPlan) = checkSparkAnswer(sql)
       if (usingDataSourceExec) {
@@ -35,10 +37,40 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase {
     }
   }
 
-  test("count distinct group by multiple column") {
+  // Aggregate by complex columns not yet supported
+  // https://github.com/apache/datafusion-comet/issues/2382
+  test("count distinct - complex columns") {
     val df = spark.read.parquet(filename)
     df.createOrReplaceTempView("t1")
-    for (col <- df.columns) {
+    for (col <- df.schema.fields.filter(f => isComplexType(f.dataType)).map(_.name)) {
+      val sql = s"SELECT count(distinct $col) FROM t1"
+      val (_, cometPlan) = checkSparkAnswer(sql)
+      if (usingDataSourceExec) {
+        assert(1 == collectNativeScans(cometPlan).length)
+      }
+    }
+  }
+
+  test("count distinct group by multiple column - simple columns ") {
+    val df = spark.read.parquet(filename)
+    df.createOrReplaceTempView("t1")
+    for (col <- df.schema.fields.filterNot(f => isComplexType(f.dataType)).map(_.name)) {
+      val sql = s"SELECT c1, c2, c3, count(distinct $col) FROM t1 group by c1, c2, c3"
+      val (_, cometPlan) = checkSparkAnswer(sql)
+      if (usingDataSourceExec) {
+        assert(1 == collectNativeScans(cometPlan).length)
+      }
+
+      checkSparkAnswerAndOperator(sql)
+    }
+  }
+
+  // Aggregate by complex columns not yet supported
+  // https://github.com/apache/datafusion-comet/issues/2382
+  test("count distinct group by multiple column - complex columns ") {
+    val df = spark.read.parquet(filename)
+    df.createOrReplaceTempView("t1")
+    for (col <- df.schema.fields.filter(f => isComplexType(f.dataType)).map(_.name)) {
       val sql = s"SELECT c1, c2, c3, count(distinct $col) FROM t1 group by c1, c2, c3"
       val (_, cometPlan) = checkSparkAnswer(sql)
       if (usingDataSourceExec) {
@@ -47,6 +79,8 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase {
     }
   }
 
+  // Not yet supported
+  // https://github.com/apache/datafusion-comet/issues/2292
   test("count distinct multiple values and group by multiple column") {
     val df = spark.read.parquet(filename)
     df.createOrReplaceTempView("t1")