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

dev/diffs/3.4.3.diff

@@ -1545,7 +1545,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..a5811d8e10d 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._
@@ -1852,7 +1852,15 @@ index 593bd7bb4ba..b327d84d5cc 100644
        }.size == 1)
        checkAnswer(ds, testData)
      }
-@@ -2028,7 +2064,8 @@ class AdaptiveQueryExecSuite
+@@ -1901,6 +1937,7 @@ class AdaptiveQueryExecSuite
+       df.collect()
+       assert(collect(df.queryExecution.executedPlan) {
+         case u: UnionExec => u
++        case u: CometUnionExec => u
+       }.size == numUnion)
+       assert(collect(df.queryExecution.executedPlan) {
+         case r: AQEShuffleReadExec => r
+@@ -2028,7 +2065,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1862,7 +1870,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 +2152,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1872,7 +1880,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2213,7 +2251,7 @@ class AdaptiveQueryExecSuite
+@@ -2213,7 +2252,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 +1889,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 +2265,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2")
              val shuffles2 = collect(adaptive2) {
@@ -1890,7 +1898,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
              }
              if (hasRequiredDistribution) {
                assert(shuffles2.size == 3)
-@@ -2260,7 +2298,8 @@ class AdaptiveQueryExecSuite
+@@ -2260,7 +2299,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1900,15 +1908,15 @@ index 593bd7bb4ba..b327d84d5cc 100644
      CostEvaluator.instantiate(
        classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
      intercept[IllegalArgumentException] {
-@@ -2404,6 +2443,7 @@ class AdaptiveQueryExecSuite
+@@ -2404,6 +2444,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 +2462,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1918,7 +1926,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 +2575,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value2 = value3")
          val shuffles1 = collect(adaptive1) {
@@ -1927,7 +1935,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
          }
          assert(shuffles1.size == 4)
          val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2544,7 +2585,7 @@ class AdaptiveQueryExecSuite
+@@ -2544,7 +2586,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value1 = value3")
          val shuffles2 = collect(adaptive2) {
@@ -1936,7 +1944,7 @@ index 593bd7bb4ba..b327d84d5cc 100644
          }
          assert(shuffles2.size == 4)
          val smj2 = findTopLevelSortMergeJoin(adaptive2)
-@@ -2703,7 +2744,8 @@ class AdaptiveQueryExecSuite
+@@ -2703,7 +2745,8 @@ class AdaptiveQueryExecSuite
      }
    }
 

dev/diffs/3.5.8.diff

@@ -1511,7 +1511,7 @@ index 5a413c77754..207b66e1d7b 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 2f8e401e743..7849c685b19 100644
+index 2f8e401e743..b6ccfe63778 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
 @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._
@@ -1825,7 +1825,15 @@ index 2f8e401e743..7849c685b19 100644
        }.size == 1)
        checkAnswer(ds, testData)
      }
