You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/05 22:47:17 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-4565-fa25924d719d8d748d06ac1d89fe7fe68235ff9e created (now d16f7bbb7e)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-4565-fa25924d719d8d748d06ac1d89fe7fe68235ff9e
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at d16f7bbb7e feat: support SortAggregateExec (#4565)
This branch includes the following new commits:
new d16f7bbb7e feat: support SortAggregateExec (#4565)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: feat: support SortAggregateExec (#4565)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-4565-fa25924d719d8d748d06ac1d89fe7fe68235ff9e
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit d16f7bbb7ef5aa566626bc98b6655785724b4da4
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Oct 5 18:51:41 2026 +0000
feat: support SortAggregateExec (#4565)
* feat: support SortAggregateExec [skip ci]
Wire SortAggregateExec through the existing CometBaseAggregate.doConvert
path so that queries planned with useObjectHashAggregate=false (or with
TypedImperativeAggregate functions whose buffer formats prevent hash
aggregation) can run their Partial->Final pair natively instead of falling
back to Spark. CollectSet was the motivating function; the same wiring
covers any other natively-implemented TypedImperativeAggregate.
No proto changes: DataFusion's AggregateExec::try_new auto-detects
InputOrderMode::Sorted from the child's output ordering, so a sorted
input naturally produces sorted output.
Wrapper reuse: createExec produces a CometHashAggregateExec (matching
CometObjectHashAggregateExec). CometExec.outputOrdering already defaults
to originalPlan.outputOrdering, so SortAggregateExec.outputOrdering
flows through unchanged for downstream operators that elided a sort
against it.
Cleanups landed in passing:
- Lifted baseAggregateSupportLevel and adjustOutputForNativeState onto
the CometBaseAggregate trait (was duplicated 3x and 2x).
- Collapsed isAggregate / findCometPartialAgg's Spark-side branches /
canAggregateBeConverted's shuffle guard to BaseAggregateExec.
* test: add comprehensive SortAggregate coverage via SQL file tests
Add a SQL file test (sort_aggregate.sql) that forces collect_set through
SortAggregateExec by disabling ObjectHashAggregate, covering global and
grouped aggregation, multiple grouping keys, expression keys, NULL/empty/
single-row groups, mixed aggregates, multiple collect_set, DISTINCT,
HAVING, a representative data-type spread, and a clean fallback when the
aggregate input is incompatible.
Add a global (no grouping keys) plan-shape assertion to CometAggregateSuite
alongside the existing grouped one, since the SQL framework cannot assert
that a SortAggregateExec was actually planned.
* fix: distinguish CometSortAggregateExec and update tests for SortAggregate support
Introduce a dedicated CometSortAggregateExec wrapper instead of reusing
CometHashAggregateExec for converted SortAggregateExec nodes, so the executed
plan reflects whether Spark planned a hash or a sort aggregate. A shared
CometBaseAggregateExec base carries the common rendering and serialization, and
findCometPartialAgg matches the base type so a partial sort aggregate (e.g.
collect_set) is still paired with its final aggregate instead of falling back.
Update the Spark test diffs to recognize the new wrapper: ReplaceHashWithSortAgg
counts CometSortAggregateExec as a sort aggregate, and the generic aggregate
collectors in AdaptiveQueryExecSuite and StreamingAggregationDistributionSuite
match CometBaseAggregateExec so they see through both wrappers.
Refresh the affected SQL file tests now that SortAggregate runs natively:
- min_max.sql: string min/max still falls back, but the reason is now the
StringType limitation rather than SortAggregate being unsupported.
- first_last.sql: the multi-type first/last IGNORE NULLS queries now run
natively; shape test_types so each group has a single non-null value, since
first/last are non-deterministic across engines with multiple non-null rows.
* fix: widen post-merge aggregate guards to CometBaseAggregateExec
Self-review after merging main. Three safety predicates landed on main after
this branch was cut and matched only CometHashAggregateExec, so a
CometSortAggregateExec Partial slipped past all of them:
- CometExecRule.preserveSparkAggregateBuffers: a Celeborn exchange that falls
back to Spark must restore the Comet Partial below it. Missing the sort
variant let a native collect_list buffer reach a Spark Final. Added a
regression test in CometCelebornShufflePlanningSuite that fails without this
change.
- RevertNativeForTransitionHeavyStages.hasUnsafeMixedAggregateAtStageBoundary:
same class of guard at a stage boundary.
- CometMetricNode.withoutAggregateMetrics: range-partition sampling would
double-count sort-aggregate spill and memory metrics.
Also drop the duplicate adjustOutputForNativeState that the merge left behind
(main grew its own copy handling CollectList and PartialMerge; the branch's
older copy handled neither), move main's aggregate metrics override onto the
shared CometBaseAggregateExec base, rename COMET_EXEC_SHUFFLE_ENABLED to
COMET_SHUFFLE_ENABLED, and mark SortAggregateExec supported in operators.md.
* refactor: share the aggregate serde and plan node equality
CometBaseAggregate is now parameterized on the Spark aggregate type and
supplies enabledConfig, getSupportLevel and convert for all three aggregate
serdes. Each serde object keeps only its operator-specific support check, a
requiresCometShuffle flag and createExec. The shared getSupportLevel notes
that the partial and final test knobs gate sort aggregates too.
equals and hashCode move into CometBaseAggregateExec and require the same
class, so the two wrappers keep only withNewChildInternal.
* fix: walk through sorts when pairing sort aggregate buffers
Spark puts a SortExec between each SortAggregateExec and the exchange
below it. The passes that pair a buffer-consuming aggregate with the
Partial that produces its buffers stopped at that sort, so a native
collect_list Partial could stay native under a Spark Final or merge
stage. Spark then read Comet's array buffer where it expected binary, and
a distinct collect_list chain slipped past the issue #4724 guard and
failed natively.
findPartialAggInPlan, revertUnsafePartialAggregates and
hasUnrepairedNativeBuffer now walk through SortExec and CometSortExec.
* fix: keep sort aggregate output in grouping-key order natively
CometSortAggregateExec reports SortAggregateExec's grouping-key output
ordering, and Spark may remove sorts above it on that basis. DataFusion
emits groups in input order only when it sees the input sorted on them.
When the ordering comes from outside the native plan, as with a cached
sorted relation read through an unordered scan, its vectorized grouping
can emit groups out of order: NULL and an empty array hash alike, so an
ORDER BY over array keys returned [1] before [].
A new HashAggregate field, ordered_by_grouping_keys, is set for sort
aggregates. The native planner then sorts the aggregate output on the
grouping columns unless the aggregate's output ordering already satisfies
them. It also normalizes float grouping keys the way native sorts do, so
that DataFusion sees the sort below a float-keyed final aggregate and
streams it instead of falling back to the hash path.
* test: run the collect_set fixtures through SortAggregateExec
collect_set.sql and collect_set_floating_fallback.sql gain a ConfigMatrix
over useObjectHashAggregateExec, so their whole fixtures run through both
ObjectHashAggregateExec and SortAggregateExec. sort_aggregate.sql keeps only
the shapes those files do not cover (several grouping keys, an expression
key, hashable aggregates riding along) and adds the decimal SUM fallback.
Its floating-point fallback query, which expected a fallback that Spark 4.2
no longer has, moves to the matrix.
A new fixture checks that SortAggregateExec stays in Spark, with the
shuffle reason, when Comet shuffle is disabled. min_by.sql and max_by.sql
drop the claim that Comet never converts SortAggregate and check the type
fallback without the TypedImperativeAggregate too. The two native
collect_set Scala tests become one test over both queries.
* docs: note how first and last can differ under sort aggregation
Spark plans a sort aggregate for first and last over a column whose buffer
cannot use hash aggregation, such as a string, and Comet now runs it
natively. Within a group the row order is not defined, and Spark and Comet
can order rows with equal grouping keys differently, so a group with more
than one candidate value can return a different, equally valid, value.
* test: run the mixed-engine collect_list tests with sort aggregates
Each mixed-engine collect_list test now also runs with
useObjectHashAggregateExec=false, so both halves of the split are covered
for SortAggregateExec, where a sort sits between each half and its
exchange. The tests count every Comet aggregate wrapper rather than only
CometHashAggregateExec, and check which aggregate operator Spark planned.
* fix: keep decimal AVG with a maximum-precision sum out of sort aggregates
Spark's Average keeps its decimal sum at precision p + 10, capped at 38.
When a sort aggregate's buffer has a field that is not mutable, such as a
string FIRST, Spark holds that sum in a generic row where it is unbounded,
so 0.6 + 0.6 - 0.4 over DECIMAL(38,38) recovers to 0.8. The native AVG
records the overflow at 1.2 and returns NULL, or raises under ANSI. These
queries stayed in Spark before sort aggregates converted.
CometSortAggregateExec now declines a decimal AVG or TRY_AVG whose sum is
at maximum precision, next to the existing decimal SUM check. The
compatibility guide notes both fallbacks.
* test: expect the distinct collect_list sort aggregate chain to run natively
#4727 added native support for collect_list and collect_set alongside a
distinct aggregate, so the PartialMerge stages of the distinct rewrite no
longer force the chain back to Spark (#4724). Under sort aggregates the
whole chain now runs natively and matches Spark. Check that, and keep the
fallback case by checking that the chain stays in Spark, through the sort
below each stage, when its final cannot convert.
---
dev/diffs/3.4.3.diff | 25 +-
dev/diffs/3.5.9.diff | 25 +-
dev/diffs/4.0.4.diff | 25 +-
dev/diffs/4.1.3.diff | 25 +-
dev/diffs/4.2.0.diff | 25 +-
.../user-guide/latest/compatibility/operators.md | 17 ++
docs/source/user-guide/latest/operators.md | 2 +-
native/core/src/execution/planner.rs | 208 ++++++++++++++-
native/proto/src/proto/operator.proto | 4 +
.../org/apache/comet/rules/CometExecRule.scala | 56 ++--
.../RevertNativeForTransitionHeavyStages.scala | 6 +-
.../apache/spark/sql/comet/CometMetricNode.scala | 2 +-
.../org/apache/spark/sql/comet/operators.scala | 284 ++++++++++++++-------
.../expressions/aggregate/collect_set.sql | 3 +
.../aggregate/collect_set_floating_fallback.sql | 3 +
.../sql-tests/expressions/aggregate/first_last.sql | 15 +-
.../sql-tests/expressions/aggregate/max_by.sql | 22 +-
.../sql-tests/expressions/aggregate/min_by.sql | 16 +-
.../sql-tests/expressions/aggregate/min_max.sql | 5 +-
.../expressions/aggregate/sort_aggregate.sql | 79 ++++++
...ql => sort_aggregate_without_comet_shuffle.sql} | 16 +-
.../apache/comet/exec/CometAggregateSuite.scala | 174 +++++++++++--
.../comet/exec/CometInMemoryCacheSuite.scala | 31 ++-
.../apache/comet/rules/CometExecRuleSuite.scala | 57 ++++-
.../CometCelebornShufflePlanningSuite.scala | 33 ++-
25 files changed, 941 insertions(+), 217 deletions(-)
diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff
index cb0c8aaffa..5c23919a2e 100644
--- a/dev/diffs/3.4.3.diff
+++ b/dev/diffs/3.4.3.diff
@@ -1647,26 +1647,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
@@ -1737,7 +1742,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..47285d87416 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._
@@ -1797,7 +1802,7 @@ index 593bd7bb4ba..b327d84d5cc 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]
}
}
@@ -3104,14 +3109,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
@@ -3119,7 +3124,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 {
@@ -3127,7 +3132,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.9.diff b/dev/diffs/3.5.9.diff
index 53bbd92c87..d83fe66edc 100644
--- a/dev/diffs/3.5.9.diff
+++ b/dev/diffs/3.5.9.diff
@@ -1645,26 +1645,31 @@ index 46ed8fdfd21..4585dbab5b8 100644
private def checkWindowGroupLimits(query: String, count: 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
@@ -1735,7 +1740,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..f7923d2b9a7 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._
@@ -1795,7 +1800,7 @@ index 2f8e401e743..7849c685b19 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]
}
}
@@ -3116,14 +3121,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
@@ -3131,7 +3136,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 {
@@ -3139,7 +3144,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.4.diff b/dev/diffs/4.0.4.diff
index e5c5793fd1..f12fdc9e35 100644
--- a/dev/diffs/4.0.4.diff
+++ b/dev/diffs/4.0.4.diff
@@ -2070,26 +2070,31 @@ index 46ed8fdfd21..4585dbab5b8 100644
private def checkWindowGroupLimits(query: String, count: 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
@@ -2382,7 +2387,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 fb8fab6a80f..403eb411920 100644
+index fb8fab6a80f..c5768935805 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
@@ -2454,7 +2459,7 @@ index fb8fab6a80f..403eb411920 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]
}
}
@@ -3872,14 +3877,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
@@ -3887,7 +3892,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 {
@@ -3895,7 +3900,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.3.diff b/dev/diffs/4.1.3.diff
index 57a6c69607..c99c1c242d 100644
--- a/dev/diffs/4.1.3.diff
+++ b/dev/diffs/4.1.3.diff
@@ -2204,26 +2204,31 @@ index 46ed8fdfd21..4585dbab5b8 100644
private def checkWindowGroupLimits(query: String, count: 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
@@ -2516,7 +2521,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 188a28ff1c0..8fdccf31749 100644
+index 188a28ff1c0..2717d45f3d0 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
@@ -2588,7 +2593,7 @@ index 188a28ff1c0..8fdccf31749 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]
}
}
@@ -4158,14 +4163,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
@@ -4173,7 +4178,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 {
@@ -4181,7 +4186,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/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff
index 3fed2699cf..01a29c48ae 100644
--- a/dev/diffs/4.2.0.diff
+++ b/dev/diffs/4.2.0.diff
@@ -2271,26 +2271,31 @@ index 03e2f4942ed..b0a80fcf683 100644
private def checkWindowGroupLimits(query: String, count: 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 41edb53441f..eed2cbf858c 100644
+index 41edb53441f..fabfa9c883e 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
-+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
-@@ -30,7 +31,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
+@@ -30,9 +31,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 e8a39eb8c53..a621c7a7923 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -2607,7 +2612,7 @@ index d70bd715879..074a9fa29d9 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 d6d19d21e65..751ad50a569 100644
+index d6d19d21e65..85e2f441f81 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,13 +27,15 @@ import org.apache.spark.SparkException
@@ -2680,7 +2685,7 @@ index d6d19d21e65..751ad50a569 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]
}
}
@@ -4363,14 +4368,14 @@ index a6067aaf189..76e73f1c554 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
@@ -4378,7 +4383,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 {
@@ -4386,7 +4391,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/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md
index 467e48ecb5..f3fc905399 100644
--- a/docs/source/user-guide/latest/compatibility/operators.md
+++ b/docs/source/user-guide/latest/compatibility/operators.md
@@ -60,6 +60,23 @@ Sampling with replacement (`df.sample(withReplacement = true, ...)`) falls back
it draws from a Poisson distribution that Comet does not implement natively
([#5109](https://github.com/apache/datafusion-comet/issues/5109)).
+## Sort Aggregation
+
+Comet runs `SortAggregateExec` natively when Comet shuffle is enabled and Comet supports every
+aggregate in it. The native aggregate keeps the grouping-key output order that Spark relies on.
+
+Decimal `sum` and `avg` over input precision 28 or more keep a running sum at precision 38, which
+can overflow even when the final sum fits. Whether Spark's sort aggregation recovers from such an
+overflow depends on the other aggregates in the operator and on codegen. Comet does not track
+this, so a sort aggregate that contains such a `sum` or `avg` falls back to Spark.
+
+`first` and `last` return the first or last value in the order that rows reach the aggregate,
+which Spark does not define within a group. Spark plans a sort aggregate for them when their buffer
+cannot use hash aggregation, for example over a string column. The sort below the aggregate orders
+rows by the grouping keys only, and Spark and Comet can leave rows with equal keys in different
+orders, so a group with more than one candidate value can return a different value than Spark.
+Both results are valid under Spark's semantics for these functions.
+
## Window Functions
Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of
diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md
index 5013e1028b..0155ebb8da 100644
--- a/docs/source/user-guide/latest/operators.md
+++ b/docs/source/user-guide/latest/operators.md
@@ -92,7 +92,7 @@ omitted from the tables below and may be reconsidered based on demand:
| ------------------------- | ------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `HashAggregateExec` | ✅ | |
| `ObjectHashAggregateExec` | ✅ | Runs the object-buffer aggregates Comet supports, such as `collect_list`, `collect_set`, `percentile`, `approx_percentile`, `mode`, `bloom_filter_agg`, and (Spark 4.0+) `listagg`. Falls back when Comet shuffle is disabled, which would otherwise split the aggregate across Comet and Spark. See [Shuffle](tuning/shuffle.md). |
-| `SortAggregateExec` | 🔜 | Falls back today; Comet currently accelerates hash aggregates. |
+| `SortAggregateExec` | ✅ | Runs when Comet supports every aggregate in it, including the object-buffer aggregates listed for `ObjectHashAggregateExec`. Falls back when Comet shuffle is disabled, which would otherwise split the aggregate across Comet and Spark. See [Shuffle](tuning/shuffle.md). |
## Joins
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 98f9c2a7cb..38d2e3084f 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -79,7 +79,7 @@ use datafusion::{
limit::LocalLimitExec,
projection::ProjectionExec,
sorts::sort::SortExec,
- ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+ ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions,
},
prelude::SessionContext,
};
@@ -1426,8 +1426,16 @@ impl PhysicalPlanner {
.iter()
.enumerate()
.map(|(idx, expr)| {
- self.create_expr(expr, child.schema())
- .map(|r| (r, format!("col_{idx}")))
+ // A native sort normalizes float keys, so normalize float grouping keys
+ // the same way for DataFusion to see that the sort below a sort aggregate
+ // orders them. Spark already normalized these values, so this changes
+ // none of them.
+ if agg.ordered_by_grouping_keys {
+ self.create_normalized_key_expr(expr, child.schema())
+ } else {
+ self.create_expr(expr, child.schema())
+ }
+ .map(|r| (r, format!("col_{idx}")))
})
.collect();
let group_by = PhysicalGroupBy::new_single(group_exprs?);
@@ -1547,6 +1555,46 @@ impl PhysicalPlanner {
)?,
);
+ // Spark's SortAggregateExec reports its output as ordered by the grouping keys,
+ // and Spark may have removed a sort above it on that basis. DataFusion emits
+ // groups in input order only when it sees the input sorted on them. When that
+ // ordering comes from outside this native plan, for example from a cached sorted
+ // relation that reaches native code as an unordered ScanExec, the hash table can
+ // emit groups out of order, so sort the aggregate output instead.
+ if agg.ordered_by_grouping_keys && !agg.grouping_exprs.is_empty() {
+ let aggregate_schema = aggregate.schema();
+ let ordering: Vec<PhysicalSortExpr> = (0..agg.grouping_exprs.len())
+ .map(|idx| {
+ PhysicalSortExpr::new(
+ Arc::new(Column::new(aggregate_schema.field(idx).name(), idx)),
+ SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ )
+ })
+ .collect();
+ if !aggregate
+ .equivalence_properties()
+ .ordering_satisfy(ordering.clone())?
+ {
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+ LexOrdering::new(ordering).unwrap(),
+ Arc::clone(&aggregate),
+ ));
+ return Ok((
+ scans,
+ shuffle_scans,
+ Arc::new(SparkPlan::new_with_additional(
+ spark_plan.plan_id,
+ sort,
+ vec![child],
+ vec![aggregate],
+ )),
+ ));
+ }
+ }
+
Ok((
scans,
shuffle_scans,
@@ -5830,6 +5878,7 @@ mod tests {
mode: spark_operator::AggregateMode::Partial as i32,
expr_modes: vec![],
initial_input_buffer_offset: 0,
+ ordered_by_grouping_keys: false,
})),
};
let projection = Operator {
@@ -5855,6 +5904,159 @@ mod tests {
assert_eq!("ScanExec", projection_exec.children[0].native_plan.name());
}
+ #[tokio::test]
+ async fn sort_aggregate_output_is_ordered_by_grouping_keys() {
+ use arrow::datatypes::Int32Type;
+ use datafusion::common::tree_node::{Transformed, TreeNode};
+ use spark_expression::data_type::{data_type_info::DatatypeStruct, DataTypeInfo, ListInfo};
+
+ let array_type = spark_expression::DataType {
+ type_id: 14,
+ type_info: Some(Box::new(DataTypeInfo {
+ datatype_struct: Some(DatatypeStruct::List(Box::new(ListInfo {
+ element_type: Some(Box::new(create_proto_datatype())),
+ contains_null: true,
+ element_field_id: None,
+ }))),
+ })),
+ };
+ let key = Expr {
+ expr_struct: Some(Bound(spark_expression::BoundReference {
+ index: 0,
+ datatype: Some(array_type.clone()),
+ })),
+ ..Default::default()
+ };
+ let scan = Operator {
+ op_struct: Some(OpStruct::Scan(spark_operator::Scan {
+ fields: vec![array_type],
+ source: String::new(),
+ })),
+ ..Default::default()
+ };
+ let sort = Operator {
+ children: vec![scan.clone()],
+ op_struct: Some(OpStruct::Sort(spark_operator::Sort {
+ sort_orders: vec![Expr {
+ expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder {
+ child: Some(Box::new(key.clone())),
+ direction: spark_expression::SortDirection::Ascending as i32,
+ null_ordering: spark_expression::NullOrdering::NullsFirst as i32,
+ }))),
+ ..Default::default()
+ }],
+ ..Default::default()
+ })),
+ ..Default::default()
+ };
+ let aggregate = |child: Operator| Operator {
+ plan_id: 1,
+ children: vec![child],
+ op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate {
+ grouping_exprs: vec![key.clone()],
+ mode: spark_operator::AggregateMode::Final as i32,
+ ordered_by_grouping_keys: true,
+ ..Default::default()
+ })),
+ ..Default::default()
+ };
+ let ctx = SessionContext::new();
+ let task_ctx = ctx.task_ctx();
+ let planner = PhysicalPlanner::new(Arc::new(ctx), 0);
+
+ // DataFusion sees a sort in the same native plan and emits the groups in input order.
+ let (_, _, planned) = planner
+ .create_plan(&aggregate(sort), &mut vec![], 1)
+ .unwrap();
+ assert_eq!("AggregateExec", planned.native_plan.name());
+
+ // An input ordered outside the native plan, as a cached sorted relation is, is not visible
+ // to DataFusion. NULL and an empty list hash alike, so its vectorized grouping would emit
+ // the empty list after [1] if the planner did not sort the output.
+ let (_, _, planned) = planner
+ .create_plan(&aggregate(scan), &mut vec![], 1)
+ .unwrap();
+ assert_eq!("SortExec", planned.native_plan.name());
+ assert_eq!(1, planned.additional_native_plans.len());
+ assert_eq!("AggregateExec", planned.additional_native_plans[0].name());
+
+ let keys = Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ None,
+ Some(vec![]),
+ Some(vec![Some(1)]),
+ ])) as ArrayRef;
+ let batch =
+ RecordBatch::try_new(planned.children[0].schema(), vec![Arc::clone(&keys)]).unwrap();
+ let input: Arc<dyn ExecutionPlan> =
+ MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap();
+ let native_plan = Arc::clone(&planned.native_plan)
+ .transform_up(|node| {
+ Ok(if node.name() == "ScanExec" {
+ Transformed::yes(Arc::clone(&input))
+ } else {
+ Transformed::no(node)
+ })
+ })
+ .unwrap()
+ .data;
+ let results = collect(native_plan.execute(0, task_ctx).unwrap())
+ .await
+ .unwrap();
+ let output = arrow::compute::concat_batches(&results[0].schema(), &results).unwrap();
+ assert_eq!(keys.as_ref(), output.column(0).as_ref());
+ }
+
+ #[test]
+ fn sort_aggregate_float_key_streams_over_native_sort() {
+ // A native sort normalizes a float key, so the aggregate above it must group on the
+ // normalized key too for DataFusion to see the ordering and skip the output sort.
+ let double_type = spark_expression::DataType {
+ type_id: 6,
+ type_info: None,
+ };
+ let key = Expr {
+ expr_struct: Some(Bound(spark_expression::BoundReference {
+ index: 0,
+ datatype: Some(double_type.clone()),
+ })),
+ ..Default::default()
+ };
+ let op = Operator {
+ plan_id: 1,
+ children: vec![Operator {
+ children: vec![Operator {
+ op_struct: Some(OpStruct::Scan(spark_operator::Scan {
+ fields: vec![double_type],
+ source: String::new(),
+ })),
+ ..Default::default()
+ }],
+ op_struct: Some(OpStruct::Sort(spark_operator::Sort {
+ sort_orders: vec![Expr {
+ expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder {
+ child: Some(Box::new(key.clone())),
+ direction: spark_expression::SortDirection::Ascending as i32,
+ null_ordering: spark_expression::NullOrdering::NullsFirst as i32,
+ }))),
+ ..Default::default()
+ }],
+ ..Default::default()
+ })),
+ ..Default::default()
+ }],
+ op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate {
+ grouping_exprs: vec![key],
+ mode: spark_operator::AggregateMode::Final as i32,
+ ordered_by_grouping_keys: true,
+ ..Default::default()
+ })),
+ ..Default::default()
+ };
+ let planner = PhysicalPlanner::default();
+ let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap();
+ assert_eq!("AggregateExec", planned.native_plan.name());
+ }
+
#[tokio::test]
async fn projection_prunes_filter_output_and_preserves_metrics() {
use crate::execution::metrics::utils::to_native_metric_node;
diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
index 68b881cdcb..527461f47a 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -828,6 +828,10 @@ message HashAggregate {
// Offset in the child's output where aggregate buffer attributes start.
// Used by PartialMerge to locate state fields in the input.
int32 initial_input_buffer_offset = 7;
+ // Set for Spark's SortAggregateExec, which reports its output as ordered by the grouping
+ // keys (ascending, nulls first). The native planner sorts the aggregate output on the
+ // grouping columns unless the aggregate already emits its groups in that order.
+ bool ordered_by_grouping_keys = 9;
}
message Limit {
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 d26ca2a624..736143baf6 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -38,7 +38,7 @@ import org.apache.spark.sql.comet.shims.{ShimCometEmptyRelation, ShimCometOneRow
import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, 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.columnar.InMemoryTableScanExec
import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec}
@@ -113,6 +113,7 @@ object CometExecRule {
classOf[GenerateExec] -> CometExplodeExec,
classOf[HashAggregateExec] -> CometHashAggregateExec,
classOf[ObjectHashAggregateExec] -> CometObjectHashAggregateExec,
+ classOf[SortAggregateExec] -> CometSortAggregateExec,
classOf[BroadcastHashJoinExec] -> CometBroadcastHashJoinExec,
classOf[BroadcastNestedLoopJoinExec] -> CometBroadcastNestedLoopJoinExec,
classOf[ShuffledHashJoinExec] -> CometHashJoinExec,
@@ -227,8 +228,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
* 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 =>
@@ -280,7 +280,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
* Restore a Spark Partial while retaining its current children. The tag prevents reconversion
* when AQE replans the exchange without its Final, and records why the Partial stays in Spark.
*/
- private def restoreSparkPartial(agg: CometHashAggregateExec, reason: String): SparkPlan = {
+ private def restoreSparkPartial(agg: CometBaseAggregateExec, reason: String): SparkPlan = {
val partial = agg.originalPlan.withNewChildren(agg.children)
partial.setTagValue(CometExecRule.COMET_UNSAFE_PARTIAL, reason)
withFallbackReason(partial, reason)
@@ -301,12 +301,12 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
def restore(plan: SparkPlan): SparkPlan = plan match {
// Do not rewrite data that an earlier stage may already have materialized.
case _: QueryStageExec | _: ShuffleExchangeLike | _: BroadcastExchangeLike => plan
- case agg: CometHashAggregateExec
+ case agg: CometBaseAggregateExec
if agg.modes == Seq(Partial) &&
!QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) =>
restoreSparkPartial(agg, reason)
// Final output is ordinary SQL data; any partial below it belongs to another aggregate.
- case agg: CometHashAggregateExec if agg.modes.contains(Final) => agg
+ case agg: CometBaseAggregateExec if agg.modes.contains(Final) => agg
case agg: BaseAggregateExec if agg.aggregateExpressions.exists(_.mode == Final) => agg
case placeholder: CometSinkPlaceHolder =>
val child = restore(placeholder.child)
@@ -1538,13 +1538,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
* Inspect a failed repair's buffer path without rewriting it or materializing any stage. Report
* only a native Partial/PartialMerge whose emitted state is not known to be Spark-compatible.
* Spark Partials and completed aggregates establish new buffers, so stop there rather than
- * finding an unrelated native producer below them. Only known aggregate and exchange wrappers
- * forward the same buffer path; an arbitrary operator is not evidence of a mixed boundary.
+ * finding an unrelated native producer below them. Only known aggregate, sort and exchange
+ * wrappers forward the same buffer path; an arbitrary operator is not evidence of a mixed
+ * boundary.
*/
private def hasUnrepairedNativeBuffer(plan: SparkPlan): Boolean = plan match {
- case agg: CometHashAggregateExec if agg.aggregateExpressions.isEmpty =>
+ case agg: CometBaseAggregateExec if agg.aggregateExpressions.isEmpty =>
hasUnrepairedNativeBuffer(agg.child)
- case agg: CometHashAggregateExec =>
+ case agg: CometBaseAggregateExec =>
agg.modes.forall(m => m == Partial || m == PartialMerge) &&
!QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions)
case agg: BaseAggregateExec
@@ -1555,6 +1556,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
agg.aggregateExpressions.forall(e => e.mode == Partial || e.mode == PartialMerge) &&
hasUnrepairedNativeBuffer(agg.child)
case placeholder: CometSinkPlaceHolder => hasUnrepairedNativeBuffer(placeholder.child)
+ // A sort aggregate reads its buffers through the sort that orders them.
+ case sort: SortExec => hasUnrepairedNativeBuffer(sort.child)
+ case sort: CometSortExec => hasUnrepairedNativeBuffer(sort.child)
case read: AQEShuffleReadExec => hasUnrepairedNativeBuffer(read.child)
case stage: ShuffleQueryStageExec => hasUnrepairedNativeBuffer(stage.plan)
case reused: ReusedExchangeExec => hasUnrepairedNativeBuffer(reused.child)
@@ -1566,22 +1570,22 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
/**
* The early tagging pass cannot know whether a Final's child will become native. Check the
* actual conversion result before serialization or AQE stage creation, restoring the feeding
- * aggregate/exchange chain while keeping native work below its Partial. When the Partial is
- * already in Spark, restore the native merge stages above it instead. Return the repaired plan,
- * or preserve an unrepairable path and record one warning on its Spark Final if an unsafe
+ * aggregate/sort/exchange chain while keeping native work below its Partial. When the Partial
+ * is already in Spark, restore the native merge stages above it instead. Return the repaired
+ * plan, or preserve an unrepairable path and record one warning on its Spark Final if an unsafe
* native producer remains. Existing stages and their buffers are never rewritten by this pass.
*/
private[rules] def revertUnsafePartialAggregates(plan: SparkPlan): SparkPlan = {
// `belowNativeMerge` is set once the walk has passed a native aggregate that merges buffers.
def revertChain(node: SparkPlan, belowNativeMerge: Boolean): Option[SparkPlan] = node match {
- case agg: CometHashAggregateExec if agg.modes == Seq(Partial) =>
+ case agg: CometBaseAggregateExec if agg.modes == Seq(Partial) =>
Some(
restoreSparkPartial(
agg,
"Partial aggregate disabled: corresponding final aggregate " +
"cannot be converted to Comet and intermediate buffer formats are incompatible"))
- case agg: CometHashAggregateExec
+ case agg: CometBaseAggregateExec
if agg.modes.forall(m => m == Partial || m == PartialMerge) =>
val merges = agg.aggregateExpressions.nonEmpty
revertChain(agg.child, belowNativeMerge || merges).map { child =>
@@ -1609,6 +1613,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
if agg.aggregateExpressions.forall(e => e.mode == Partial || e.mode == PartialMerge) =>
revertChain(agg.child, belowNativeMerge).map(child => agg.withNewChildren(Seq(child)))
+ // A sort aggregate reads its buffers through the sort that orders them.
+ case sort: CometSortExec =>
+ revertChain(sort.child, belowNativeMerge)
+ .map(child => sort.originalPlan.withNewChildren(Seq(child)))
+ case sort: SortExec =>
+ revertChain(sort.child, belowNativeMerge)
+ .map(child => sort.withNewChildren(Seq(child)))
+
case CometSinkPlaceHolder(_, _, shuffle: CometShuffleExchangeExec) =>
revertChain(shuffle, belowNativeMerge)
case shuffle: CometShuffleExchangeExec =>
@@ -1651,12 +1663,13 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
/**
* Look for the bottom Partial-mode aggregate that feeds into the given plan (the child of a
* Final or PartialMerge), and return it after the aggregates passed on the way down. Walks
- * through exchanges and AQE stages, and continues down through intermediate aggregate stages
- * whose modes are all Partial / PartialMerge - these are the PartialMerge (and mixed
- * Partial/PartialMerge) stages that Spark's distinct-aggregate rewrite inserts between the
- * Partial and the Final. Stops at anything else, returning nothing if no Partial was reached.
- * Requires `aggregateExpressions.nonEmpty` so that group-by-only dedup stages are traversed
- * rather than mistaken for the partial.
+ * through exchanges, AQE stages and the sorts that Spark inserts below sort aggregates, and
+ * continues down through intermediate aggregate stages whose modes are all Partial /
+ * PartialMerge - these are the PartialMerge (and mixed Partial/PartialMerge) stages that
+ * Spark's distinct-aggregate rewrite inserts between the Partial and the Final. Stops at
+ * anything else, returning nothing if no Partial was reached. Requires
+ * `aggregateExpressions.nonEmpty` so that group-by-only dedup stages are traversed rather than
+ * mistaken for the partial.
*/
private def findBufferProducersInPlan(plan: SparkPlan): Seq[BaseAggregateExec] = plan match {
case agg: BaseAggregateExec
@@ -1672,6 +1685,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
case a: AQEShuffleReadExec => findBufferProducersInPlan(a.child)
case s: ShuffleQueryStageExec => findBufferProducersInPlan(s.plan)
case e: ShuffleExchangeExec => findBufferProducersInPlan(e.child)
+ case s: SortExec => findBufferProducersInPlan(s.child)
case other =>
logDebug(
s"findBufferProducersInPlan: stopping at ${other.nodeName}; not a known passthrough")
diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
index 622f24a228..1181a573bc 100644
--- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
@@ -23,7 +23,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometExec, CometHashAggregateExec, CometIcebergNativeScanExec, CometLocalTopKExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometIcebergNativeScanExec, CometLocalTopKExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometSparkToColumnarExec}
import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.QueryStageExec
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike}
@@ -121,13 +121,13 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan
private def hasUnsafeMixedAggregateAtStageBoundary(stagePlan: SparkPlan): Boolean = {
def reachesBoundaryBeforeAggregate(plan: SparkPlan): Boolean = plan match {
case _ if isStageBoundary(plan) => true
- case _: CometHashAggregateExec => false
+ case _: CometBaseAggregateExec => false
case _ => plan.children.exists(reachesBoundaryBeforeAggregate)
}
def visit(plan: SparkPlan): Boolean = plan match {
case _ if isStageBoundary(plan) => false
- case aggregate: CometHashAggregateExec
+ case aggregate: CometBaseAggregateExec
if !QueryPlanSerde
.allAggsSupportNativePartialToSparkFinal(aggregate.aggregateExpressions) ||
QueryPlanSerde
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala
index 17dc34dcb6..7d86de43c4 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala
@@ -96,7 +96,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
*/
private[comet] def withoutAggregateMetrics(plan: SparkPlan): CometMetricNode =
CometMetricNode(
- if (plan.isInstanceOf[CometHashAggregateExec]) {
+ if (plan.isInstanceOf[CometBaseAggregateExec]) {
metrics -- CometMetricNode.aggregateMetricNames
} else {
metrics
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 70f6ce1874..909539bf2b 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
@@ -31,7 +31,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, Expression, ExpressionSet, Generator, LeafExpression, Literal, NamedExpression, SortOrder, XXH64}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum}
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
import org.apache.spark.sql.catalyst.plans._
@@ -41,7 +41,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, QueryStageExec, 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}
@@ -2047,7 +2047,67 @@ case class CometUnionExec(
override def hashCode(): Int = Objects.hashCode(output, children)
}
-trait CometBaseAggregate {
+/**
+ * Shared serde for Spark's aggregate operators. HashAggregateExec, ObjectHashAggregateExec and
+ * SortAggregateExec all serialize to the same native HashAggregate, so each serde object supplies
+ * only its operator-specific support checks and the Comet plan node it creates.
+ */
+trait CometBaseAggregate[T <: BaseAggregateExec] extends CometOperatorSerde[T] {
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+ CometConf.COMET_EXEC_AGGREGATE_ENABLED)
+
+ /**
+ * Whether the operator carries aggregates whose intermediate buffer formats differ between
+ * Spark and Comet, so that its Partial and Final must both run in Comet, with Comet shuffle
+ * between them.
+ */
+ protected def requiresCometShuffle: Boolean
+
+ /** Operator-specific support checks, run after the shared ones in `getSupportLevel`. */
+ protected def operatorSupportLevel(op: T): SupportLevel = Compatible()
+
+ /**
+ * Whether Spark reports the operator's output as ordered by its grouping keys, which the native
+ * aggregate must then preserve.
+ */
+ protected def orderedByGroupingKeys: Boolean = false
+
+ override final def getSupportLevel(op: T): SupportLevel = {
+ // Some unit tests disable partial or final aggregate conversion to check that CometExecRule
+ // does not allow mixed Spark/Comet aggregates. Despite their names, these knobs gate every
+ // aggregate operator, sort aggregates included.
+ 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"))
+ }
+ // When Comet shuffle is disabled we do not want to convert the aggregate, because we would
+ // probably get partial Comet aggregation and final Spark aggregation. The reason text stays
+ // mode-neutral because `getSupportLevel` runs for every stage, so the Final node receives it
+ // too and must not be told it is the partial. Declining here rather than in `convert` is what
+ // lets CometExecRule record the reason centrally, so the fallback is explained rather than
+ // silent - see https://github.com/apache/datafusion-comet/issues/5500. The message
+ // deliberately does not name a config key: `isCometShuffleEnabled` is a conjunction of the
+ // shuffle config, the configured shuffle manager and the Celeborn compatibility check, so
+ // naming one of them would misdirect when another is the cause.
+ if (requiresCometShuffle && !isCometShuffleEnabled(op.conf)) {
+ return Unsupported(
+ Some(
+ s"Comet shuffle is not enabled, so converting ${op.nodeName} would split the " +
+ "aggregate across Comet and Spark"))
+ }
+ operatorSupportLevel(op)
+ }
+
+ override def convert(
+ aggregate: T,
+ builder: Operator.Builder,
+ childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] =
+ doConvert(aggregate, builder, childOp: _*)
/**
* Whether a decimal SUM's result precision is DecimalType.MAX_PRECISION, the only case with no
@@ -2063,6 +2123,21 @@ trait CometBaseAggregate {
case _ => false
})
+ /**
+ * Whether a decimal AVG accumulates its sum at DecimalType.MAX_PRECISION, which Spark does for
+ * input precision 28 and up. As with SUM, that is where an intermediate overflow can change the
+ * answer.
+ */
+ protected def hasMaxPrecisionDecimalAvg(op: BaseAggregateExec): Boolean =
+ op.aggregateExpressions.exists(_.aggregateFunction match {
+ case avg: Average =>
+ avg.sumDataType match {
+ case decimal: DecimalType => decimal.precision == DecimalType.MAX_PRECISION
+ case _ => false
+ }
+ case _ => false
+ })
+
def doConvert(
aggregate: BaseAggregateExec,
builder: Operator.Builder,
@@ -2186,6 +2261,7 @@ trait CometBaseAggregate {
if (aggregateExpressions.isEmpty) {
val hashAggBuilder = OperatorOuterClass.HashAggregate.newBuilder()
hashAggBuilder.addAllGroupingExprs(groupingExprs.map(_.get).asJava)
+ hashAggBuilder.setOrderedByGroupingKeys(orderedByGroupingKeys)
// Spark has no expression mode to serialize here. An empty aggregate with a required child
// distribution must fully deduplicate its keys (Final, or a pre-distinct PartialMerge), so
// use native Final to keep skip-partial disabled.
@@ -2250,6 +2326,7 @@ trait CometBaseAggregate {
hashAggBuilder.addAllGroupingExprs(groupingExprs.map(_.get).asJava)
hashAggBuilder.addAllAggExprs(aggExprs.map(_.get).asJava)
hashAggBuilder.setModeValue(mode.getNumber)
+ hashAggBuilder.setOrderedByGroupingKeys(orderedByGroupingKeys)
// Send per-expression modes and buffer offset for PartialMerge handling
if (hasPartialMerge) {
@@ -2389,21 +2466,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)
@@ -2413,24 +2488,11 @@ trait CometBaseAggregate {
}
-object CometHashAggregateExec
- extends CometOperatorSerde[HashAggregateExec]
- with CometBaseAggregate {
+object CometHashAggregateExec extends CometBaseAggregate[HashAggregateExec] {
- override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
- CometConf.COMET_EXEC_AGGREGATE_ENABLED)
+ override protected def requiresCometShuffle: Boolean = false
- 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"))
- }
+ override protected def operatorSupportLevel(op: HashAggregateExec): SupportLevel = {
// Without codegen Spark buffers an ungrouped aggregate in an UnsafeRow, which latches a
// decimal sum that leaves the precision, while the native accumulator keeps it unbounded.
// Spark turns codegen off by config, for an imperative aggregate, for a non-leaf
@@ -2454,13 +2516,6 @@ object CometHashAggregateExec
Compatible()
}
- override def convert(
- aggregate: HashAggregateExec,
- builder: Operator.Builder,
- childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
- doConvert(aggregate, builder, childOp: _*)
- }
-
override def createExec(nativeOp: Operator, op: HashAggregateExec): CometNativeExec = {
CometHashAggregateExec(
nativeOp,
@@ -2476,40 +2531,11 @@ object CometHashAggregateExec
}
}
-object CometObjectHashAggregateExec
- extends CometOperatorSerde[ObjectHashAggregateExec]
- with CometBaseAggregate {
+object CometObjectHashAggregateExec extends CometBaseAggregate[ObjectHashAggregateExec] {
- override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
- CometConf.COMET_EXEC_AGGREGATE_ENABLED)
+ override protected def requiresCometShuffle: Boolean = true
- 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"))
- }
- // When Comet shuffle is disabled we do not want to transform the ObjectHashAggregate to
- // CometHashAggregate, because we would probably get partial Comet aggregation and final
- // Spark aggregation. The reason text stays mode-neutral because `getSupportLevel` runs for
- // every stage, so the Final node receives it too and must not be told it is the partial.
- // Declining here rather than in `convert` is what lets CometExecRule
- // record the reason centrally, so the fallback is explained rather than silent - see
- // https://github.com/apache/datafusion-comet/issues/5500. The message deliberately does not
- // name a config key: `isCometShuffleEnabled` is a conjunction of the shuffle config, the
- // configured shuffle manager and the Celeborn compatibility check, so naming one of them
- // would misdirect when another is the cause.
- if (!isCometShuffleEnabled(op.conf)) {
- return Unsupported(
- Some(
- "Comet shuffle is not enabled, so converting ObjectHashAggregate would split the " +
- "aggregate across Comet and Spark"))
- }
+ override protected def operatorSupportLevel(op: ObjectHashAggregateExec): SupportLevel = {
// Spark's object aggregation buffers the intermediate decimal sum unbounded, while Comet's
// grouped accumulator latches to null once a running sum leaves the precision, so the
// grouped case declines. Decimal AVG has the same gap and is tracked separately.
@@ -2522,13 +2548,6 @@ object CometObjectHashAggregateExec
Compatible()
}
- override def convert(
- aggregate: ObjectHashAggregateExec,
- builder: Operator.Builder,
- childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
- doConvert(aggregate, builder, childOp: _*)
- }
-
override def createExec(nativeOp: Operator, op: ObjectHashAggregateExec): CometNativeExec = {
CometHashAggregateExec(
nativeOp,
@@ -2544,24 +2563,79 @@ object CometObjectHashAggregateExec
}
}
-case class CometHashAggregateExec(
- override val nativeOp: Operator,
- override val originalPlan: SparkPlan,
- override val output: Seq[Attribute],
- groupingExpressions: Seq[NamedExpression],
- aggregateExpressions: Seq[AggregateExpression],
- aggregateAttributes: Seq[Attribute],
- resultExpressions: Seq[NamedExpression],
- input: Seq[Attribute],
- child: SparkPlan,
- override val serializedPlanOpt: SerializedPlan)
+object CometSortAggregateExec extends CometBaseAggregate[SortAggregateExec] {
+
+ // Spark plans SortAggregate for TypedImperativeAggregate functions whose intermediate buffer
+ // formats differ between Spark and Comet, the same risk as ObjectHashAggregate.
+ override protected def requiresCometShuffle: Boolean = true
+
+ override protected def operatorSupportLevel(op: SortAggregateExec): SupportLevel = {
+ // Spark's sort aggregation buffers in an UnsafeRow, which latches a decimal sum that leaves
+ // the precision, only when every buffer field is mutable and codegen does not apply;
+ // otherwise the sum stays unbounded. Comet latches only when grouped, so decline rather
+ // than track every combination.
+ if (hasMaxPrecisionDecimalSum(op)) {
+ return Unsupported(
+ Some("Decimal SUM at maximum precision cannot match Spark's sort aggregation buffer"))
+ }
+ // Decimal AVG keeps the same kind of running sum, and the native AVG records an overflow as
+ // soon as that sum leaves the precision, grouped or not. Where Spark's sum stays unbounded,
+ // as it does beside a string FIRST, Comet would return NULL, or raise under ANSI, for an
+ // average that Spark computes.
+ if (hasMaxPrecisionDecimalAvg(op)) {
+ return Unsupported(
+ Some(
+ "Decimal AVG with a maximum-precision sum cannot match Spark's " +
+ "sort aggregation buffer"))
+ }
+ Compatible()
+ }
+
+ // CometExec.outputOrdering reports SortAggregateExec's grouping-key ordering, and Spark may have
+ // removed sorts above it on that basis, so the native aggregate must emit its groups in that
+ // order. DataFusion does so on its own when it sees its input sorted on the grouping keys, which
+ // holds when the sort below runs in the same native plan. Otherwise, for example over a cached
+ // sorted relation that reaches native code as an unordered scan, its hash table can emit groups
+ // in any order, and the native planner sorts the aggregate output instead.
+ override protected def orderedByGroupingKeys: Boolean = true
+
+ override def createExec(nativeOp: Operator, op: SortAggregateExec): CometNativeExec = {
+ CometSortAggregateExec(
+ nativeOp,
+ op,
+ adjustOutputForNativeState(op),
+ op.groupingExpressions,
+ op.aggregateExpressions,
+ op.aggregateAttributes,
+ op.resultExpressions,
+ op.child.output,
+ op.child,
+ SerializedPlan(None))
+ }
+}
+
+/**
+ * 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 so the executed plan reflects whether Spark planned a HashAggregateExec /
+ * ObjectHashAggregateExec or a SortAggregateExec. The sort variant also keeps the grouping-key
+ * output ordering that SortAggregateExec reports; see `CometSortAggregateExec`.
+ */
+abstract class CometBaseAggregateExec
extends CometUnaryExec
with PartitioningPreservingUnaryExecNode {
+ def groupingExpressions: Seq[NamedExpression]
+ def aggregateExpressions: Seq[AggregateExpression]
+ def aggregateAttributes: Seq[Attribute]
+ 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: subclass constructor parameters are not assigned yet while this class initializes.
+ lazy val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct
// Match Spark's aggregate canonicalization, including the original result attributes that
// rewritten DISTINCT aggregate expressions do not necessarily retain in their resultIds.
@@ -2573,9 +2647,6 @@ case class CometHashAggregateExec(
override def producedAttributes: AttributeSet =
outputSet ++ AttributeSet(resultExpressions) ++ AttributeSet(aggregateAttributes)
- override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
- this.copy(child = newChild)
-
override def verboseStringWithOperatorId(): String = {
s"""
|$formattedNodeName
@@ -2588,9 +2659,10 @@ case class CometHashAggregateExec(
override def stringArgs: Iterator[Any] =
Iterator(input, modes, groupingExpressions, aggregateExpressions, child)
+ // The wrappers differ only in type, so equality also requires the same class.
override def equals(obj: Any): Boolean = {
obj match {
- case other: CometHashAggregateExec =>
+ case other: CometBaseAggregateExec if other.getClass == getClass =>
this.output == other.output &&
this.groupingExpressions == other.groupingExpressions &&
this.aggregateExpressions == other.aggregateExpressions &&
@@ -2628,6 +2700,40 @@ case class CometHashAggregateExec(
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],
+ aggregateAttributes: Seq[Attribute],
+ 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)
+}
+
+case class CometSortAggregateExec(
+ override val nativeOp: Operator,
+ override val originalPlan: SparkPlan,
+ override val output: Seq[Attribute],
+ groupingExpressions: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ 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)
+}
+
trait CometHashJoin {
// Only BroadcastHashJoinExec can be null-aware (NOT IN subqueries).
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
index 6c37fae562..c2f04e37f4 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set.sql
@@ -16,6 +16,9 @@
-- under the License.
-- ConfigMatrix: parquet.enable.dictionary=false,true
+-- Disabling ObjectHashAggregate makes Spark plan SortAggregateExec for collect_set instead, so the
+-- matrix runs the whole fixture through both aggregate operators.
+-- ConfigMatrix: spark.sql.execution.useObjectHashAggregateExec=true,false
-- ============================================================
-- Setup: tables
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_floating_fallback.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_floating_fallback.sql
index 55bdf13c27..6c8df53388 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_floating_fallback.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_set_floating_fallback.sql
@@ -17,6 +17,9 @@
-- MaxSparkVersion: 4.1
-- Config: spark.comet.exec.strictFloatingPoint=true
+-- Disabling ObjectHashAggregate makes Spark plan SortAggregateExec for collect_set instead, so the
+-- matrix checks that both aggregate operators fall back.
+-- ConfigMatrix: spark.sql.execution.useObjectHashAggregateExec=true,false
statement
CREATE TABLE cs_fallback_float(v float, grp string) USING parquet
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/max_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
index 9c8354f06a..89ac7f0243 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
@@ -17,9 +17,8 @@
-- max_by(x, y) returns the value of x associated with the maximum value of y.
--
--- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate
--- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length
--- value types such as string force SortAggregate and fall back to Spark.
+-- The value (x) must be a fixed-length type. A variable-length value type such as string makes
+-- Spark plan SortAggregate rather than HashAggregate, and Comet falls back for it.
--
-- Ordering values are kept unique within each group so results are deterministic (max_by is
-- non-deterministic when several rows tie on the maximum ordering).
@@ -262,12 +261,11 @@ SELECT grp, max_by(v, ord) FROM mb_signed_zero GROUP BY grp ORDER BY grp
-- ============================================================
-- Variable-length value or ordering falls back to Spark
--
--- A plain max_by over a string is planned as SortAggregate, which Comet never converts, so the
--- serde's type check is not what stops it. Pairing it with a TypedImperativeAggregate switches
--- Spark to ObjectHashAggregate, which Comet does convert, and then getSupportLevel is the only
--- thing that keeps the aggregate off the native path. That matters most for a string *ordering*:
--- Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys, so this
--- check is load-bearing for correctness and not just an optimisation.
+-- A plain max_by over a string is planned as SortAggregate, and pairing it with a
+-- TypedImperativeAggregate switches Spark to ObjectHashAggregate. Comet converts both, so
+-- getSupportLevel is the only thing that keeps the aggregate off the native path. That matters
+-- most for a string *ordering*: Arrow's row format compares raw UTF-8 bytes, while Spark compares
+-- collation sort keys, so this check is load-bearing for correctness and not just an optimisation.
-- ============================================================
statement
@@ -282,3 +280,9 @@ SELECT grp, max_by(s, v), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER B
query expect_fallback(Unsupported ordering data type)
SELECT grp, max_by(v, s), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER BY grp
+
+query expect_fallback(Unsupported value data type)
+SELECT grp, max_by(s, v) FROM mb_varlen GROUP BY grp ORDER BY grp
+
+query expect_fallback(Unsupported ordering data type)
+SELECT grp, max_by(v, s) FROM mb_varlen GROUP BY grp ORDER BY grp
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
index d015acc35d..3d40eee4cd 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
@@ -17,9 +17,8 @@
-- min_by(x, y) returns the value of x associated with the minimum value of y.
--
--- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate
--- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length
--- value types such as string force SortAggregate and fall back to Spark.
+-- The value (x) must be a fixed-length type. A variable-length value type such as string makes
+-- Spark plan SortAggregate rather than HashAggregate, and Comet falls back for it.
--
-- Ordering values are kept unique within each group so results are deterministic (min_by is
-- non-deterministic when several rows tie on the minimum ordering).
@@ -247,9 +246,8 @@ SELECT grp, min_by(v, ord) FROM mnb_signed_zero GROUP BY grp ORDER BY grp
-- ============================================================
-- Variable-length value or ordering falls back to Spark
--
--- See the equivalent section in max_by.sql for why this needs a TypedImperativeAggregate
--- alongside it: without one, Spark plans SortAggregate and Comet never sees the aggregate at all,
--- so the serde's type check would go untested.
+-- See the equivalent section in max_by.sql: with or without a TypedImperativeAggregate
+-- alongside it, the serde's type check is what keeps the aggregate off the native path.
-- ============================================================
statement
@@ -264,3 +262,9 @@ SELECT grp, min_by(s, v), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER
query expect_fallback(Unsupported ordering data type)
SELECT grp, min_by(v, s), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER BY grp
+
+query expect_fallback(Unsupported value data type)
+SELECT grp, min_by(s, v) FROM mnb_varlen GROUP BY grp ORDER BY grp
+
+query expect_fallback(Unsupported ordering data type)
+SELECT grp, min_by(v, s) FROM mnb_varlen GROUP BY grp ORDER BY grp
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..f4cc20d96c
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql
@@ -0,0 +1,79 @@
+-- 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. collect_set.sql runs its whole fixture through
+-- SortAggregateExec as well, so this file keeps only the shapes it does not cover. collect_set
+-- returns elements in nondeterministic order, so every result is wrapped in sort_array.
+-- Config: spark.sql.execution.useObjectHashAggregateExec=false
+
+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)
+
+-- ============================================================
+-- 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
+
+-- ============================================================
+-- 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
+
+-- ============================================================
+-- Decimal SUM at maximum precision falls back: depending on the buffer
+-- schema and codegen, Spark's sort aggregation keeps the running sum
+-- unbounded or nulls it on overflow, and Comet does not track which
+-- ============================================================
+
+statement
+CREATE TABLE sa_dec38(d decimal(38,0), i int, g string) USING parquet
+
+statement
+INSERT INTO sa_dec38 VALUES (1, 1, 'a'), (2, 2, 'a'), (3, 3, 'b'), (NULL, 4, 'b')
+
+query expect_fallback(Decimal SUM at maximum precision cannot match Spark's sort aggregation buffer)
+SELECT g, sum(d), sort_array(collect_set(i)) FROM sa_dec38 GROUP BY g ORDER BY g
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate_without_comet_shuffle.sql
similarity index 55%
copy from spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql
copy to spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate_without_comet_shuffle.sql
index dbec88e2f6..83b274ca0d 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate_without_comet_shuffle.sql
@@ -15,14 +15,16 @@
-- specific language governing permissions and limitations
-- under the License.
-statement
-CREATE TABLE test_min_max(i int, d double, s string, grp string) USING parquet
+-- Without Comet shuffle, a sort aggregate's Partial and Final would be split across Comet and
+-- Spark, and Spark cannot read Comet's collect_set buffer, so SortAggregateExec stays in Spark.
+-- Config: spark.sql.execution.useObjectHashAggregateExec=false
+-- Config: spark.comet.shuffle.enabled=false
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')
+CREATE TABLE sa_no_shuffle(i int, g string) USING parquet
-query expect_fallback(SortAggregate is not supported)
-SELECT min(i), max(i), min(d), max(d), min(s), max(s) FROM test_min_max
+statement
+INSERT INTO sa_no_shuffle VALUES (1, 'a'), (2, 'a'), (1, 'a'), (3, 'b'), (NULL, 'b')
-query
-SELECT grp, min(i), max(i) FROM test_min_max GROUP BY grp ORDER BY grp
+query expect_fallback(Comet shuffle is not enabled, so converting SortAggregate would split the aggregate across Comet and Spark)
+SELECT g, sort_array(collect_set(i)) FROM sa_no_shuffle 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 f0b37eef88..bcdca09172 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -31,12 +31,12 @@ import org.apache.spark.sql.catalyst.expressions.Cast
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.optimizer.EliminateSorts
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning}
-import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec}
+import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortAggregateExec}
import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution.LocalTableScanExec
import org.apache.spark.sql.execution.SQLExecution
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec}
-import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
import org.apache.spark.sql.functions.{avg, col, collect_list, collect_set, count_distinct, expr, sort_array, sum}
import org.apache.spark.sql.internal.SQLConf
@@ -956,26 +956,34 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
// cascade so that neither half runs in Comet - a Spark final cannot read a Comet-produced list,
// and adjustOutputForNativeState would misinterpret Spark's Binary buffer as a list if a Comet
// final ran above a Spark partial. count(*) is included so the aggregate also carries a
- // buffer-compatible function, which is the shape most likely to be split by mistake.
- Seq(
- ("Comet partial + Spark final", CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE),
- ("Spark partial + Comet final", CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE)).foreach {
- case (name, disabledHalf) =>
- test(s"mixed engine collect_list: $name matches Spark") {
- val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7))
- withParquetTable(data, "tbl") {
- withSQLConf(
- disabledHalf.key -> "false",
- CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
- val df = sql("SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2")
- checkSparkAnswer(df)
- // Without the cascade the surviving half would still convert, so pinning this at zero
- // is what keeps the test from passing on a regression.
- assert(getNumCometHashAggregate(df) == 0)
- }
+ // buffer-compatible function, which is the shape most likely to be split by mistake. Disabling
+ // ObjectHashAggregate plans the same aggregate as SortAggregateExec, which puts a sort between
+ // each half and its exchange.
+ for {
+ (name, disabledHalf) <- Seq(
+ ("Comet partial + Spark final", CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE),
+ ("Spark partial + Comet final", CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE))
+ objectHash <- Seq(true, false)
+ } {
+ val suffix = if (objectHash) "" else " with sort aggregates"
+ test(s"mixed engine collect_list: $name matches Spark$suffix") {
+ val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7))
+ withParquetTable(data, "tbl") {
+ withSQLConf(
+ disabledHalf.key -> "false",
+ SQLConf.USE_OBJECT_HASH_AGG.key -> objectHash.toString,
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
+ val df = sql("SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2")
+ checkSparkAnswer(df)
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(plan.find(_.isInstanceOf[SortAggregateExec]).isDefined != objectHash, plan)
+ // Without the cascade the surviving half would still convert, so pinning this at zero
+ // is what keeps the test from passing on a regression.
+ assert(plan.collect { case agg: CometBaseAggregateExec => agg }.isEmpty, plan)
}
}
+ }
}
test("Aggregation without aggregate expressions should use correct result expressions") {
@@ -3505,6 +3513,132 @@ 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 tests collect_set.sql (whose ConfigMatrix runs it
+ // through SortAggregateExec) and sort_aggregate.sql; this test additionally asserts that Spark
+ // actually planned a SortAggregateExec, which the SQL framework cannot check.
+ private def assertSortAggregateRunsNatively(query: String): Unit = {
+ withSQLConf(
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ CometConf.COMET_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") {
+ Seq(
+ "SELECT g, sort_array(collect_set(v)) FROM tbl GROUP BY g ORDER BY g",
+ // Empty grouping is a distinct plan shape: no pre-aggregate sort, empty output ordering,
+ // and adjustOutputForNativeState with zero grouping columns.
+ "SELECT sort_array(collect_set(v)) FROM tbl").foreach(assertSortAggregateRunsNatively)
+ }
+
+ // Spark puts a SortExec between each sort aggregate and the exchange below it, so the passes
+ // that pair a buffer-consuming aggregate with its Partial must walk through it. If they stop
+ // there, a native collect_list Partial stays native under a Spark consumer.
+ private def assertSortAggregateStaysInSpark(query: String, confs: (String, String)*): Unit = {
+ Seq("false", "true").foreach { aqe =>
+ withSQLConf(
+ Seq(
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true") ++ confs: _*) {
+ withParquetTable((0 until 8).map(i => (i % 2, i)), "tbl") {
+ val df = sql(query)
+ checkSparkAnswer(df)
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case agg: SortAggregateExec => agg }.nonEmpty, plan)
+ assert(collect(plan) { case agg: CometSortAggregateExec => agg }.isEmpty, plan)
+ }
+ }
+ }
+ }
+
+ test("SortAggregate keeps a collect_list partial in Spark when its final cannot convert") {
+ assertSortAggregateStaysInSpark(
+ "SELECT _1, sort_array(collect_list(_2)) FROM tbl GROUP BY _1",
+ CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false")
+ }
+
+ test("SortAggregate runs a distinct collect_list chain natively") {
+ // Each stage of the distinct rewrite runs as a native sort aggregate, including the
+ // collect_list PartialMerge that carries the native list state between the shuffles.
+ val query = "SELECT _1, sort_array(collect_list(_2)), count(DISTINCT _2) FROM tbl GROUP BY _1"
+ Seq("false", "true").foreach { aqe =>
+ withSQLConf(
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true") {
+ withParquetTable((0 until 8).map(i => (i % 2, i)), "tbl") {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query))
+ assert(collect(cometPlan) { case a: CometSortAggregateExec => a }.nonEmpty, cometPlan)
+ assert(cometPlan.toString.contains("merge_collect_list"), cometPlan)
+ }
+ }
+ }
+ }
+
+ test("SortAggregate keeps a distinct collect_list chain in Spark under a Spark final") {
+ // Every buffer producer between the final and the bottom Partial, the PartialMerge stages of
+ // the distinct rewrite included, must stay in Spark, through the sort below each of them.
+ assertSortAggregateStaysInSpark(
+ "SELECT _1, sort_array(collect_list(_2)), count(DISTINCT _2) FROM tbl GROUP BY _1",
+ CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false")
+ }
+
+ test("SortAggregate keeps a decimal AVG with a maximum-precision sum in Spark") {
+ // A string FIRST cannot use a hash aggregation buffer, so Spark plans sort aggregates and
+ // keeps the AVG sum in a generic row, where it is unbounded: 0.6 + 0.6 leaves DECIMAL(38,38)
+ // and -0.4 brings it back. The native AVG records that overflow and would return NULL, or
+ // raise under ANSI. Sorting on ord fixes the order in which the partial adds the values.
+ withTempPath { dir =>
+ Seq((1, 1, "0.6"), (1, 2, "0.6"), (1, 3, "-0.4"))
+ .toDF("g", "ord", "raw")
+ .selectExpr("g", "ord", "CAST(raw AS DECIMAL(38,38)) AS v", "'x' AS label")
+ .coalesce(1)
+ .write
+ .parquet(dir.toString)
+ withParquetTable(dir.toString, "dec_avg") {
+ Seq("false", "true").foreach { ansi =>
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> ansi,
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ "SELECT g, avg(v), first(label) FROM (SELECT * FROM dec_avg SORT BY g, ord) " +
+ "GROUP BY g")
+ val (sparkPlan, cometPlan) = checkSparkAnswerAndFallbackReason(
+ df,
+ "Decimal AVG with a maximum-precision sum cannot match Spark's " +
+ "sort aggregation buffer")
+ assert(collect(sparkPlan) { case agg: SortAggregateExec => agg }.nonEmpty, sparkPlan)
+ assert(
+ collect(cometPlan) { case agg: CometSortAggregateExec => agg }.isEmpty,
+ cometPlan)
+ checkAnswer(
+ df,
+ Row(1, new java.math.BigDecimal("0.26666666666666666666666666666666666667"), "x"))
+ }
+ }
+ }
+ }
+ }
+
// 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
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index 4c90612c3a..0eca67e613 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -37,7 +37,7 @@ import org.apache.spark.SparkConf
import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row}
import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal}
import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch}
-import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec}
+import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec}
import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper}
import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager
import org.apache.spark.sql.comet.util.Utils
@@ -910,6 +910,35 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ test("sort aggregate over a sorted cache keeps its grouping-key order") {
+ // The cached relation reports its sort order, so Spark plans both sort aggregates and the
+ // ORDER BY without a SortExec, and the native aggregates read the cache through a scan that
+ // reports no order. NULL and an empty array hash alike, so DataFusion's grouping can emit
+ // the empty array after [1] unless the aggregate output is sorted natively.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ try {
+ Seq[(Option[Seq[Int]], String)]((None, "n"), (Some(Seq.empty), "e"), (Some(Seq(1)), "o"))
+ .toDF("k", "v")
+ .coalesce(1)
+ .sortWithinPartitions("k")
+ .cache()
+ .createOrReplaceTempView("sorted_cache")
+
+ val df = sql("SELECT k, first(v) FROM sorted_cache GROUP BY k ORDER BY k")
+ checkAnswer(df, Seq(Row(null, "n"), Row(Seq.empty[Int], "e"), Row(Seq(1), "o")))
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan)
+ assert(collect(plan) { case a: CometSortAggregateExec => a }.size == 2, plan)
+ assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan)
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+
test("Comet in-memory cache supports shuffle after cache read") {
withConversions(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index f0202c25b5..a882f93125 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -40,7 +40,7 @@ import org.apache.spark.sql.comet.CometExec.nativeLeaves
import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec, SimpleCost, SimpleCostEvaluator}
-import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
import org.apache.spark.sql.internal.SQLConf
@@ -1442,6 +1442,61 @@ class CometExecRuleSuite extends CometTestBase {
}
}
+ test("buffer repair walks through the sort below a sort aggregate") {
+ // Spark puts a SortExec between a sort aggregate and its exchange, converted or not. Repair
+ // must rebuild through it to restore the native Partial, and when repair stops at a stage it
+ // must still find the native Partial below the sort and warn.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+ withTempView("test_data") {
+ createTestDataFrame.createOrReplaceTempView("test_data")
+ val plan = applyCometExecRule(
+ createSparkPlan(spark, "SELECT collect_list(id) FROM test_data GROUP BY (id % 3)"))
+ val nativeFinal = plan
+ .collectFirst { case agg: CometSortAggregateExec if agg.modes == Seq(Final) => agg }
+ .getOrElse(fail(s"Expected a native final sort aggregate in:\n$plan"))
+ val partial = plan
+ .collectFirst { case agg: CometSortAggregateExec if agg.modes == Seq(Partial) => agg }
+ .getOrElse(fail(s"Expected a native partial sort aggregate in:\n$plan"))
+ val sparkFinal = nativeFinal.originalPlan.asInstanceOf[SortAggregateExec]
+ val cometSort = nativeFinal.child.asInstanceOf[CometSortExec]
+ val sorts: Seq[SparkPlan => SparkPlan] = Seq(
+ child => cometSort.copy(child = child),
+ child => cometSort.originalPlan.withNewChildren(Seq(child)))
+ val rule = CometExecRule(spark)
+
+ sorts.foreach { sortOver =>
+ val repaired =
+ rule.revertUnsafePartialAggregates(sparkFinal.copy(child = sortOver(cometSort.child)))
+ assert(repaired.collect { case agg: CometSortAggregateExec => agg }.isEmpty, repaired)
+ val sparkPartial = repaired.collectFirst {
+ case agg: SortAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) =>
+ agg
+ }.get
+ assert(sparkPartial.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined)
+
+ val exchange = ShuffleExchangeExec(
+ org.apache.spark.sql.catalyst.plans.physical.SinglePartition,
+ partial)
+ val stage = ShuffleQueryStageExec(0, exchange, exchange.canonicalized)
+ val placeholder = CometSinkPlaceHolder(
+ org.apache.comet.serde.OperatorOuterClass.Operator.getDefaultInstance,
+ stage,
+ stage)
+ val consumer = sparkFinal.copy(child = sortOver(placeholder))
+ assert(rule.revertUnsafePartialAggregates(consumer) eq consumer)
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(consumer)
+ assert(reasons.exists(_.contains("could not restore a native intermediate buffer")))
+ }
+ }
+ }
+ }
+
test("CometExecRule should not allow decimal SUM mixed execution") {
withTempView("test_data") {
createTestDataFrame.createOrReplaceTempView("test_data")
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
index 0e2b47e15b..a63560f337 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Attribut
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition}
-import org.apache.spark.sql.comet.{CometCollectLimitExec, CometHashAggregateExec, CometLocalTableScanExec, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSortExec, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometCollectLimitExec, CometHashAggregateExec, CometLocalTableScanExec, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSortAggregateExec, CometSortExec, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec}
import org.apache.spark.sql.execution.{CollectLimitExec, ColumnarToRowTransition, LocalTableScanExec, SortExec, SparkPlan, TakeOrderedAndProjectExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
@@ -453,6 +453,37 @@ class CometCelebornShufflePlanningSuite extends CometTestBase {
}
}
+ // SortAggregateExec carries the same TypedImperativeAggregate buffer formats as
+ // ObjectHashAggregateExec, so preserveSparkAggregateBuffers must restore a Comet sort-aggregate
+ // Partial too. Regression guard: that predicate matched only CometHashAggregateExec before
+ // CometSortAggregateExec existed, which would have let a native collect_list buffer reach a
+ // Spark Final across a fallen-back Celeborn exchange.
+ test("a Comet sort-aggregate Partial is restored when a Celeborn exchange falls back") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false") {
+ val partial = collect(collectionAggregatePlan) {
+ case aggregate: CometSortAggregateExec if aggregate.modes == Seq(Partial) => aggregate
+ }.head
+
+ manager.withPlanningSupport(
+ CelebornNativeShufflePlanningSupport(fallbackPartitionThreshold = 2L)) {
+ val exchange =
+ ShuffleExchangeExec(HashPartitioning(Seq(partial.output.head), 4), partial)
+ val restored = CometExecRule(spark)(exchange)
+ assertSparkExchange(restored)
+ assert(
+ collect(restored) { case aggregate: CometBaseAggregateExec => aggregate }.isEmpty,
+ s"$restored")
+ val sparkPartial = collect(restored) { case aggregate: BaseAggregateExec =>
+ aggregate
+ }.head
+ assert(sparkPartial.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined)
+ }
+ }
+ }
+
test("transition reversion preserves native aggregate buffers across a Celeborn exchange") {
manager.withPlanningSupport(CelebornNativeShufflePlanningSupport()) {
withSQLConf(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org