From 38a09585a4811e97ef8854f9b05e24fb40b874e4 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Tue, 2 Jun 2026 08:11:32 -0600
Subject: [PATCH 1/3] 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.
---
 .../apache/comet/rules/CometExecRule.scala    |  14 +-
 .../apache/spark/sql/comet/operators.scala    | 163 +++++++++++-------
 .../comet/exec/CometAggregateSuite.scala      |  27 +++
 3 files changed, 135 insertions(+), 69 deletions(-)

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 d116d2f407..8c94411001 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -33,7 +33,7 @@ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, Comet
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
-import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
 import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
 import org.apache.spark.sql.execution.datasources.WriteFilesExec
 import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat
@@ -80,6 +80,7 @@ object CometExecRule {
       classOf[GenerateExec] -> CometExplodeExec,
       classOf[HashAggregateExec] -> CometHashAggregateExec,
       classOf[ObjectHashAggregateExec] -> CometObjectHashAggregateExec,
+      classOf[SortAggregateExec] -> CometSortAggregateExec,
       classOf[BroadcastHashJoinExec] -> CometBroadcastHashJoinExec,
       classOf[ShuffledHashJoinExec] -> CometHashJoinExec,
       classOf[SortMergeJoinExec] -> CometSortMergeJoinExec,
@@ -149,8 +150,7 @@ case class CometExecRule(session: SparkSession)
    * Comet columnar shuffle.
    */
   private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = {
-    def isAggregate(p: SparkPlan): Boolean =
-      p.isInstanceOf[HashAggregateExec] || p.isInstanceOf[ObjectHashAggregateExec]
+    def isAggregate(p: SparkPlan): Boolean = p.isInstanceOf[BaseAggregateExec]
 
     def isRedundantShuffle(child: SparkPlan): Boolean = child match {
       case s: CometShuffleExchangeExec =>
@@ -858,9 +858,13 @@ case class CometExecRule(session: SparkSession)
     val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]]
     if (!isOperatorEnabled(serde, agg.asInstanceOf[SparkPlan])) return false
 
-    // ObjectHashAggregate has an extra shuffle-enabled guard in its convert method
+    // ObjectHashAggregate / SortAggregate carry TypedImperativeAggregate functions whose
+    // intermediate buffer formats differ between Spark and Comet, so the Partial->Final pair
+    // must travel via Comet shuffle.
     agg match {
-      case _: ObjectHashAggregateExec if !isCometShuffleEnabled(agg.conf) => return false
+      case _: ObjectHashAggregateExec | _: SortAggregateExec
+          if !isCometShuffleEnabled(agg.conf) =>
+        return false
       case _ =>
     }
 
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index 8cbf7c9189..1b1fd6744f 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
@@ -38,7 +38,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
-import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
@@ -1405,6 +1405,58 @@ case class CometUnionExec(
 
 trait CometBaseAggregate {
 
+  /**
+   * Shared support-level check used by every aggregate serde: honor the unit-test knobs that
+   * selectively disable Comet conversion for partial or final aggregates.
+   */
+  protected def baseAggregateSupportLevel(op: BaseAggregateExec): SupportLevel = {
+    if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) &&
+      op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) {
+      return Unsupported(Some("Partial aggregates disabled via test config"))
+    }
+    if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) &&
+      op.aggregateExpressions.exists(_.mode == Final)) {
+      return Unsupported(Some("Final aggregates disabled via test config"))
+    }
+    Compatible()
+  }
+
+  /**
+   * For Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet),
+   * the Spark-side output declares buffer columns as BinaryType (Spark serializes state to
+   * binary). However, the native Comet aggregate produces the actual state type (e.g.,
+   * ArrayType(elementType) for CollectSet). This corrects the output schema to match the native
+   * state types so the shuffle exchange schema is consistent with the actual data.
+   *
+   * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a
+   * case branch here mapping it to the native state type.
+   */
+  protected def adjustOutputForNativeState(op: BaseAggregateExec): Seq[Attribute] = {
+    val modes = op.aggregateExpressions.map(_.mode).distinct
+    if (modes != Seq(Partial)) {
+      return op.output
+    }
+
+    val numGrouping = op.groupingExpressions.length
+    val output = op.output.toArray
+
+    var bufferIdx = numGrouping
+    for (aggExpr <- op.aggregateExpressions) {
+      val aggFunc = aggExpr.aggregateFunction
+      val bufferAttrs = aggFunc.aggBufferAttributes
+      aggFunc match {
+        case cs: CollectSet =>
+          val elementType = cs.children.head.dataType
+          val nativeStateType = ArrayType(elementType, containsNull = true)
+          output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType)
+        case _ =>
+      }
+      bufferIdx += bufferAttrs.length
+    }
+
+    output.toSeq
+  }
+
   def doConvert(
       aggregate: BaseAggregateExec,
       builder: Operator.Builder,
@@ -1625,8 +1677,9 @@ 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] = {
     def isPartialOrMerge(mode: AggregateMode): Boolean =
@@ -1636,10 +1689,7 @@ trait CometBaseAggregate {
       case agg: CometHashAggregateExec
           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)
@@ -1656,19 +1706,8 @@ object CometHashAggregateExec
   override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
     CometConf.COMET_EXEC_AGGREGATE_ENABLED)
 
-  override def getSupportLevel(op: HashAggregateExec): SupportLevel = {
-    // some unit tests need to disable partial or final hash aggregate support to test that
-    // CometExecRule does not allow mixed Spark/Comet aggregates
-    if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) &&
-      op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) {
-      return Unsupported(Some("Partial aggregates disabled via test config"))
-    }
-    if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) &&
-      op.aggregateExpressions.exists(_.mode == Final)) {
-      return Unsupported(Some("Final aggregates disabled via test config"))
-    }
-    Compatible()
-  }
+  override def getSupportLevel(op: HashAggregateExec): SupportLevel =
+    baseAggregateSupportLevel(op)
 
   override def convert(
       aggregate: HashAggregateExec,
@@ -1698,19 +1737,8 @@ object CometObjectHashAggregateExec
   override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
     CometConf.COMET_EXEC_AGGREGATE_ENABLED)
 
-  override def getSupportLevel(op: ObjectHashAggregateExec): SupportLevel = {
-    // Mirror the same test-knobs as CometHashAggregateExec so that mixed-execution
-    // unit tests can selectively disable partial or final ObjectHashAggregateExec conversion.
-    if (!CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.get(op.conf) &&
-      op.aggregateExpressions.exists(expr => expr.mode == Partial || expr.mode == PartialMerge)) {
-      return Unsupported(Some("Partial aggregates disabled via test config"))
-    }
-    if (!CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.get(op.conf) &&
-      op.aggregateExpressions.exists(_.mode == Final)) {
-      return Unsupported(Some("Final aggregates disabled via test config"))
-    }
-    Compatible()
-  }
+  override def getSupportLevel(op: ObjectHashAggregateExec): SupportLevel =
+    baseAggregateSupportLevel(op)
 
   override def convert(
       aggregate: ObjectHashAggregateExec,
@@ -1739,42 +1767,49 @@ object CometObjectHashAggregateExec
       op.child,
       SerializedPlan(None))
   }
+}
 
-  /**
-   * For Partial mode aggregates containing TypedImperativeAggregate functions (like CollectSet),
-   * the Spark-side output declares buffer columns as BinaryType (since Spark serializes state to
-   * binary). However, the native Comet aggregate produces the actual state type (e.g.,
-   * ArrayType(elementType) for CollectSet). This method corrects the output schema to match the
-   * native state types so the shuffle exchange schema is consistent with the actual data.
-   *
-   * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a
-   * case branch here mapping it to the native state type.
-   */
-  private def adjustOutputForNativeState(op: ObjectHashAggregateExec): Seq[Attribute] = {
-    // This adjustment only applies to pure-Partial aggregates (checked below).
-    val modes = op.aggregateExpressions.map(_.mode).distinct
-    if (modes != Seq(Partial)) {
-      return op.output
-    }
+object CometSortAggregateExec
+    extends CometOperatorSerde[SortAggregateExec]
+    with CometBaseAggregate {
 
-    val numGrouping = op.groupingExpressions.length
-    val output = op.output.toArray
+  override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+    CometConf.COMET_EXEC_AGGREGATE_ENABLED)
 
-    var bufferIdx = numGrouping
-    for (aggExpr <- op.aggregateExpressions) {
-      val aggFunc = aggExpr.aggregateFunction
-      val bufferAttrs = aggFunc.aggBufferAttributes
-      aggFunc match {
-        case cs: CollectSet =>
-          val elementType = cs.children.head.dataType
-          val nativeStateType = ArrayType(elementType, containsNull = true)
-          output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType)
-        case _ =>
-      }
-      bufferIdx += bufferAttrs.length
+  override def getSupportLevel(op: SortAggregateExec): SupportLevel =
+    baseAggregateSupportLevel(op)
+
+  override def convert(
+      aggregate: SortAggregateExec,
+      builder: Operator.Builder,
+      childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
+
+    // SortAggregate is planned for TypedImperativeAggregate functions whose intermediate
+    // buffer formats differ between Spark and Comet (same risk as ObjectHashAggregate).
+    // Require Comet shuffle so a Partial->Final pair never spans the JVM/native boundary.
+    if (!isCometShuffleEnabled(aggregate.conf)) {
+      return None
     }
 
-    output.toSeq
+    doConvert(aggregate, builder, childOp: _*)
+  }
+
+  override def createExec(nativeOp: Operator, op: SortAggregateExec): CometNativeExec = {
+    // Reuse CometHashAggregateExec as the wrapper. The native AggregateExec auto-detects
+    // Sorted input mode from the child's output ordering and produces output sorted by the
+    // grouping keys; CometExec.outputOrdering defaults to originalPlan.outputOrdering, which
+    // is SortAggregateExec's grouping-key ordering, so downstream operators that elided a
+    // sort against it still see a satisfying ordering without a dedicated wrapper class.
+    CometHashAggregateExec(
+      nativeOp,
+      op,
+      adjustOutputForNativeState(op),
+      op.groupingExpressions,
+      op.aggregateExpressions,
+      op.resultExpressions,
+      op.child.output,
+      op.child,
+      SerializedPlan(None))
   }
 }
 
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 cd0beb56cc..1ae80523b0 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.Cast
 import org.apache.spark.sql.catalyst.optimizer.EliminateSorts
 import org.apache.spark.sql.comet.CometHashAggregateExec
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.aggregate.SortAggregateExec
 import org.apache.spark.sql.functions.{avg, col, count_distinct, sum}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
@@ -2109,4 +2110,30 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("SortAggregate with collect_set is converted to native") {
+    // useObjectHashAggregateExec=false forces Spark to plan SortAggregateExec for
+    // TypedImperativeAggregate functions like collect_set. Comet converts those just like
+    // ObjectHashAggregateExec via the shared CometBaseAggregate path.
+    withSQLConf(
+      "spark.sql.execution.useObjectHashAggregateExec" -> "false",
+      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
+      withTempView("tbl") {
+        Seq((1, "a"), (2, "a"), (1, "a"), (3, "b"), (4, "b"), (4, "b"))
+          .toDF("v", "g")
+          .createOrReplaceTempView("tbl")
+        val query = "SELECT g, sort_array(collect_set(v)) FROM tbl GROUP BY g ORDER BY g"
+        // 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))
+      }
+    }
+  }
+
 }

From 640976a0b5808de13ea71685d569f25870372126 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Tue, 30 Jun 2026 10:12:57 -0600
Subject: [PATCH 2/3] 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.
---
 .../expressions/aggregate/sort_aggregate.sql  | 214 ++++++++++++++++++
 .../comet/exec/CometAggregateSuite.scala      |  24 +-
 2 files changed, 233 insertions(+), 5 deletions(-)
 create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql

diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql
new file mode 100644
index 0000000000..18b0dac7a0
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql
@@ -0,0 +1,214 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements.  See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership.  The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License.  You may obtain a copy of the License at
+--
+--   http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied.  See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- Disabling ObjectHashAggregate forces Spark to plan SortAggregateExec for the
+-- TypedImperativeAggregate function collect_set, which is the operator this file exercises.
+-- (The plan-shape assertion that a SortAggregateExec is actually produced lives in
+-- CometAggregateSuite; here the default `query` mode asserts the whole plan, including the
+-- pre-aggregate sorts and the Comet shuffle, runs natively and matches Spark.)
+-- collect_set returns elements in nondeterministic order, so every result is wrapped in
+-- sort_array for a stable comparison against Spark.
+-- Config: spark.sql.execution.useObjectHashAggregateExec=false
+-- Config: spark.comet.exec.strictFloatingPoint=true
+
+-- ============================================================
+-- Setup: tables
+-- ============================================================
+
+statement
+CREATE TABLE sa_int(i int, g string) USING parquet
+
+statement
+INSERT INTO sa_int VALUES
+  (1, 'a'), (2, 'a'), (1, 'a'), (3, 'a'),
+  (4, 'b'), (4, 'b'), (NULL, 'b'), (5, 'b'),
+  (NULL, 'c'), (NULL, 'c')
+
+statement
+CREATE TABLE sa_multikey(a int, k1 string, k2 int) USING parquet
+
+statement
+INSERT INTO sa_multikey VALUES
+  (1, 'x', 10), (2, 'x', 10), (1, 'x', 10),
+  (3, 'x', 20), (4, 'y', 10), (NULL, 'y', 10)
+
+statement
+CREATE TABLE sa_nulls(v int, g string) USING parquet
+
+statement
+INSERT INTO sa_nulls VALUES (NULL, 'a'), (NULL, 'a'), (NULL, 'b'), (1, 'b')
+
+statement
+CREATE TABLE sa_empty(v int, g string) USING parquet
+
+statement
+CREATE TABLE sa_single(v int) USING parquet
+
+statement
+INSERT INTO sa_single VALUES (42)
+
+-- ============================================================
+-- Global aggregate (no GROUP BY): single SortAggregate group
+-- ============================================================
+
+query
+SELECT sort_array(collect_set(i)) FROM sa_int
+
+-- ============================================================
+-- Single grouping key
+-- ============================================================
+
+query
+SELECT g, sort_array(collect_set(i)) FROM sa_int GROUP BY g ORDER BY g
+
+-- ============================================================
+-- Multiple grouping keys (sort ordering spans both keys)
+-- ============================================================
+
+query
+SELECT k1, k2, sort_array(collect_set(a))
+FROM sa_multikey GROUP BY k1, k2 ORDER BY k1, k2
+
+-- ============================================================
+-- Grouping key derived from an expression
+-- ============================================================
+
+query
+SELECT upper(g) AS gg, sort_array(collect_set(i))
+FROM sa_int GROUP BY upper(g) ORDER BY gg
+
+-- ============================================================
+-- All-NULL group returns an empty array
+-- ============================================================
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_nulls GROUP BY g ORDER BY g
+
+-- ============================================================
+-- Empty table
+-- ============================================================
+
+query
+SELECT sort_array(collect_set(v)) FROM sa_empty
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_empty GROUP BY g ORDER BY g
+
+-- ============================================================
+-- Single row
+-- ============================================================
+
+query
+SELECT sort_array(collect_set(v)) FROM sa_single
+
+-- ============================================================
+-- collect_set mixed with hashable aggregates: collect_set forces the
+-- whole aggregate onto SortAggregate, so count/sum/min/max ride along
+-- ============================================================
+
+query
+SELECT g, sort_array(collect_set(i)), count(*), count(i), sum(i), min(i), max(i)
+FROM sa_int GROUP BY g ORDER BY g
+
+-- ============================================================
+-- Multiple collect_set in one query
+-- ============================================================
+
+query
+SELECT k1, sort_array(collect_set(a)), sort_array(collect_set(k2))
+FROM sa_multikey GROUP BY k1 ORDER BY k1
+
+-- ============================================================
+-- DISTINCT (distinct-aggregate planner path)
+-- ============================================================
+
+query
+SELECT g, sort_array(collect_set(DISTINCT i)) FROM sa_int GROUP BY g ORDER BY g
+
+-- ============================================================
+-- HAVING over a collect_set-derived predicate
+-- ============================================================
+
+query
+SELECT g, sort_array(collect_set(i))
+FROM sa_int GROUP BY g HAVING size(collect_set(i)) > 1 ORDER BY g
+
+-- ============================================================
+-- Representative data types through the SortAggregate path
+-- ============================================================
+
+statement
+CREATE TABLE sa_string(v string, g string) USING parquet
+
+statement
+INSERT INTO sa_string VALUES
+  ('hello', 'a'), ('world', 'a'), ('hello', 'a'), (NULL, 'a'),
+  ('', 'b'), ('x', 'b'), ('', 'b'), (NULL, 'b')
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_string GROUP BY g ORDER BY g
+
+statement
+CREATE TABLE sa_bigint(v bigint, g string) USING parquet
+
+statement
+INSERT INTO sa_bigint VALUES
+  (1000000000000, 'a'), (2000000000000, 'a'), (1000000000000, 'a'), (NULL, 'a'),
+  (3000000000000, 'b'), (NULL, 'b')
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_bigint GROUP BY g ORDER BY g
+
+statement
+CREATE TABLE sa_decimal(v decimal(10,2), g string) USING parquet
+
+statement
+INSERT INTO sa_decimal VALUES
+  (1.50, 'a'), (2.50, 'a'), (1.50, 'a'), (NULL, 'a'),
+  (0.00, 'b'), (99999999.99, 'b'), (NULL, 'b')
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_decimal GROUP BY g ORDER BY g
+
+statement
+CREATE TABLE sa_date(v date, g string) USING parquet
+
+statement
+INSERT INTO sa_date VALUES
+  (DATE '2024-01-01', 'a'), (DATE '2024-06-15', 'a'), (DATE '2024-01-01', 'a'), (NULL, 'a'),
+  (DATE '1970-01-01', 'b'), (NULL, 'b')
+
+query
+SELECT g, sort_array(collect_set(v)) FROM sa_date GROUP BY g ORDER BY g
+
+-- ============================================================
+-- Unsupported aggregate input under the SortAggregate path:
+-- with strictFloatingPoint=true, collect_set on floating-point types is
+-- incompatible, so the SortAggregateExec must fall back to Spark cleanly
+-- rather than crash.
+-- ============================================================
+
+statement
+CREATE TABLE sa_double(v double, g string) USING parquet
+
+statement
+INSERT INTO sa_double VALUES
+  (1.1, 'a'), (2.2, 'a'), (1.1, 'a'), (NULL, 'a'),
+  (CAST('NaN' AS DOUBLE), 'b'), (CAST('NaN' AS DOUBLE), 'b'), (1.0, 'b')
+
+query expect_fallback(not fully compatible with Spark)
+SELECT g, sort_array(collect_set(v)) FROM sa_double GROUP BY g ORDER BY g
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
index 4079be4529..db4c93bf7d 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -2110,10 +2110,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
-  test("SortAggregate with collect_set is converted to native") {
-    // useObjectHashAggregateExec=false forces Spark to plan SortAggregateExec for
-    // TypedImperativeAggregate functions like collect_set. Comet converts those just like
-    // ObjectHashAggregateExec via the shared CometBaseAggregate path.
+  // useObjectHashAggregateExec=false forces Spark to plan SortAggregateExec for
+  // TypedImperativeAggregate functions like collect_set. Comet converts those just like
+  // ObjectHashAggregateExec via the shared CometBaseAggregate path. Broader data-type and
+  // edge-case coverage lives in the SQL file test
+  // spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql; these Scala
+  // tests additionally assert that Spark actually planned a SortAggregateExec, which the SQL
+  // framework cannot check.
+  private def assertSortAggregateRunsNatively(query: String): Unit = {
     withSQLConf(
       "spark.sql.execution.useObjectHashAggregateExec" -> "false",
       CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
@@ -2122,7 +2126,6 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         Seq((1, "a"), (2, "a"), (1, "a"), (3, "b"), (4, "b"), (4, "b"))
           .toDF("v", "g")
           .createOrReplaceTempView("tbl")
-        val query = "SELECT g, sort_array(collect_set(v)) FROM tbl GROUP BY g ORDER BY g"
         // 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") {
@@ -2136,6 +2139,17 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("SortAggregate with collect_set is converted to native") {
+    assertSortAggregateRunsNatively(
+      "SELECT g, sort_array(collect_set(v)) FROM tbl GROUP BY g ORDER BY g")
+  }
+
+  test("SortAggregate global collect_set (no grouping keys) is converted to native") {
+    // Empty grouping is a distinct plan shape: no pre-aggregate sort, empty output ordering,
+    // and adjustOutputForNativeState with zero grouping columns.
+    assertSortAggregateRunsNatively("SELECT sort_array(collect_set(v)) FROM tbl")
+  }
+
   // Regression: Catalyst prunes `HashAggregateExec.resultExpressions` to
   // empty for EXISTS / row-existence-only subqueries. The native HashAggregate's natural
   // output (the grouping keys) then disagrees with the pruned JVM `output`, leaking through

From 3a627ae88393b31fd461312f6156528c95381e41 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Tue, 30 Jun 2026 14:00:02 -0600
Subject: [PATCH 3/3] 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.
---
 dev/diffs/3.4.3.diff                          | 25 +++--
 dev/diffs/3.5.8.diff                          | 25 +++--
 dev/diffs/4.0.2.diff                          | 25 +++--
 dev/diffs/4.1.2.diff                          | 25 +++--
 .../apache/spark/sql/comet/operators.scala    | 95 ++++++++++++++-----
 .../expressions/aggregate/first_last.sql      | 15 ++-
 .../expressions/aggregate/min_max.sql         |  5 +-
 7 files changed, 148 insertions(+), 67 deletions(-)

diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff
index a905f32e3c..87844d0e12 100644
--- a/dev/diffs/3.4.3.diff
+++ b/dev/diffs/3.4.3.diff
@@ -1473,26 +1473,31 @@ index 30ce940b032..0d3f6c6c934 100644
  
    private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-index 47679ed7865..9ffbaecb98e 100644
+index 47679ed7865..cbcabe00992 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 @@ -18,6 +18,7 @@
  package org.apache.spark.sql.execution
  
  import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec}
  import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
  import org.apache.spark.sql.internal.SQLConf
-@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
+@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase
    private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
      val plan = df.queryExecution.executedPlan
      assert(collectWithSubqueries(plan) {
 -      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
 +      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
      }.length == hashAggCount)
-     assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
+-    assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
++    assert(collectWithSubqueries(plan) {
++      case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s
++    }.length == sortAggCount)
    }
+ 
+   private def checkAggs(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
 index eec396b2e39..bf3f1c769d6 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -1563,7 +1568,7 @@ index ac710c32296..2854b433dd3 100644
  
    import testImplicits._
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 593bd7bb4ba..32af28b0238 100644
+index 593bd7bb4ba..2ae8b702068 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 @@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._
@@ -1623,7 +1628,7 @@ index 593bd7bb4ba..32af28b0238 100644
    private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
      collect(plan) {
        case agg: BaseAggregateExec => agg
-+      case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
++      case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
      }
    }
  
@@ -2856,14 +2861,14 @@ index ef5b8a769fe..84fe1bfabc9 100644
        }
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-index b4c4ec7acbf..20579284856 100644
+index b4c4ec7acbf..e57aff1f170 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils
  import org.scalatest.Assertions
  
  import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.CometBaseAggregateExec
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
  import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec}
  import org.apache.spark.sql.functions.count
@@ -2871,7 +2876,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = query.lastExecution.executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
@@ -2879,7 +2884,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff
index 0cbf239981..d822091229 100644
--- a/dev/diffs/3.5.8.diff
+++ b/dev/diffs/3.5.8.diff
@@ -1439,26 +1439,31 @@ index 005e764cc30..92ec088efab 100644
  
    private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-index 47679ed7865..9ffbaecb98e 100644
+index 47679ed7865..cbcabe00992 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 @@ -18,6 +18,7 @@
  package org.apache.spark.sql.execution
  
  import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec}
  import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
  import org.apache.spark.sql.internal.SQLConf
-@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
+@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase
    private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
      val plan = df.queryExecution.executedPlan
      assert(collectWithSubqueries(plan) {
 -      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
 +      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
      }.length == hashAggCount)
-     assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
+-    assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
++    assert(collectWithSubqueries(plan) {
++      case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s
++    }.length == sortAggCount)
    }
+ 
+   private def checkAggs(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
 index eec396b2e39..bf3f1c769d6 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -1529,7 +1534,7 @@ index 5a413c77754..207b66e1d7b 100644
  
    import testImplicits._
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 2f8e401e743..dbcf3171946 100644
+index 2f8e401e743..fc712552c4f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._
@@ -1589,7 +1594,7 @@ index 2f8e401e743..dbcf3171946 100644
    private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
      collect(plan) {
        case agg: BaseAggregateExec => agg
-+      case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
++      case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
      }
    }
  
@@ -2806,14 +2811,14 @@ index c97979a57a5..45a998db0e0 100644
        }
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-index b4c4ec7acbf..20579284856 100644
+index b4c4ec7acbf..e57aff1f170 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils
  import org.scalatest.Assertions
  
  import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.CometBaseAggregateExec
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
  import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec}
  import org.apache.spark.sql.functions.count
@@ -2821,7 +2826,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = query.lastExecution.executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
@@ -2829,7 +2834,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff
index f6d01f4f6a..d2244acc24 100644
--- a/dev/diffs/4.0.2.diff
+++ b/dev/diffs/4.0.2.diff
@@ -1841,26 +1841,31 @@ index 005e764cc30..92ec088efab 100644
  
    private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-index 47679ed7865..9ffbaecb98e 100644
+index 47679ed7865..cbcabe00992 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 @@ -18,6 +18,7 @@
  package org.apache.spark.sql.execution
  
  import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec}
  import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
  import org.apache.spark.sql.internal.SQLConf
-@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
+@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase
    private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
      val plan = df.queryExecution.executedPlan
      assert(collectWithSubqueries(plan) {
 -      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
 +      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
      }.length == hashAggCount)
-     assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
+-    assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
++    assert(collectWithSubqueries(plan) {
++      case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s
++    }.length == sortAggCount)
    }
+ 
+   private def checkAggs(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
 index aa619c5cde8..4b5a42e683f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -2153,7 +2158,7 @@ index a3cfdc5a240..3793b6191bf 100644
        })
        checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 272be70f9fe..12daa1f5932 100644
+index 272be70f9fe..a246594c73c 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 @@ -28,12 +28,14 @@ import org.apache.spark.SparkException
@@ -2216,7 +2221,7 @@ index 272be70f9fe..12daa1f5932 100644
    private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
      collect(plan) {
        case agg: BaseAggregateExec => agg
-+      case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
++      case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
      }
    }
  
@@ -3470,14 +3475,14 @@ index b0967d5ffdf..3d567f913de 100644
        }
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-index b4c4ec7acbf..20579284856 100644
+index b4c4ec7acbf..e57aff1f170 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 @@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils
  import org.scalatest.Assertions
  
  import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.CometBaseAggregateExec
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
  import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec}
  import org.apache.spark.sql.functions.count
@@ -3485,7 +3490,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = query.lastExecution.executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
@@ -3493,7 +3498,7 @@ index b4c4ec7acbf..20579284856 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff
index e084bc9913..65a278bcb9 100644
--- a/dev/diffs/4.1.2.diff
+++ b/dev/diffs/4.1.2.diff
@@ -1964,26 +1964,31 @@ index 005e764cc30..92ec088efab 100644
  
    private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-index 47679ed7865..9ffbaecb98e 100644
+index 47679ed7865..cbcabe00992 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
 @@ -18,6 +18,7 @@
  package org.apache.spark.sql.execution
  
  import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec}
  import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
  import org.apache.spark.sql.internal.SQLConf
-@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
+@@ -31,9 +32,11 @@ abstract class ReplaceHashWithSortAggSuiteBase
    private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
      val plan = df.queryExecution.executedPlan
      assert(collectWithSubqueries(plan) {
 -      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
 +      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
      }.length == hashAggCount)
-     assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
+-    assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
++    assert(collectWithSubqueries(plan) {
++      case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s
++    }.length == sortAggCount)
    }
+ 
+   private def checkAggs(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
 index aa619c5cde8..4b5a42e683f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -2276,7 +2281,7 @@ index a3cfdc5a240..3793b6191bf 100644
        })
        checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 3e7d26f74bd..7e70e72fa3e 100644
