feat: re-enable COUNT for mixed Spark partial / Comet final aggregates

dev/diffs/3.4.3.diff

@@ -236,6 +236,30 @@ index 0efe0877e9b..423d3b3d76d 100644
  --
  -- SELECT_HAVING
  -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
+index c7c09bf7c79..5eaa5222142 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
+@@ -280,7 +280,8 @@ abstract class CTEInlineSuiteBase
+     }
+   }
+
+-  test("CTE Predicate push-down and column pruning") {
++  test("CTE Predicate push-down and column pruning",
++      IgnoreComet("Comet changes the exchange reuse count asserted by this test")) {
+     withTempView("t") {
+       Seq((0, 1), (1, 2)).toDF("c1", "c2").createOrReplaceTempView("t")
+       val df = sql(
+@@ -330,7 +331,8 @@ abstract class CTEInlineSuiteBase
+     }
+   }
+
+-  test("CTE Predicate push-down and column pruning - combined predicate") {
++  test("CTE Predicate push-down and column pruning - combined predicate",
++      IgnoreComet("Comet changes the exchange reuse count asserted by this test")) {
+     withTempView("t") {
+       Seq((0, 1, 2), (1, 2, 3)).toDF("c1", "c2", "c3").createOrReplaceTempView("t")
+       val df = sql(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
 index cf40e944c09..bdd5be4f462 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -1545,7 +1569,7 @@ index ac710c32296..2854b433dd3 100644
 
    import testImplicits._
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 593bd7bb4ba..b327d84d5cc 100644
+index 593bd7bb4ba..3ec90d8a44c 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 @@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._
@@ -1813,7 +1837,27 @@ index 593bd7bb4ba..b327d84d5cc 100644
      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
        val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
          "SELECT key FROM testData GROUP BY key")
-@@ -1599,7 +1628,7 @@ class AdaptiveQueryExecSuite
+@@ -1541,7 +1570,8 @@ class AdaptiveQueryExecSuite
+     }
+   }
+
+-  test("SPARK-35442: Support propagate empty relation through aggregate") {
++  test("SPARK-35442: Support propagate empty relation through aggregate",
++      IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) {
+     withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+       val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult(
+         "SELECT key, count(*) FROM testData WHERE value = 'no_match' GROUP BY key")
+@@ -1560,7 +1590,8 @@ class AdaptiveQueryExecSuite
+     }
+   }
+
+-  test("SPARK-35442: Support propagate empty relation through union") {
++  test("SPARK-35442: Support propagate empty relation through union",
++      IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) {
+     def checkNumUnion(plan: SparkPlan, numUnion: Int): Unit = {
+       assert(
+         collect(plan) {
+@@ -1599,7 +1630,7 @@ class AdaptiveQueryExecSuite
          val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
            "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id")
          assert(collect(adaptivePlan) {
@@ -1822,7 +1866,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
          }.length == 1)
        }
      }
-@@ -1679,7 +1708,8 @@ class AdaptiveQueryExecSuite
+@@ -1679,7 +1710,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1832,7 +1876,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
      def hasRepartitionShuffle(plan: SparkPlan): Boolean = {
        find(plan) {
          case s: ShuffleExchangeLike =>
-@@ -1864,6 +1894,9 @@ class AdaptiveQueryExecSuite
+@@ -1864,6 +1896,9 @@ class AdaptiveQueryExecSuite
      def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = {
        assert(collect(ds.queryExecution.executedPlan) {
          case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
@@ -1842,7 +1886,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
        }.size == 1)
        ds.collect()
        val plan = ds.queryExecution.executedPlan
-@@ -1872,6 +1905,9 @@ class AdaptiveQueryExecSuite
+@@ -1872,6 +1907,9 @@ class AdaptiveQueryExecSuite
        }.isEmpty)
        assert(collect(plan) {
          case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
@@ -1852,7 +1896,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
        }.size == 1)
        checkAnswer(ds, testData)
      }
-@@ -2028,7 +2064,8 @@ class AdaptiveQueryExecSuite
+@@ -2028,7 +2066,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1862,7 +1906,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
      withTempView("t1", "t2") {
        def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
          Seq("100", "100000").foreach { size =>
-@@ -2114,7 +2151,8 @@ class AdaptiveQueryExecSuite
+@@ -2114,7 +2153,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1872,7 +1916,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2213,7 +2251,7 @@ class AdaptiveQueryExecSuite
+@@ -2213,7 +2253,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
              val shuffles1 = collect(adaptive1) {
@@ -1881,7 +1925,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
              }
              assert(shuffles1.size == 3)
              // shuffles1.head is the top-level shuffle under the Aggregate operator
-@@ -2226,7 +2264,7 @@ class AdaptiveQueryExecSuite
+@@ -2226,7 +2266,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2")
              val shuffles2 = collect(adaptive2) {
@@ -1890,7 +1934,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
              }
              if (hasRequiredDistribution) {
                assert(shuffles2.size == 3)
-@@ -2260,7 +2298,8 @@ class AdaptiveQueryExecSuite
+@@ -2260,7 +2300,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1900,15 +1944,15 @@ index 593bd7bb4ba..b327d84d5cc 100644
      CostEvaluator.instantiate(
        classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
      intercept[IllegalArgumentException] {
-@@ -2404,6 +2443,7 @@ class AdaptiveQueryExecSuite
+@@ -2404,6 +2445,7 @@ class AdaptiveQueryExecSuite
            val (_, adaptive) = runAdaptiveAndVerifyResult(query)
            assert(adaptive.collect {
              case sort: SortExec => sort
 +            case sort: CometSortExec => sort
            }.size == 1)
            val read = collect(adaptive) {
              case read: AQEShuffleReadExec => read
-@@ -2421,7 +2461,8 @@ class AdaptiveQueryExecSuite
+@@ -2421,7 +2463,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1918,7 +1962,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
-@@ -2533,7 +2574,7 @@ class AdaptiveQueryExecSuite
+@@ -2533,7 +2576,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value2 = value3")
          val shuffles1 = collect(adaptive1) {
@@ -1927,7 +1971,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
          }
          assert(shuffles1.size == 4)
          val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2544,7 +2585,7 @@ class AdaptiveQueryExecSuite
+@@ -2544,7 +2587,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value1 = value3")
          val shuffles2 = collect(adaptive2) {
@@ -1936,7 +1980,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
          }
          assert(shuffles2.size == 4)
          val smj2 = findTopLevelSortMergeJoin(adaptive2)
-@@ -2703,7 +2744,8 @@ class AdaptiveQueryExecSuite
+@@ -2703,7 +2746,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2408,7 +2452,7 @@ index bf5c51b89bb..f7402b7d883 100644
        val e = testSchemaMismatch(dir.getCanonicalPath, vectorizedReaderEnabled = true)
        assert(e.getCause.isInstanceOf[SparkException])
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
-index 3a0bd35cb70..b28f06a757f 100644
+index 3a0bd35cb70..99b70606261 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
 @@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.debug
@@ -2419,7 +2463,27 @@ index 3a0bd35cb70..b28f06a757f 100644
  import org.apache.spark.sql.catalyst.InternalRow
  import org.apache.spark.sql.catalyst.expressions.Attribute
  import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
-@@ -124,7 +125,8 @@ class DebuggingSuite extends DebuggingSuiteBase with DisableAdaptiveExecutionSui
+@@ -41,7 +42,8 @@ abstract class DebuggingSuiteBase extends SharedSparkSession {
+     testData.as[TestData].debug()
+   }
+
+-  test("debugCodegen") {
++  test("debugCodegen",
++      IgnoreComet("Comet changes the WholeStageCodegen subtree count")) {
+     val df = spark.range(10).groupBy(col("id") * 2).count()
+     df.collect()
+     val res = codegenString(df.queryExecution.executedPlan)
+@@ -50,7 +52,8 @@ abstract class DebuggingSuiteBase extends SharedSparkSession {
+     assert(res.contains("Object[]"))
+   }
+
+-  test("debugCodegenStringSeq") {
++  test("debugCodegenStringSeq",
++      IgnoreComet("Comet changes the WholeStageCodegen subtree count")) {
+     val df = spark.range(10).groupBy(col("id") * 2).count()
+     df.collect()
+     val res = codegenStringSeq(df.queryExecution.executedPlan)
+@@ -124,7 +127,8 @@ class DebuggingSuite extends DebuggingSuiteBase with DisableAdaptiveExecutionSui
           | id LongType: {}""".stripMargin))
    }
 
@@ -2529,6 +2593,28 @@ index d083cac48ff..3c11bcde807 100644
 
    import testImplicits._
 
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
+index fdc633f3556..f1e1751171a 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
+@@ -36,6 +36,7 @@ import org.apache.spark.internal.config.Status._
+ import org.apache.spark.rdd.RDD
+ import org.apache.spark.resource.ResourceProfile
+ import org.apache.spark.scheduler._
++import org.apache.spark.sql.IgnoreComet
+ import org.apache.spark.sql.{DataFrame, SparkSession}
+ import org.apache.spark.sql.catalyst.InternalRow
+ import org.apache.spark.sql.catalyst.expressions.Attribute
+@@ -699,7 +700,8 @@ abstract class SQLAppStatusListenerSuite extends SharedSparkSession with JsonTes
+   }
+
+   test("SPARK-29894 test Codegen Stage Id in SparkPlanInfo",
+-    DisableAdaptiveExecution("WSCG rule is applied later in AQE")) {
++    DisableAdaptiveExecution("WSCG rule is applied later in AQE"),
++    IgnoreComet("Comet changes the WholeStageCodegen subtree count")) {
+     // with AQE on, the WholeStageCodegen rule is applied when running QueryStageExec.
+     val df = createTestDataFrame.select(count("*"))
+     val sparkPlanInfo = SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
 index 266bb343526..f8ad838e2b2 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala