update plans

Author: comphead

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -1180,14 +1180,12 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
-  // TODO enable once https://github.com/apache/datafusion-comet/issues/1267 is implemented
-  ignore("distinct") {
+  test("distinct") {
     withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
       Seq("native", "jvm").foreach { cometShuffleMode =>
         withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
           Seq(true, false).foreach { dictionary =>
             withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
-              val cometColumnShuffleEnabled = cometShuffleMode == "jvm"
               val table = "test"
               withTable(table) {
                 sql(s"create table $table(col1 int, col2 int, col3 int) using parquet")
@@ -1196,38 +1194,25 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
                 var expectedNumOfCometAggregates = 2
 
-                checkSparkAnswerAndNumOfAggregates(
-                  s"SELECT DISTINCT(col2) FROM $table",
-                  expectedNumOfCometAggregates)
-
-                expectedNumOfCometAggregates = 4
+                checkSparkAnswerAndOperator(s"SELECT DISTINCT(col2) FROM $table")
 
-                checkSparkAnswerAndNumOfAggregates(
-                  s"SELECT COUNT(distinct col2) FROM $table",
-                  expectedNumOfCometAggregates)
+                checkSparkAnswerAndOperator(s"SELECT COUNT(distinct col2) FROM $table")
 
-                checkSparkAnswerAndNumOfAggregates(
-                  s"SELECT COUNT(distinct col2), col1 FROM $table group by col1",
-                  expectedNumOfCometAggregates)
+                checkSparkAnswerAndOperator(
+                  s"SELECT COUNT(distinct col2), col1 FROM $table group by col1")
 
-                checkSparkAnswerAndNumOfAggregates(
-                  s"SELECT SUM(distinct col2) FROM $table",
-                  expectedNumOfCometAggregates)
+                checkSparkAnswerAndOperator(s"SELECT SUM(distinct col2) FROM $table")
 
-                checkSparkAnswerAndNumOfAggregates(
-                  s"SELECT SUM(distinct col2), col1 FROM $table group by col1",
-                  expectedNumOfCometAggregates)
+                checkSparkAnswerAndOperator(
+                  s"SELECT SUM(distinct col2), col1 FROM $table group by col1")
 
-                checkSparkAnswerAndNumOfAggregates(
+                checkSparkAnswerAndOperator(
                   "SELECT COUNT(distinct col2), SUM(distinct col2), col1, COUNT(distinct col2)," +
-                    s" SUM(distinct col2) FROM $table group by col1",
-                  expectedNumOfCometAggregates)
+                    s" SUM(distinct col2) FROM $table group by col1")
 
-                expectedNumOfCometAggregates = if (cometColumnShuffleEnabled) 2 else 1
-                checkSparkAnswerAndNumOfAggregates(
+                checkSparkAnswerAndOperator(
                   "SELECT COUNT(col2), MIN(col2), COUNT(DISTINCT col2), SUM(col2)," +
-                    s" SUM(DISTINCT col2), COUNT(DISTINCT col2), col1 FROM $table group by col1",
-                  expectedNumOfCometAggregates)
+                    s" SUM(DISTINCT col2), COUNT(DISTINCT col2), col1 FROM $table group by col1")
               }
             }
           }