diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index a905f32e3c..87844d0e12 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -1473,26 +1473,31 @@ index 30ce940b032..0d3f6c6c934 100644 private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala -index 47679ed7865..9ffbaecb98e 100644 +index 47679ed7865..cbcabe00992 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{DataFrame, QueryTest} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf -@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase +@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = { val plan = df.queryExecution.executedPlan assert(collectWithSubqueries(plan) { - case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s + case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s }.length == hashAggCount) - assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) +- assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) ++ assert(collectWithSubqueries(plan) { ++ case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s ++ }.length == sortAggCount) } + + private def checkAggs( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala index eec396b2e39..bf3f1c769d6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala @@ -1563,7 +1568,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..32af28b0238 100644 +index 593bd7bb4ba..2ae8b702068 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._ @@ -1623,7 +1628,7 @@ index 593bd7bb4ba..32af28b0238 100644 private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = { collect(plan) { case agg: BaseAggregateExec => agg -+ case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] ++ case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] } } @@ -2856,14 +2861,14 @@ index ef5b8a769fe..84fe1bfabc9 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala -index b4c4ec7acbf..20579284856 100644 +index b4c4ec7acbf..e57aff1f170 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils import org.scalatest.Assertions import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.CometBaseAggregateExec import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec} import org.apache.spark.sql.functions.count @@ -2871,7 +2876,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = query.lastExecution.executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { @@ -2879,7 +2884,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 0cbf239981..d822091229 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -1439,26 +1439,31 @@ index 005e764cc30..92ec088efab 100644 private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala -index 47679ed7865..9ffbaecb98e 100644 +index 47679ed7865..cbcabe00992 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{DataFrame, QueryTest} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf -@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase +@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = { val plan = df.queryExecution.executedPlan assert(collectWithSubqueries(plan) { - case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s + case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s }.length == hashAggCount) - assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) +- assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) ++ assert(collectWithSubqueries(plan) { ++ case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s ++ }.length == sortAggCount) } + + private def checkAggs( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala index eec396b2e39..bf3f1c769d6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala @@ -1529,7 +1534,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..dbcf3171946 100644 +index 2f8e401e743..fc712552c4f 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._ @@ -1589,7 +1594,7 @@ index 2f8e401e743..dbcf3171946 100644 private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = { collect(plan) { case agg: BaseAggregateExec => agg -+ case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] ++ case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] } } @@ -2806,14 +2811,14 @@ index c97979a57a5..45a998db0e0 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala -index b4c4ec7acbf..20579284856 100644 +index b4c4ec7acbf..e57aff1f170 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils import org.scalatest.Assertions import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.CometBaseAggregateExec import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec} import org.apache.spark.sql.functions.count @@ -2821,7 +2826,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = query.lastExecution.executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { @@ -2829,7 +2834,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff index f6d01f4f6a..d2244acc24 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -1841,26 +1841,31 @@ index 005e764cc30..92ec088efab 100644 private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala -index 47679ed7865..9ffbaecb98e 100644 +index 47679ed7865..cbcabe00992 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{DataFrame, QueryTest} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf -@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase +@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = { val plan = df.queryExecution.executedPlan assert(collectWithSubqueries(plan) { - case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s + case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s }.length == hashAggCount) - assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) +- assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) ++ assert(collectWithSubqueries(plan) { ++ case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s ++ }.length == sortAggCount) } + + private def checkAggs( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala index aa619c5cde8..4b5a42e683f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala @@ -2153,7 +2158,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..12daa1f5932 100644 +index 272be70f9fe..a246594c73c 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 @@ -2216,7 +2221,7 @@ index 272be70f9fe..12daa1f5932 100644 private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = { collect(plan) { case agg: BaseAggregateExec => agg -+ case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] ++ case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] } } @@ -3470,14 +3475,14 @@ index b0967d5ffdf..3d567f913de 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala -index b4c4ec7acbf..20579284856 100644 +index b4c4ec7acbf..e57aff1f170 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils import org.scalatest.Assertions import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.CometBaseAggregateExec import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec} import org.apache.spark.sql.functions.count @@ -3485,7 +3490,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = query.lastExecution.executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { @@ -3493,7 +3498,7 @@ index b4c4ec7acbf..20579284856 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff index e084bc9913..65a278bcb9 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -1964,26 +1964,31 @@ index 005e764cc30..92ec088efab 100644 private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala -index 47679ed7865..9ffbaecb98e 100644 +index 47679ed7865..cbcabe00992 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{DataFrame, QueryTest} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf -@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase +@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = { val plan = df.queryExecution.executedPlan assert(collectWithSubqueries(plan) { - case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s + case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s }.length == hashAggCount) - assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) +- assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) ++ assert(collectWithSubqueries(plan) { ++ case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s ++ }.length == sortAggCount) } + + private def checkAggs( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala index aa619c5cde8..4b5a42e683f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala @@ -2276,7 +2281,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 3e7d26f74bd..7e70e72fa3e 100644 +index 3e7d26f74bd..3a4f383777e 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,12 +27,14 @@ import org.apache.spark.SparkException @@ -2339,7 +2344,7 @@ index 3e7d26f74bd..7e70e72fa3e 100644 private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = { collect(plan) { case agg: BaseAggregateExec => agg -+ case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] ++ case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec] } } @@ -3713,14 +3718,14 @@ index e57c4e1e665..1dd8b80a006 100644 // Necessary to make ScalaTest 3.x interrupt a thread on the JVM like ScalaTest 2.2.x implicit val defaultSignaler: Signaler = ThreadSignaler diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala -index 465da3cd469..92ac998929d 100644 +index 465da3cd469..2e6bcb6ca0d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala @@ -22,6 +22,7 @@ import java.io.File import org.scalatest.Assertions import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.CometBaseAggregateExec import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.execution.streaming.operators.stateful.{StateStoreRestoreExec, StateStoreSaveExec} import org.apache.spark.sql.execution.streaming.runtime.MemoryStream @@ -3728,7 +3733,7 @@ index 465da3cd469..92ac998929d 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = query.lastExecution.executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { @@ -3736,7 +3741,7 @@ index 465da3cd469..92ac998929d 100644 // verify aggregations in between, except partial aggregation val allAggregateExecs = executedPlan.collect { case a: BaseAggregateExec => a -+ case c: CometHashAggregateExec => c.originalPlan ++ case c: CometBaseAggregateExec => c.originalPlan } val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 52f39c59ad..a97c22aace 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, Comet import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat @@ -80,6 +80,7 @@ object CometExecRule { classOf[GenerateExec] -> CometExplodeExec, classOf[HashAggregateExec] -> CometHashAggregateExec, classOf[ObjectHashAggregateExec] -> CometObjectHashAggregateExec, + classOf[SortAggregateExec] -> CometSortAggregateExec, classOf[BroadcastHashJoinExec] -> CometBroadcastHashJoinExec, classOf[BroadcastNestedLoopJoinExec] -> CometBroadcastNestedLoopJoinExec, classOf[ShuffledHashJoinExec] -> CometHashJoinExec, @@ -150,8 +151,7 @@ case class CometExecRule(session: SparkSession) * Comet columnar shuffle. */ private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = { - def isAggregate(p: SparkPlan): Boolean = - p.isInstanceOf[HashAggregateExec] || p.isInstanceOf[ObjectHashAggregateExec] + def isAggregate(p: SparkPlan): Boolean = p.isInstanceOf[BaseAggregateExec] def isRedundantShuffle(child: SparkPlan): Boolean = child match { case s: CometShuffleExchangeExec => @@ -881,9 +881,13 @@ case class CometExecRule(session: SparkSession) val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]] if (!isOperatorEnabled(serde, agg.asInstanceOf[SparkPlan])) return false - // ObjectHashAggregate has an extra shuffle-enabled guard in its convert method + // ObjectHashAggregate / SortAggregate carry TypedImperativeAggregate functions whose + // intermediate buffer formats differ between Spark and Comet, so the Partial->Final pair + // must travel via Comet shuffle. agg match { - case _: ObjectHashAggregateExec if !isCometShuffleEnabled(agg.conf) => return false + case _: ObjectHashAggregateExec | _: SortAggregateExec + if !isCometShuffleEnabled(agg.conf) => + return false case _ => } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index e4d6b53770..3aaf5983d3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -1515,6 +1515,63 @@ case class CometUnionExec( trait CometBaseAggregate { + /** + * Shared support-level check used by every aggregate serde: honor the unit-test knobs that + * selectively disable Comet conversion for partial or final aggregates. + */ + protected def baseAggregateSupportLevel(op: BaseAggregateExec): SupportLevel = { + if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) && + op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) { + return Unsupported(Some("Partial aggregates disabled via test config")) + } + if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) && + op.aggregateExpressions.exists(_.mode == Final)) { + return Unsupported(Some("Final aggregates disabled via test config")) + } + Compatible() + } + + /** + * For Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet), + * the Spark-side output declares buffer columns as BinaryType (Spark serializes state to + * binary). However, the native Comet aggregate produces the actual state type (e.g., + * ArrayType(elementType) for CollectSet). This corrects the output schema to match the native + * state types so the shuffle exchange schema is consistent with the actual data. + * + * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a + * case branch here mapping it to the native state type. + */ + protected def adjustOutputForNativeState(op: BaseAggregateExec): Seq[Attribute] = { + val modes = op.aggregateExpressions.map(_.mode).distinct + if (modes != Seq(Partial)) { + return op.output + } + + val numGrouping = op.groupingExpressions.length + val output = op.output.toArray + + var bufferIdx = numGrouping + for (aggExpr <- op.aggregateExpressions) { + val aggFunc = aggExpr.aggregateFunction + val bufferAttrs = aggFunc.aggBufferAttributes + aggFunc match { + case cs: CollectSet => + val elementType = cs.children.head.dataType + val nativeStateType = ArrayType(elementType, containsNull = true) + output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) + case _: Percentile => + // DataFusion's percentile_cont keeps all values in a List partial state. + // Comet casts the child to double, so the native state is ArrayType(DoubleType). + val nativeStateType = ArrayType(DoubleType, containsNull = true) + output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) + case _ => + } + bufferIdx += bufferAttrs.length + } + + output.toSeq + } + def doConvert( aggregate: BaseAggregateExec, builder: Operator.Builder, @@ -1776,21 +1833,19 @@ trait CometBaseAggregate { } /** - * Find the first Comet partial aggregate in the plan. If it reaches a Spark HashAggregate with - * partial or partial-merge mode, it will return None. + * Find the first Comet partial aggregate in the plan. If it reaches a Spark BaseAggregateExec + * (HashAggregate / ObjectHashAggregate / SortAggregate) with partial or partial-merge mode, it + * will return None. */ - private def findCometPartialAgg(plan: SparkPlan): Option[CometHashAggregateExec] = { + private def findCometPartialAgg(plan: SparkPlan): Option[CometBaseAggregateExec] = { def isPartialOrMerge(mode: AggregateMode): Boolean = mode == Partial || mode == PartialMerge plan.collectFirst { - case agg: CometHashAggregateExec + case agg: CometBaseAggregateExec if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) => Some(agg) - case agg: HashAggregateExec - if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) => - None - case agg: ObjectHashAggregateExec + case agg: BaseAggregateExec if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) => None case a: AQEShuffleReadExec => findCometPartialAgg(a.child) @@ -1807,19 +1862,8 @@ object CometHashAggregateExec override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_AGGREGATE_ENABLED) - override def getSupportLevel(op: HashAggregateExec): SupportLevel = { - // some unit tests need to disable partial or final hash aggregate support to test that - // CometExecRule does not allow mixed Spark/Comet aggregates - if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) && - op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) { - return Unsupported(Some("Partial aggregates disabled via test config")) - } - if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) && - op.aggregateExpressions.exists(_.mode == Final)) { - return Unsupported(Some("Final aggregates disabled via test config")) - } - Compatible() - } + override def getSupportLevel(op: HashAggregateExec): SupportLevel = + baseAggregateSupportLevel(op) override def convert( aggregate: HashAggregateExec, @@ -1849,19 +1893,8 @@ object CometObjectHashAggregateExec override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_AGGREGATE_ENABLED) - override def getSupportLevel(op: ObjectHashAggregateExec): SupportLevel = { - // Mirror the same test-knobs as CometHashAggregateExec so that mixed-execution - // unit tests can selectively disable partial or final ObjectHashAggregateExec conversion. - if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) && - op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) { - return Unsupported(Some("Partial aggregates disabled via test config")) - } - if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) && - op.aggregateExpressions.exists(_.mode == Final)) { - return Unsupported(Some("Final aggregates disabled via test config")) - } - Compatible() - } + override def getSupportLevel(op: ObjectHashAggregateExec): SupportLevel = + baseAggregateSupportLevel(op) override def convert( aggregate: ObjectHashAggregateExec, @@ -1890,73 +1923,74 @@ object CometObjectHashAggregateExec op.child, SerializedPlan(None)) } +} - /** - * For Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet), - * the Spark-side output declares buffer columns as BinaryType (since Spark serializes state to - * binary). However, the native Comet aggregate produces the actual state type (e.g., - * ArrayType(elementType) for CollectSet). This method corrects the output schema to match the - * native state types so the shuffle exchange schema is consistent with the actual data. - * - * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a - * case branch here mapping it to the native state type. - */ - private def adjustOutputForNativeState(op: ObjectHashAggregateExec): Seq[Attribute] = { - // This adjustment only applies to pure-Partial aggregates (checked below). - val modes = op.aggregateExpressions.map(_.mode).distinct - if (modes != Seq(Partial)) { - return op.output - } +object CometSortAggregateExec + extends CometOperatorSerde[SortAggregateExec] + with CometBaseAggregate { - val numGrouping = op.groupingExpressions.length - val output = op.output.toArray + override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( + CometConf.COMET_EXEC_AGGREGATE_ENABLED) - var bufferIdx = numGrouping - for (aggExpr <- op.aggregateExpressions) { - val aggFunc = aggExpr.aggregateFunction - val bufferAttrs = aggFunc.aggBufferAttributes - aggFunc match { - case cs: CollectSet => - val elementType = cs.children.head.dataType - val nativeStateType = ArrayType(elementType, containsNull = true) - output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) - case _: Percentile => - // DataFusion's percentile_cont keeps all values in a List partial state. - // Comet casts the child to double, so the native state is ArrayType(DoubleType). - val nativeStateType = ArrayType(DoubleType, containsNull = true) - output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType) - case _ => - } - bufferIdx += bufferAttrs.length + override def getSupportLevel(op: SortAggregateExec): SupportLevel = + baseAggregateSupportLevel(op) + + override def convert( + aggregate: SortAggregateExec, + builder: Operator.Builder, + childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { + + // SortAggregate is planned for TypedImperativeAggregate functions whose intermediate + // buffer formats differ between Spark and Comet (same risk as ObjectHashAggregate). + // Require Comet shuffle so a Partial->Final pair never spans the JVM/native boundary. + if (!isCometShuffleEnabled(aggregate.conf)) { + return None } - output.toSeq + doConvert(aggregate, builder, childOp: _*) + } + + override def createExec(nativeOp: Operator, op: SortAggregateExec): CometNativeExec = { + // The native AggregateExec auto-detects Sorted input mode from the child's output ordering + // and produces output sorted by the grouping keys; CometExec.outputOrdering defaults to + // originalPlan.outputOrdering, which is SortAggregateExec's grouping-key ordering, so + // downstream operators that elided a sort against it still see a satisfying ordering. + CometSortAggregateExec( + nativeOp, + op, + adjustOutputForNativeState(op), + op.groupingExpressions, + op.aggregateExpressions, + op.resultExpressions, + op.child.output, + op.child, + SerializedPlan(None)) } } -case class CometHashAggregateExec( - override val nativeOp: Operator, - override val originalPlan: SparkPlan, - override val output: Seq[Attribute], - groupingExpressions: Seq[NamedExpression], - aggregateExpressions: Seq[AggregateExpression], - resultExpressions: Seq[NamedExpression], - input: Seq[Attribute], - child: SparkPlan, - override val serializedPlanOpt: SerializedPlan) +/** + * Common base for Comet's aggregate wrapper operators. The hash-based and sort-based variants + * share the same native AggregateExec serialization and rendering; they are kept as distinct plan + * node types only so the executed plan reflects whether Spark planned a HashAggregateExec / + * ObjectHashAggregateExec or a SortAggregateExec (the native AggregateExec auto-detects sorted + * input mode from the child ordering, so execution is otherwise identical). + */ +abstract class CometBaseAggregateExec extends CometUnaryExec with PartitioningPreservingUnaryExecNode { + def groupingExpressions: Seq[NamedExpression] + def aggregateExpressions: Seq[AggregateExpression] + def resultExpressions: Seq[NamedExpression] + def input: Seq[Attribute] + // The aggExprs could be empty. For example, if the aggregate functions only have // distinct aggregate functions or only have group by, the aggExprs is empty and // modes is empty too. - val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct + lazy val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct override def producedAttributes: AttributeSet = outputSet ++ AttributeSet(resultExpressions) - override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = - this.copy(child = newChild) - override def verboseStringWithOperatorId(): String = { s""" |$formattedNodeName @@ -1969,6 +2003,24 @@ case class CometHashAggregateExec( override def stringArgs: Iterator[Any] = Iterator(input, modes, groupingExpressions, aggregateExpressions, child) + override protected def outputExpressions: Seq[NamedExpression] = resultExpressions +} + +case class CometHashAggregateExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + groupingExpressions: Seq[NamedExpression], + aggregateExpressions: Seq[AggregateExpression], + resultExpressions: Seq[NamedExpression], + input: Seq[Attribute], + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometBaseAggregateExec { + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + this.copy(child = newChild) + override def equals(obj: Any): Boolean = { obj match { case other: CometHashAggregateExec => @@ -1986,8 +2038,40 @@ case class CometHashAggregateExec( override def hashCode(): Int = Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child) +} - override protected def outputExpressions: Seq[NamedExpression] = resultExpressions +case class CometSortAggregateExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + groupingExpressions: Seq[NamedExpression], + aggregateExpressions: Seq[AggregateExpression], + resultExpressions: Seq[NamedExpression], + input: Seq[Attribute], + child: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometBaseAggregateExec { + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + this.copy(child = newChild) + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometSortAggregateExec => + this.output == other.output && + this.groupingExpressions == other.groupingExpressions && + this.aggregateExpressions == other.aggregateExpressions && + this.input == other.input && + this.modes == other.modes && + this.child == other.child && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => + false + } + } + + override def hashCode(): Int = + Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child) } trait CometHashJoin { diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql index 8aedc4f5f1..18c7e1968f 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql @@ -70,11 +70,16 @@ CREATE TABLE test_types( grp string ) USING parquet +-- first/last IGNORE NULLS are non-deterministic when a group has more than one non-null value, +-- because the result depends on intra-group row order, which is not defined and can differ +-- between Spark and Comet. To keep the multi-type queries below comparable across engines, each +-- group has at most one non-null value per column (surrounded by nulls so null-skipping is still +-- exercised); group 'b' is entirely null. statement INSERT INTO test_types VALUES (NULL, NULL, NULL, NULL, NULL, 'a'), (1, 100, 1.5, 'foo', true, 'a'), - (2, 200, 2.5, 'bar', false, 'a'), + (NULL, NULL, NULL, NULL, NULL, 'a'), (NULL, NULL, NULL, NULL, NULL, 'b'), (NULL, NULL, NULL, NULL, NULL, 'b') @@ -183,7 +188,9 @@ SELECT first(val) IGNORE NULLS FROM test_single_null -- first IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +-- Spark plans this as SortAggregateExec (the string buffer is not hash-friendly), which Comet +-- now runs natively. +query SELECT grp, first(i_val) IGNORE NULLS, first(l_val) IGNORE NULLS, @@ -327,7 +334,9 @@ SELECT last(val) IGNORE NULLS FROM test_single_null -- last IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +-- Spark plans this as SortAggregateExec (the string buffer is not hash-friendly), which Comet +-- now runs natively. +query SELECT grp, last(i_val) IGNORE NULLS, last(l_val) IGNORE NULLS, diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql index dbec88e2f6..5aacc9a10f 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql @@ -21,7 +21,10 @@ CREATE TABLE test_min_max(i int, d double, s string, grp string) USING parquet statement INSERT INTO test_min_max VALUES (1, 1.5, 'b', 'x'), (3, 3.5, 'a', 'x'), (2, 2.5, 'c', 'y'), (NULL, NULL, NULL, 'y'), (-1, -1.5, 'z', 'x') -query expect_fallback(SortAggregate is not supported) +-- min/max over StringType is not yet supported natively, so this falls back to Spark. +-- (Spark plans this global aggregate as SortAggregateExec because the string min/max buffer +-- is not hash-friendly; that operator is now supported, but the string aggregate itself is not.) +query expect_fallback(Unsupported data type: StringType) SELECT min(i), max(i), min(d), max(d), min(s), max(s) FROM test_min_max query diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql new file mode 100644 index 0000000000..18b0dac7a0 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql @@ -0,0 +1,214 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Disabling ObjectHashAggregate forces Spark to plan SortAggregateExec for the +-- TypedImperativeAggregate function collect_set, which is the operator this file exercises. +-- (The plan-shape assertion that a SortAggregateExec is actually produced lives in +-- CometAggregateSuite; here the default `query` mode asserts the whole plan, including the +-- pre-aggregate sorts and the Comet shuffle, runs natively and matches Spark.) +-- collect_set returns elements in nondeterministic order, so every result is wrapped in +-- sort_array for a stable comparison against Spark. +-- Config: spark.sql.execution.useObjectHashAggregateExec=false +-- Config: spark.comet.exec.strictFloatingPoint=true + +-- ============================================================ +-- Setup: tables +-- ============================================================ + +statement +CREATE TABLE sa_int(i int, g string) USING parquet + +statement +INSERT INTO sa_int VALUES + (1, 'a'), (2, 'a'), (1, 'a'), (3, 'a'), + (4, 'b'), (4, 'b'), (NULL, 'b'), (5, 'b'), + (NULL, 'c'), (NULL, 'c') + +statement +CREATE TABLE sa_multikey(a int, k1 string, k2 int) USING parquet + +statement +INSERT INTO sa_multikey VALUES + (1, 'x', 10), (2, 'x', 10), (1, 'x', 10), + (3, 'x', 20), (4, 'y', 10), (NULL, 'y', 10) + +statement +CREATE TABLE sa_nulls(v int, g string) USING parquet + +statement +INSERT INTO sa_nulls VALUES (NULL, 'a'), (NULL, 'a'), (NULL, 'b'), (1, 'b') + +statement +CREATE TABLE sa_empty(v int, g string) USING parquet + +statement +CREATE TABLE sa_single(v int) USING parquet + +statement +INSERT INTO sa_single VALUES (42) + +-- ============================================================ +-- Global aggregate (no GROUP BY): single SortAggregate group +-- ============================================================ + +query +SELECT sort_array(collect_set(i)) FROM sa_int + +-- ============================================================ +-- Single grouping key +-- ============================================================ + +query +SELECT g, sort_array(collect_set(i)) FROM sa_int GROUP BY g ORDER BY g + +-- ============================================================ +-- Multiple grouping keys (sort ordering spans both keys) +-- ============================================================ + +query +SELECT k1, k2, sort_array(collect_set(a)) +FROM sa_multikey GROUP BY k1, k2 ORDER BY k1, k2 + +-- ============================================================ +-- Grouping key derived from an expression +-- ============================================================ + +query +SELECT upper(g) AS gg, sort_array(collect_set(i)) +FROM sa_int GROUP BY upper(g) ORDER BY gg + +-- ============================================================ +-- All-NULL group returns an empty array +-- ============================================================ + +query +SELECT g, sort_array(collect_set(v)) FROM sa_nulls GROUP BY g ORDER BY g + +-- ============================================================ +-- Empty table +-- ============================================================ + +query +SELECT sort_array(collect_set(v)) FROM sa_empty + +query +SELECT g, sort_array(collect_set(v)) FROM sa_empty GROUP BY g ORDER BY g + +-- ============================================================ +-- Single row +-- ============================================================ + +query +SELECT sort_array(collect_set(v)) FROM sa_single + +-- ============================================================ +-- collect_set mixed with hashable aggregates: collect_set forces the +-- whole aggregate onto SortAggregate, so count/sum/min/max ride along +-- ============================================================ + +query +SELECT g, sort_array(collect_set(i)), count(*), count(i), sum(i), min(i), max(i) +FROM sa_int GROUP BY g ORDER BY g + +-- ============================================================ +-- Multiple collect_set in one query +-- ============================================================ + +query +SELECT k1, sort_array(collect_set(a)), sort_array(collect_set(k2)) +FROM sa_multikey GROUP BY k1 ORDER BY k1 + +-- ============================================================ +-- DISTINCT (distinct-aggregate planner path) +-- ============================================================ + +query +SELECT g, sort_array(collect_set(DISTINCT i)) FROM sa_int GROUP BY g ORDER BY g + +-- ============================================================ +-- HAVING over a collect_set-derived predicate +-- ============================================================ + +query +SELECT g, sort_array(collect_set(i)) +FROM sa_int GROUP BY g HAVING size(collect_set(i)) > 1 ORDER BY g + +-- ============================================================ +-- Representative data types through the SortAggregate path +-- ============================================================ + +statement +CREATE TABLE sa_string(v string, g string) USING parquet + +statement +INSERT INTO sa_string VALUES + ('hello', 'a'), ('world', 'a'), ('hello', 'a'), (NULL, 'a'), + ('', 'b'), ('x', 'b'), ('', 'b'), (NULL, 'b') + +query +SELECT g, sort_array(collect_set(v)) FROM sa_string GROUP BY g ORDER BY g + +statement +CREATE TABLE sa_bigint(v bigint, g string) USING parquet + +statement +INSERT INTO sa_bigint VALUES + (1000000000000, 'a'), (2000000000000, 'a'), (1000000000000, 'a'), (NULL, 'a'), + (3000000000000, 'b'), (NULL, 'b') + +query +SELECT g, sort_array(collect_set(v)) FROM sa_bigint GROUP BY g ORDER BY g + +statement +CREATE TABLE sa_decimal(v decimal(10,2), g string) USING parquet + +statement +INSERT INTO sa_decimal VALUES + (1.50, 'a'), (2.50, 'a'), (1.50, 'a'), (NULL, 'a'), + (0.00, 'b'), (99999999.99, 'b'), (NULL, 'b') + +query +SELECT g, sort_array(collect_set(v)) FROM sa_decimal GROUP BY g ORDER BY g + +statement +CREATE TABLE sa_date(v date, g string) USING parquet + +statement +INSERT INTO sa_date VALUES + (DATE '2024-01-01', 'a'), (DATE '2024-06-15', 'a'), (DATE '2024-01-01', 'a'), (NULL, 'a'), + (DATE '1970-01-01', 'b'), (NULL, 'b') + +query +SELECT g, sort_array(collect_set(v)) FROM sa_date GROUP BY g ORDER BY g + +-- ============================================================ +-- Unsupported aggregate input under the SortAggregate path: +-- with strictFloatingPoint=true, collect_set on floating-point types is +-- incompatible, so the SortAggregateExec must fall back to Spark cleanly +-- rather than crash. +-- ============================================================ + +statement +CREATE TABLE sa_double(v double, g string) USING parquet + +statement +INSERT INTO sa_double VALUES + (1.1, 'a'), (2.2, 'a'), (1.1, 'a'), (NULL, 'a'), + (CAST('NaN' AS DOUBLE), 'b'), (CAST('NaN' AS DOUBLE), 'b'), (1.0, 'b') + +query expect_fallback(not fully compatible with Spark) +SELECT g, sort_array(collect_set(v)) FROM sa_double GROUP BY g ORDER BY g diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index ae14c68207..db4c93bf7d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.comet.CometHashAggregateExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.functions.{avg, col, count_distinct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -2109,6 +2110,46 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + // useObjectHashAggregateExec=false forces Spark to plan SortAggregateExec for + // TypedImperativeAggregate functions like collect_set. Comet converts those just like + // ObjectHashAggregateExec via the shared CometBaseAggregate path. Broader data-type and + // edge-case coverage lives in the SQL file test + // spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql; these Scala + // tests additionally assert that Spark actually planned a SortAggregateExec, which the SQL + // framework cannot check. + private def assertSortAggregateRunsNatively(query: String): Unit = { + withSQLConf( + "spark.sql.execution.useObjectHashAggregateExec" -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withTempView("tbl") { + Seq((1, "a"), (2, "a"), (1, "a"), (3, "b"), (4, "b"), (4, "b")) + .toDF("v", "g") + .createOrReplaceTempView("tbl") + // Spark must actually plan a SortAggregateExec for this query; otherwise the test + // would pass without exercising the new code path. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val plan = stripAQEPlan(sql(query).queryExecution.executedPlan) + assert( + plan.find(_.isInstanceOf[SortAggregateExec]).isDefined, + s"Expected SortAggregateExec in Spark-only plan but got:\n$plan") + } + checkSparkAnswerAndOperator(sql(query)) + } + } + } + + test("SortAggregate with collect_set is converted to native") { + assertSortAggregateRunsNatively( + "SELECT g, sort_array(collect_set(v)) FROM tbl GROUP BY g ORDER BY g") + } + + test("SortAggregate global collect_set (no grouping keys) is converted to native") { + // Empty grouping is a distinct plan shape: no pre-aggregate sort, empty output ordering, + // and adjustOutputForNativeState with zero grouping columns. + assertSortAggregateRunsNatively("SELECT sort_array(collect_set(v)) FROM tbl") + } + // Regression: Catalyst prunes `HashAggregateExec.resultExpressions` to // empty for EXISTS / row-existence-only subqueries. The native HashAggregate's natural // output (the grouping keys) then disagrees with the pruned JVM `output`, leaking through