+index 3e7d26f74bd..3a4f383777e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 @@ -27,12 +27,14 @@ import org.apache.spark.SparkException
@@ -2339,7 +2344,7 @@ index 3e7d26f74bd..7e70e72fa3e 100644
    private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
      collect(plan) {
        case agg: BaseAggregateExec => agg
-+      case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
++      case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
      }
    }
  
@@ -3713,14 +3718,14 @@ index e57c4e1e665..1dd8b80a006 100644
    // Necessary to make ScalaTest 3.x interrupt a thread on the JVM like ScalaTest 2.2.x
    implicit val defaultSignaler: Signaler = ThreadSignaler
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-index 465da3cd469..92ac998929d 100644
+index 465da3cd469..2e6bcb6ca0d 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 @@ -22,6 +22,7 @@ import java.io.File
  import org.scalatest.Assertions
  
  import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
-+import org.apache.spark.sql.comet.CometHashAggregateExec
++import org.apache.spark.sql.comet.CometBaseAggregateExec
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
  import org.apache.spark.sql.execution.streaming.operators.stateful.{StateStoreRestoreExec, StateStoreSaveExec}
  import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
@@ -3728,7 +3733,7 @@ index 465da3cd469..92ac998929d 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = query.lastExecution.executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
@@ -3736,7 +3741,7 @@ index 465da3cd469..92ac998929d 100644
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = executedPlan.collect {
            case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
++          case c: CometBaseAggregateExec => c.originalPlan
          }
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index eff6269ba8..d9b5a1fff0 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
@@ -1832,12 +1832,12 @@ trait CometBaseAggregate {
    * (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: BaseAggregateExec
@@ -1946,12 +1946,11 @@ object CometSortAggregateExec
   }
 
   override def createExec(nativeOp: Operator, op: SortAggregateExec): CometNativeExec = {
-    // Reuse CometHashAggregateExec as the wrapper. The native AggregateExec auto-detects
-    // Sorted input mode from the child's output ordering and produces output sorted by the
-    // grouping keys; CometExec.outputOrdering defaults to originalPlan.outputOrdering, which
-    // is SortAggregateExec's grouping-key ordering, so downstream operators that elided a
-    // sort against it still see a satisfying ordering without a dedicated wrapper class.
-    CometHashAggregateExec(
+    // The native AggregateExec auto-detects Sorted input mode from the child's output ordering
+    // and produces output sorted by the grouping keys; CometExec.outputOrdering defaults to
+    // originalPlan.outputOrdering, which is SortAggregateExec's grouping-key ordering, so
+    // downstream operators that elided a sort against it still see a satisfying ordering.
+    CometSortAggregateExec(
       nativeOp,
       op,
       adjustOutputForNativeState(op),
@@ -1964,29 +1963,29 @@ object CometSortAggregateExec
   }
 }
 
-case class CometHashAggregateExec(
-    override val nativeOp: Operator,
-    override val originalPlan: SparkPlan,
-    override val output: Seq[Attribute],
-    groupingExpressions: Seq[NamedExpression],
-    aggregateExpressions: Seq[AggregateExpression],
-    resultExpressions: Seq[NamedExpression],
-    input: Seq[Attribute],
-    child: SparkPlan,
-    override val serializedPlanOpt: SerializedPlan)
+/**
+ * Common base for Comet's aggregate wrapper operators. The hash-based and sort-based variants
+ * share the same native AggregateExec serialization and rendering; they are kept as distinct plan
+ * node types only so the executed plan reflects whether Spark planned a HashAggregateExec /
+ * ObjectHashAggregateExec or a SortAggregateExec (the native AggregateExec auto-detects sorted
+ * input mode from the child ordering, so execution is otherwise identical).
+ */
+abstract class CometBaseAggregateExec
     extends CometUnaryExec
     with PartitioningPreservingUnaryExecNode {
 
+  def groupingExpressions: Seq[NamedExpression]
+  def aggregateExpressions: Seq[AggregateExpression]
+  def resultExpressions: Seq[NamedExpression]
+  def input: Seq[Attribute]
+
   // The aggExprs could be empty. For example, if the aggregate functions only have
   // distinct aggregate functions or only have group by, the aggExprs is empty and
   // modes is empty too.
-  val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct
+  lazy val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct
 
   override def producedAttributes: AttributeSet = outputSet ++ AttributeSet(resultExpressions)
 
-  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
-    this.copy(child = newChild)
-
   override def verboseStringWithOperatorId(): String = {
     s"""
        |$formattedNodeName
@@ -1999,6 +1998,24 @@ case class CometHashAggregateExec(
   override def stringArgs: Iterator[Any] =
     Iterator(input, modes, groupingExpressions, aggregateExpressions, child)
 
+  override protected def outputExpressions: Seq[NamedExpression] = resultExpressions
+}
+
+case class CometHashAggregateExec(
+    override val nativeOp: Operator,
+    override val originalPlan: SparkPlan,
+    override val output: Seq[Attribute],
+    groupingExpressions: Seq[NamedExpression],
+    aggregateExpressions: Seq[AggregateExpression],
+    resultExpressions: Seq[NamedExpression],
+    input: Seq[Attribute],
+    child: SparkPlan,
+    override val serializedPlanOpt: SerializedPlan)
+    extends CometBaseAggregateExec {
+
+  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+    this.copy(child = newChild)
+
   override def equals(obj: Any): Boolean = {
     obj match {
       case other: CometHashAggregateExec =>
@@ -2016,8 +2033,40 @@ case class CometHashAggregateExec(
 
   override def hashCode(): Int =
     Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child)
+}
 
-  override protected def outputExpressions: Seq[NamedExpression] = resultExpressions
+case class CometSortAggregateExec(
+    override val nativeOp: Operator,
+    override val originalPlan: SparkPlan,
+    override val output: Seq[Attribute],
+    groupingExpressions: Seq[NamedExpression],
+    aggregateExpressions: Seq[AggregateExpression],
+    resultExpressions: Seq[NamedExpression],
+    input: Seq[Attribute],
+    child: SparkPlan,
+    override val serializedPlanOpt: SerializedPlan)
+    extends CometBaseAggregateExec {
+
+  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+    this.copy(child = newChild)
+
+  override def equals(obj: Any): Boolean = {
+    obj match {
+      case other: CometSortAggregateExec =>
+        this.output == other.output &&
+        this.groupingExpressions == other.groupingExpressions &&
+        this.aggregateExpressions == other.aggregateExpressions &&
+        this.input == other.input &&
+        this.modes == other.modes &&
+        this.child == other.child &&
+        this.serializedPlanOpt == other.serializedPlanOpt
+      case _ =>
+        false
+    }
+  }
+
+  override def hashCode(): Int =
+    Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child)
 }
 
 trait CometHashJoin {
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql
index 8aedc4f5f1..18c7e1968f 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql
@@ -70,11 +70,16 @@ CREATE TABLE test_types(
   grp string
 ) USING parquet
 
+-- first/last IGNORE NULLS are non-deterministic when a group has more than one non-null value,
+-- because the result depends on intra-group row order, which is not defined and can differ
+-- between Spark and Comet. To keep the multi-type queries below comparable across engines, each
+-- group has at most one non-null value per column (surrounded by nulls so null-skipping is still
+-- exercised); group 'b' is entirely null.
 statement
 INSERT INTO test_types VALUES
   (NULL, NULL,  NULL,  NULL,  NULL,  'a'),
   (1,    100,   1.5,   'foo', true,  'a'),
-  (2,    200,   2.5,   'bar', false, 'a'),
+  (NULL, NULL,  NULL,  NULL,  NULL,  'a'),
   (NULL, NULL,  NULL,  NULL,  NULL,  'b'),
   (NULL, NULL,  NULL,  NULL,  NULL,  'b')
 
@@ -183,7 +188,9 @@ SELECT first(val) IGNORE NULLS FROM test_single_null
 -- first IGNORE NULLS: multiple data types
 -- ============================================================
 
-query expect_fallback(SortAggregate is not supported)
+-- Spark plans this as SortAggregateExec (the string buffer is not hash-friendly), which Comet
+-- now runs natively.
+query
 SELECT grp,
   first(i_val) IGNORE NULLS,
   first(l_val) IGNORE NULLS,
@@ -327,7 +334,9 @@ SELECT last(val) IGNORE NULLS FROM test_single_null
 -- last IGNORE NULLS: multiple data types
 -- ============================================================
 
-query expect_fallback(SortAggregate is not supported)
+-- Spark plans this as SortAggregateExec (the string buffer is not hash-friendly), which Comet
+-- now runs natively.
+query
 SELECT grp,
   last(i_val) IGNORE NULLS,
   last(l_val) IGNORE NULLS,
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql
index dbec88e2f6..5aacc9a10f 100644
--- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql
@@ -21,7 +21,10 @@ CREATE TABLE test_min_max(i int, d double, s string, grp string) USING parquet
 statement
 INSERT INTO test_min_max VALUES (1, 1.5, 'b', 'x'), (3, 3.5, 'a', 'x'), (2, 2.5, 'c', 'y'), (NULL, NULL, NULL, 'y'), (-1, -1.5, 'z', 'x')
 
-query expect_fallback(SortAggregate is not supported)
+-- min/max over StringType is not yet supported natively, so this falls back to Spark.
+-- (Spark plans this global aggregate as SortAggregateExec because the string min/max buffer
+-- is not hash-friendly; that operator is now supported, but the string aggregate itself is not.)
+query expect_fallback(Unsupported data type: StringType)
 SELECT min(i), max(i), min(d), max(d), min(s), max(s) FROM test_min_max
 
 query