-@@ -2054,7 +2091,8 @@ class AdaptiveQueryExecSuite
+@@ -1927,6 +1964,7 @@ class AdaptiveQueryExecSuite
+       df.collect()
+       assert(collect(df.queryExecution.executedPlan) {
+         case u: UnionExec => u
++        case u: CometUnionExec => u
+       }.size == numUnion)
+       assert(collect(df.queryExecution.executedPlan) {
+         case r: AQEShuffleReadExec => r
+@@ -2054,7 +2092,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1835,7 +1843,7 @@ index 2f8e401e743..7849c685b19 100644
      withTempView("t1", "t2") {
        def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
          Seq("100", "100000").foreach { size =>
-@@ -2140,7 +2178,8 @@ class AdaptiveQueryExecSuite
+@@ -2140,7 +2179,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1845,7 +1853,7 @@ index 2f8e401e743..7849c685b19 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2239,7 +2278,7 @@ class AdaptiveQueryExecSuite
+@@ -2239,7 +2279,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
              val shuffles1 = collect(adaptive1) {
@@ -1854,7 +1862,7 @@ index 2f8e401e743..7849c685b19 100644
              }
              assert(shuffles1.size == 3)
              // shuffles1.head is the top-level shuffle under the Aggregate operator
-@@ -2252,7 +2291,7 @@ class AdaptiveQueryExecSuite
+@@ -2252,7 +2292,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2")
              val shuffles2 = collect(adaptive2) {
@@ -1863,7 +1871,7 @@ index 2f8e401e743..7849c685b19 100644
              }
              if (hasRequiredDistribution) {
                assert(shuffles2.size == 3)
-@@ -2286,7 +2325,8 @@ class AdaptiveQueryExecSuite
+@@ -2286,7 +2326,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1873,15 +1881,15 @@ index 2f8e401e743..7849c685b19 100644
      CostEvaluator.instantiate(
        classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
      intercept[IllegalArgumentException] {
-@@ -2452,6 +2492,7 @@ class AdaptiveQueryExecSuite
+@@ -2452,6 +2493,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
-@@ -2469,7 +2510,8 @@ class AdaptiveQueryExecSuite
+@@ -2469,7 +2511,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1891,7 +1899,7 @@ index 2f8e401e743..7849c685b19 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
-@@ -2581,7 +2623,7 @@ class AdaptiveQueryExecSuite
+@@ -2581,7 +2624,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value2 = value3")
          val shuffles1 = collect(adaptive1) {
@@ -1900,7 +1908,7 @@ index 2f8e401e743..7849c685b19 100644
          }
          assert(shuffles1.size == 4)
          val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2592,7 +2634,7 @@ class AdaptiveQueryExecSuite
+@@ -2592,7 +2635,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value1 = value3")
          val shuffles2 = collect(adaptive2) {
@@ -1909,15 +1917,15 @@ index 2f8e401e743..7849c685b19 100644
          }
          assert(shuffles2.size == 4)
          val smj2 = findTopLevelSortMergeJoin(adaptive2)
-@@ -2850,6 +2892,7 @@ class AdaptiveQueryExecSuite
+@@ -2850,6 +2893,7 @@ class AdaptiveQueryExecSuite
          }.size == (if (firstAccess) 1 else 0))
          assert(collect(initialExecutedPlan) {
            case s: SortExec => s
 +          case s: CometSortExec => s
          }.size == (if (firstAccess) 2 else 0))
          assert(collect(initialExecutedPlan) {
            case i: InMemoryTableScanLike => i
-@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite
+@@ -2938,7 +2982,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -1927,7 +1935,7 @@ index 2f8e401e743..7849c685b19 100644
      val emptyDf = spark.range(1).where("false")
      val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1"))
      val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2"))
-@@ -2980,7 +3024,9 @@ class AdaptiveQueryExecSuite
+@@ -2980,7 +3025,9 @@ class AdaptiveQueryExecSuite
 
        val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
        assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec])

dev/diffs/4.0.2.diff

@@ -2135,7 +2135,7 @@ index a3cfdc5a240..3793b6191bf 100644
        })
        checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
 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 272be70f9fe..4a175083adf 100644
+index 272be70f9fe..0c38d555aaf 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
 @@ -28,12 +28,14 @@ import org.apache.spark.SparkException
@@ -2452,7 +2452,15 @@ index 272be70f9fe..4a175083adf 100644
        }.size == 1)
        checkAnswer(ds, testData)
      }
-@@ -2150,7 +2187,8 @@ class AdaptiveQueryExecSuite
+@@ -2023,6 +2060,7 @@ class AdaptiveQueryExecSuite
+       df.collect()
+       assert(collect(df.queryExecution.executedPlan) {
+         case u: UnionExec => u
++        case u: CometUnionExec => u
+       }.size == numUnion)
+       assert(collect(df.queryExecution.executedPlan) {
+         case r: AQEShuffleReadExec => r
+@@ -2150,7 +2188,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2462,7 +2470,7 @@ index 272be70f9fe..4a175083adf 100644
      withTempView("t1", "t2") {
        def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
          Seq("100", "100000").foreach { size =>
-@@ -2236,7 +2274,8 @@ class AdaptiveQueryExecSuite
+@@ -2236,7 +2275,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2472,7 +2480,7 @@ index 272be70f9fe..4a175083adf 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2335,7 +2374,7 @@ class AdaptiveQueryExecSuite
+@@ -2335,7 +2375,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
              val shuffles1 = collect(adaptive1) {
@@ -2481,7 +2489,7 @@ index 272be70f9fe..4a175083adf 100644
              }
              assert(shuffles1.size == 3)
              // shuffles1.head is the top-level shuffle under the Aggregate operator
-@@ -2348,7 +2387,7 @@ class AdaptiveQueryExecSuite
+@@ -2348,7 +2388,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2")
              val shuffles2 = collect(adaptive2) {
@@ -2490,7 +2498,7 @@ index 272be70f9fe..4a175083adf 100644
              }
              if (hasRequiredDistribution) {
                assert(shuffles2.size == 3)
-@@ -2382,7 +2421,8 @@ class AdaptiveQueryExecSuite
+@@ -2382,7 +2422,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2500,15 +2508,15 @@ index 272be70f9fe..4a175083adf 100644
      CostEvaluator.instantiate(
        classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
      intercept[IllegalArgumentException] {
-@@ -2548,6 +2588,7 @@ class AdaptiveQueryExecSuite
+@@ -2548,6 +2589,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
-@@ -2565,7 +2606,8 @@ class AdaptiveQueryExecSuite
+@@ -2565,7 +2607,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2518,7 +2526,7 @@ index 272be70f9fe..4a175083adf 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
-@@ -2677,7 +2719,7 @@ class AdaptiveQueryExecSuite
+@@ -2677,7 +2720,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value2 = value3")
          val shuffles1 = collect(adaptive1) {
@@ -2527,7 +2535,7 @@ index 272be70f9fe..4a175083adf 100644
          }
          assert(shuffles1.size == 4)
          val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2688,7 +2730,7 @@ class AdaptiveQueryExecSuite
+@@ -2688,7 +2731,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value1 = value3")
          val shuffles2 = collect(adaptive2) {
@@ -2536,23 +2544,23 @@ index 272be70f9fe..4a175083adf 100644
          }
          assert(shuffles2.size == 4)
          val smj2 = findTopLevelSortMergeJoin(adaptive2)
-@@ -2946,6 +2988,7 @@ class AdaptiveQueryExecSuite
+@@ -2946,6 +2989,7 @@ class AdaptiveQueryExecSuite
          }.size == (if (firstAccess) 1 else 0))
          assert(collect(initialExecutedPlan) {
            case s: SortExec => s
 +          case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
          }.size == (if (firstAccess) 2 else 0))
          assert(collect(initialExecutedPlan) {
            case i: InMemoryTableScanLike => i
-@@ -2958,6 +3001,7 @@ class AdaptiveQueryExecSuite
+@@ -2958,6 +3002,7 @@ class AdaptiveQueryExecSuite
          }.isEmpty)
          assert(collect(finalExecutedPlan) {
            case s: SortExec => s
 +          case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
          }.isEmpty)
          assert(collect(initialExecutedPlan) {
            case i: InMemoryTableScanLike => i
-@@ -3039,7 +3083,8 @@ class AdaptiveQueryExecSuite
+@@ -3039,7 +3084,8 @@ class AdaptiveQueryExecSuite
      }
    }
 
@@ -2562,7 +2570,7 @@ index 272be70f9fe..4a175083adf 100644
      val emptyDf = spark.range(1).where("false")
      val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1"))
      val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2"))
-@@ -3129,7 +3174,8 @@ class AdaptiveQueryExecSuite
+@@ -3129,7 +3175,8 @@ class AdaptiveQueryExecSuite
 
        val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
        assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec])