From 834b5fb2d8bddf95143e0264ca4af6d26e5c762f Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Fri, 26 Jun 2026 00:11:14 +0800
Subject: [PATCH 01/10] Add native collect_list aggregate support

---
 .../expression-audits/agg_funcs.md            |  8 +++++
 docs/source/user-guide/latest/expressions.md  |  4 +--
 native/core/src/execution/planner.rs          |  7 +++-
 native/proto/src/proto/expr.proto             |  6 ++++
 .../serde/CometAggregateExpressionSerde.scala |  4 +--
 .../apache/comet/serde/QueryPlanSerde.scala   |  1 +
 .../org/apache/comet/serde/aggregates.scala   | 34 ++++++++++++++++++-
 .../apache/spark/sql/comet/operators.scala    | 24 ++++++-------
 .../comet/exec/CometAggregateSuite.scala      | 23 +++++++++++++
 9 files changed, 93 insertions(+), 18 deletions(-)

diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md
index 8025f31b73..116a31df4b 100644
--- a/docs/source/contributor-guide/expression-audits/agg_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md
@@ -39,4 +39,12 @@
 - Spark 3.5.8 (2026-05-26)
 - Spark 4.0.1 (2026-05-26)
 
+## collect_list
+
+- Spark 3.4.3 (audited 2026-06-24): `CollectList` extends `Collect[ArrayBuffer[Any]]`, returns `ArrayType(child.dataType, containsNull = false)`, ignores NULL inputs in `update()` (Hive-compatible semantics), and yields an empty array as `defaultResult`. `nullable = false`. No `checkInputDataTypes` override, so any input type accepted by Spark's analyzer is accepted. Registered as both `collect_list` and `array_agg` aliases in `FunctionRegistry`.
+- Spark 3.5.8 (audited 2026-06-24): identical to 3.4.3.
+- Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class; no behavior change.
+- Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1.
+- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array. Native intermediate state is `ArrayType(child.dataType, containsNull = true)`, so Comet rewrites intermediate `Partial`/`PartialMerge` object-aggregate buffer attributes from Spark's serialized `BinaryType` to that native state type.
+
 [Spark Expression Support]: ../../user-guide/latest/expressions.md
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index 134b231eda..707b1746ac 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -71,14 +71,14 @@ The tables below list every Spark built-in expression with its current status.
 | `any` | ✅ |  |
 | `any_value` | ✅ |  |
 | `approx_count_distinct` | 🔜 | tracking [#4098](https://github.com/apache/datafusion-comet/issues/4098) |
-| `array_agg` | 🔜 | Array aggregate (related to `collect_list`, [#2524](https://github.com/apache/datafusion-comet/issues/2524)) |
+| `array_agg` | ✅ | Alias for `collect_list` |
 | `avg` | ✅ | Interval types fall back |
 | `bit_and` | ✅ |  |
 | `bit_or` | ✅ |  |
 | `bit_xor` | ✅ |  |
 | `bool_and` | ✅ |  |
 | `bool_or` | ✅ |  |
-| `collect_list` | 🔜 | [#2524](https://github.com/apache/datafusion-comet/issues/2524) |
+| `collect_list` | ✅ |  |
 | `collect_set` | ✅ |  |
 | `corr` | ✅ |  |
 | `count` | ✅ |  |
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 553fe5215c..a226c59ff4 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -75,7 +75,7 @@ use datafusion_comet_spark_expr::{
     BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SparkArraysZipFunc,
     SparkBloomFilterVersion, SumInteger, ToCsv,
 };
-use datafusion_spark::function::aggregate::collect::SparkCollectSet;
+use datafusion_spark::function::aggregate::collect::{SparkCollectList, SparkCollectSet};
 use iceberg::expr::Bind;
 
 use crate::execution::operators::ExecutionError::GeneralError;
@@ -2627,6 +2627,11 @@ impl PhysicalPlanner {
                 let func = AggregateUDF::new_from_impl(SparkCollectSet::new());
                 Self::create_aggr_func_expr("collect_set", schema, vec![child], func)
             }
+            AggExprStruct::CollectList(expr) => {
+                let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?;
+                let func = AggregateUDF::new_from_impl(SparkCollectList::new());
+                Self::create_aggr_func_expr("collect_list", schema, vec![child], func)
+            }
         }
     }
 
diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto
index 90e3d87032..e73c07a914 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -143,6 +143,7 @@ message AggExpr {
     Correlation correlation = 15;
     BloomFilterAgg bloomFilterAgg = 16;
     CollectSet collectSet = 17;
+    CollectList collectList = 18;
   }
 
   // Optional filter expression for SQL FILTER (WHERE ...) clause.
@@ -267,6 +268,11 @@ message CollectSet {
   DataType datatype = 2;
 }
 
+message CollectList {
+  Expr child = 1;
+  DataType datatype = 2;
+}
+
 enum EvalMode {
   LEGACY = 0;
   TRY = 1;
diff --git a/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala
index 9a83152168..0c34e5c0b2 100644
--- a/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala
@@ -85,8 +85,8 @@ trait CometAggregateExpressionSerde[T <: AggregateFunction] {
    * Whether this aggregate's intermediate buffer format is compatible between Spark and Comet,
    * making it safe to run the Partial in one engine and the Final in the other. Aggregates with
    * simple single-value buffers (MIN, MAX, bitwise) are safe; those with complex or
-   * differently-encoded buffers (AVG, SUM with decimals, CollectSet, Variance) are not. COUNT is
-   * intentionally excluded: mixed COUNT partial/final regressed AQE's
+   * differently-encoded buffers (AVG, SUM with decimals, CollectList, CollectSet, Variance) are
+   * not. COUNT is intentionally excluded: mixed COUNT partial/final regressed AQE's
    * PropagateEmptyRelationAfterAQE pattern (which matches BaseAggregateExec only) and the Spark
    * 4.0 count-bug decorrelation for correlated IN subqueries.
    */
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 40d8e39dbf..b91b7db901 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -380,6 +380,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[BitOrAgg] -> CometBitOrAgg,
     classOf[BitXorAgg] -> CometBitXOrAgg,
     classOf[BloomFilterAggregate] -> CometBloomFilterAggregate,
+    classOf[CollectList] -> CometCollectList,
     classOf[CollectSet] -> CometCollectSet,
     classOf[Corr] -> CometCorr,
     classOf[Count] -> CometCount,
diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
index 6221452c86..8696b39ea1 100644
--- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
@@ -22,7 +22,7 @@ package org.apache.comet.serde
 import scala.jdk.CollectionConverters._
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Literal}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, Last, Max, Min, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, BitAndAgg, BitOrAgg, BitXorAgg, BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count, Covariance, CovPopulation, CovSample, First, Last, Max, Min, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{ByteType, DecimalType, IntegerType, LongType, ShortType, StringType}
 
@@ -740,6 +740,38 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] {
   }
 }
 
+object CometCollectList extends CometAggregateExpressionSerde[CollectList] {
+
+  override def convert(
+      aggExpr: AggregateExpression,
+      expr: CollectList,
+      inputs: Seq[Attribute],
+      binding: Boolean,
+      conf: SQLConf): Option[ExprOuterClass.AggExpr] = {
+    val child = expr.children.head
+    val childExpr = exprToProto(child, inputs, binding)
+    val dataType = serializeDataType(expr.dataType)
+
+    if (childExpr.isDefined && dataType.isDefined) {
+      val builder = ExprOuterClass.CollectList.newBuilder()
+      builder.setChild(childExpr.get)
+      builder.setDatatype(dataType.get)
+
+      Some(
+        ExprOuterClass.AggExpr
+          .newBuilder()
+          .setCollectList(builder)
+          .build())
+    } else if (dataType.isEmpty) {
+      withFallbackReason(aggExpr, s"datatype ${expr.dataType} is not supported", child)
+      None
+    } else {
+      withFallbackReason(aggExpr, child)
+      None
+    }
+  }
+}
+
 object AggSerde {
   import org.apache.spark.sql.types._
 
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 af9e1df8a3..60f8656a90 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
@@ -30,7 +30,7 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectSet, Final, First, Last, Partial, PartialMerge}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge}
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.physical._
@@ -1892,19 +1892,19 @@ object CometObjectHashAggregateExec
   }
 
   /**
-   * 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.
+   * For intermediate aggregates containing TypedImperativeAggregate functions (like CollectSet or
+   * CollectList), Spark declares buffer columns as BinaryType because it serializes the JVM
+   * state. Native Comet keeps the actual state type instead: ArrayType(elementType) with
+   * containsNull true for CollectSet/CollectList. Rewrite the Spark-side output attributes for
+   * Partial, PartialMerge, and mixed {Partial, PartialMerge} stages so shuffle and downstream
+   * native aggregate stages see the schema that native execution really produces.
    *
-   * NOTE: If a new TypedImperativeAggregate function (e.g., CollectList) is added natively, add a
-   * case branch here mapping it to the native state type.
+   * Final aggregates output user-visible values rather than intermediate state, so their Spark
+   * result schema is left unchanged.
    */
   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)) {
+    if (modes.exists(mode => mode != Partial && mode != PartialMerge)) {
       return op.output
     }
 
@@ -1916,8 +1916,8 @@ object CometObjectHashAggregateExec
       val aggFunc = aggExpr.aggregateFunction
       val bufferAttrs = aggFunc.aggBufferAttributes
       aggFunc match {
-        case cs: CollectSet =>
-          val elementType = cs.children.head.dataType
+        case _: CollectSet | _: CollectList =>
+          val elementType = aggFunc.children.head.dataType
           val nativeStateType = ArrayType(elementType, containsNull = true)
           output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType)
         case _ =>
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
index ae14c68207..7dc5338ac8 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -48,6 +48,29 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   override protected def sparkConf: SparkConf =
     super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false")
 
+  test("collect_list/collect_set combined with distinct aggregate runs fully native") {
+    // SPARK-17616: combining a distinct aggregate with collect_list/collect_set creates a
+    // multi-stage aggregate plan with a PartialMerge collect stage. These collect functions declare
+    // BinaryType JVM buffers in Spark but produce ArrayType native state in Comet; the intermediate
+    // Comet aggregate outputs must advertise the native ArrayType so that the state can round-trip
+    // through shuffle and downstream native PartialMerge/Final stages. See issue #4724.
+    withParquetTable(
+      Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")),
+      "t_collect_distinct",
+      withDictionary = false) {
+      for (fn <- Seq("collect_list", "collect_set")) {
+        val query =
+          s"SELECT _1, count(DISTINCT _2), sort_array($fn(_3)) " +
+            "FROM t_collect_distinct GROUP BY _1"
+        val (_, cometPlan) =
+          checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometHashAggregateExec]))
+        assert(
+          cometPlan.toString.contains(s"merge_$fn"),
+          s"Expected $fn PartialMerge stage to run as CometHashAggregateExec; plan:\n$cometPlan")
+      }
+    }
+  }
+
   test("min/max floating point with negative zero") {
     val r = new Random(42)
     val schema = StructType(

From ed83aa5f9233683d1bc5cba6cd480c0d8f74d1b1 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Sat, 27 Jun 2026 01:03:54 +0800
Subject: [PATCH 02/10] add test

---
 .../expressions/aggregate/collect_list.sql    | 87 +++++++++++++++++++
 .../expressions/aggregate/collect_set.sql     |  9 ++
 2 files changed, 96 insertions(+)
 create mode 100644 spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql

diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
new file mode 100644
index 0000000000..ff1155d53b
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql
@@ -0,0 +1,87 @@
+-- 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.
+
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+-- ============================================================
+-- Setup: tables
+-- ============================================================
+
+statement
+CREATE TABLE cl_src(k int, v string, grp string) USING parquet
+
+statement
+INSERT INTO cl_src VALUES
+  (1, 'b', 'g1'),
+  (1, 'a', 'g1'),
+  (2, 'b', 'g1'),
+  (NULL, NULL, 'g1'),
+  (3, 'c', 'g2'),
+  (3, 'c', 'g2'),
+  (4, 'd', 'g2'),
+  (NULL, 'z', 'g3')
+
+statement
+CREATE TABLE cl_empty(v string) USING parquet
+
+-- ============================================================
+-- Basic: global collect_list keeps duplicate non-null values
+-- ============================================================
+
+query
+SELECT sort_array(collect_list(v)) FROM cl_src
+
+-- ============================================================
+-- GROUP BY: collect_list per group
+-- ============================================================
+
+query
+SELECT grp, sort_array(collect_list(v)) FROM cl_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Alias: array_agg uses the collect_list implementation
+-- ============================================================
+
+query
+SELECT grp, sort_array(array_agg(k)) FROM cl_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Empty table: returns empty array
+-- ============================================================
+
+query
+SELECT sort_array(collect_list(v)) FROM cl_empty
+
+-- ============================================================
+-- PartialMerge: collect_list combined with a distinct aggregate
+-- exercises native intermediate ArrayType state through shuffle
+-- ============================================================
+
+query
+SELECT grp, count(DISTINCT k), sort_array(collect_list(v))
+FROM cl_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- PartialMerge: multiple collect-style aggregates with distinct
+-- ============================================================
+
+query
+SELECT grp,
+       count(DISTINCT k),
+       sort_array(collect_list(v)),
+       sort_array(collect_set(v))
+FROM cl_src GROUP BY grp ORDER BY grp
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 cd528272af..e1b30d0d3f 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
@@ -291,6 +291,15 @@ FROM cs_src_multi GROUP BY grp ORDER BY grp
 query
 SELECT grp, sort_array(collect_set(DISTINCT i)) FROM cs_src_int GROUP BY grp ORDER BY grp
 
+-- ============================================================
+-- PartialMerge: collect_set combined with a distinct aggregate
+-- exercises native intermediate ArrayType state through shuffle
+-- ============================================================
+
+query
+SELECT grp, count(DISTINCT a), sort_array(collect_set(b))
+FROM cs_src_multi GROUP BY grp ORDER BY grp
+
 -- ============================================================
 -- HAVING clause with collect_set
 -- ============================================================

From 2ee3a90f5aa7f549bd4df6045dcdc9e68a5f6b25 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Sat, 27 Jun 2026 01:05:02 +0800
Subject: [PATCH 03/10] Constructs inner collect_list/collect_set accumulators
 with the original element

---
 native/core/src/execution/merge_as_partial.rs | 57 +++++++++++++++----
 1 file changed, 47 insertions(+), 10 deletions(-)

diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs
index 5ea26115bf..de4502b39f 100644
--- a/native/core/src/execution/merge_as_partial.rs
+++ b/native/core/src/execution/merge_as_partial.rs
@@ -30,7 +30,7 @@ use std::fmt::Debug;
 use std::hash::{Hash, Hasher};
 
 use arrow::array::{ArrayRef, BooleanArray};
-use arrow::datatypes::{DataType, FieldRef};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion::common::Result;
 use datafusion::logical_expr::function::AccumulatorArgs;
 use datafusion::logical_expr::function::StateFieldsArgs;
@@ -60,6 +60,8 @@ pub struct MergeAsPartialUDF {
     cached_state_fields: Vec<FieldRef>,
     /// Cached signature that accepts state field types.
     signature: Signature,
+    /// Override argument fields used only when constructing the inner accumulator.
+    accumulator_expr_fields: Option<Vec<FieldRef>>,
     /// Name for this wrapper.
     name: String,
 }
@@ -83,6 +85,8 @@ impl MergeAsPartialUDF {
         let name = format!("merge_as_partial_{}", inner_expr.name());
         let return_type = inner_expr.field().data_type().clone();
         let cached_state_fields = inner_expr.state_fields()?;
+        let accumulator_expr_fields =
+            Self::accumulator_expr_fields(inner_expr, &cached_state_fields);
 
         // Use a permissive signature since we accept state field types which
         // vary per aggregate function.
@@ -93,9 +97,43 @@ impl MergeAsPartialUDF {
             return_type,
             cached_state_fields,
             signature,
+            accumulator_expr_fields,
             name,
         })
     }
+
+    fn accumulator_expr_fields(
+        inner_expr: &AggregateFunctionExpr,
+        state_fields: &[FieldRef],
+    ) -> Option<Vec<FieldRef>> {
+        // Spark collect_list/collect_set use a single list state field but their accumulator
+        // constructors expect the original scalar input type. MergeAsPartial receives state-typed
+        // input columns, so passing those fields through would make the collect accumulator think
+        // it is aggregating arrays and create nested list state (List(List(T))).
+        match (inner_expr.fun().name(), state_fields) {
+            ("collect_list" | "collect_set", [state_field]) => match state_field.data_type() {
+                DataType::List(item_field) => Some(vec![Field::new(
+                    item_field.name(),
+                    item_field.data_type().clone(),
+                    item_field.is_nullable(),
+                )
+                .into()]),
+                _ => None,
+            },
+            _ => None,
+        }
+    }
+
+    fn inner_accumulator_args<'a>(&'a self, args: AccumulatorArgs<'a>) -> AccumulatorArgs<'a> {
+        let expr_fields = self
+            .accumulator_expr_fields
+            .as_deref()
+            .unwrap_or(args.expr_fields);
+        AccumulatorArgs {
+            expr_fields,
+            ..args
+        }
+    }
 }
 
 impl AggregateUDFImpl for MergeAsPartialUDF {
@@ -120,25 +158,24 @@ impl AggregateUDFImpl for MergeAsPartialUDF {
     }
 
     fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
-        // args.exprs are state-typed (match this wrapper's signature), not the
-        // inner aggregate's original inputs. Safe for built-ins (SUM/COUNT/
-        // MIN/MAX/AVG) which build accumulators from return_type; aggregates
-        // that inspect args.exprs types would need reconsideration.
-        let inner_acc = self.inner_udf.accumulator(args)?;
+        let inner_acc = self
+            .inner_udf
+            .accumulator(self.inner_accumulator_args(args))?;
         Ok(Box::new(MergeAsPartialAccumulator { inner: inner_acc }))
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
-        // See `accumulator`: args.exprs are state-typed.
-        self.inner_udf.groups_accumulator_supported(args)
+        self.inner_udf
+            .groups_accumulator_supported(self.inner_accumulator_args(args))
     }
 
     fn create_groups_accumulator(
         &self,
         args: AccumulatorArgs,
     ) -> Result<Box<dyn GroupsAccumulator>> {
-        // See `accumulator`: args.exprs are state-typed.
-        let inner_acc = self.inner_udf.create_groups_accumulator(args)?;
+        let inner_acc = self
+            .inner_udf
+            .create_groups_accumulator(self.inner_accumulator_args(args))?;
         Ok(Box::new(MergeAsPartialGroupsAccumulator {
             inner: inner_acc,
         }))

From dbf7ba50ce3e6ff8f21916d9122611c136b63421 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Sat, 27 Jun 2026 19:02:47 +0800
Subject: [PATCH 04/10] Fixes Spark test plan-shape assertions for native
 `collect_list`.

---
 dev/diffs/3.4.3.diff | 26 +++++++++++++++++++++++---
 dev/diffs/3.5.8.diff | 26 +++++++++++++++++++++++---
 dev/diffs/4.0.2.diff | 26 +++++++++++++++++++++++---
 dev/diffs/4.1.2.diff | 26 +++++++++++++++++++++++---
 4 files changed, 92 insertions(+), 12 deletions(-)

diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff
index bfbdca7545..48bdf69715 100644
--- a/dev/diffs/3.4.3.diff
+++ b/dev/diffs/3.4.3.diff
@@ -260,10 +260,14 @@ index cf40e944c09..bdd5be4f462 100644
  
    test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-index 1cc09c3d7fc..f031fa45c33 100644
+index 1cc09c3d7fc..3dd8d8fbcc7 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-@@ -27,7 +27,7 @@ import org.apache.spark.SparkException
+@@ -24,10 +24,11 @@ import scala.util.Random
+ import org.scalatest.matchers.must.Matchers.the
+ 
+ import org.apache.spark.SparkException
++import org.apache.spark.sql.comet.CometHashAggregateExec
  import org.apache.spark.sql.execution.WholeStageCodegenExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
@@ -272,7 +276,23 @@ index 1cc09c3d7fc..f031fa45c33 100644
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.internal.SQLConf
-@@ -755,7 +755,7 @@ class DataFrameAggregateSuite extends QueryTest
+@@ -700,7 +701,8 @@ class DataFrameAggregateSuite extends QueryTest
+         val objHashAggOrSortAggPlan =
+           stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan)
+         if (useObjectHashAgg) {
+-          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec])
++          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] ||
++            objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec])
+         } else {
+           assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec])
+         }
+@@ -750,12 +752,12 @@ class DataFrameAggregateSuite extends QueryTest
+       assert(sortAggPlans.isEmpty)
+ 
+       val objHashAggPlans = collect(aggPlan) {
+-        case objHashAgg: ObjectHashAggregateExec => objHashAgg
++        case agg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => agg
+       }
        assert(objHashAggPlans.nonEmpty)
  
        val exchangePlans = collect(aggPlan) {
diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff
index ae13b120da..8e6a04091f 100644
--- a/dev/diffs/3.5.8.diff
+++ b/dev/diffs/3.5.8.diff
@@ -241,10 +241,14 @@ index e5494726695..00937f025c2 100644
  
    test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-index 6f3090d8908..c08a60fb0c2 100644
+index 6f3090d8908..e6681f4d8e4 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.plans.logical.Expand
+@@ -25,10 +25,11 @@ import org.scalatest.matchers.must.Matchers.the
+ 
+ import org.apache.spark.{SparkException, SparkThrowable}
+ import org.apache.spark.sql.catalyst.plans.logical.Expand
++import org.apache.spark.sql.comet.CometHashAggregateExec
  import org.apache.spark.sql.execution.WholeStageCodegenExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
@@ -253,7 +257,23 @@ index 6f3090d8908..c08a60fb0c2 100644
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.internal.SQLConf
-@@ -793,7 +793,7 @@ class DataFrameAggregateSuite extends QueryTest
+@@ -738,7 +739,8 @@ class DataFrameAggregateSuite extends QueryTest
+         val objHashAggOrSortAggPlan =
+           stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan)
+         if (useObjectHashAgg) {
+-          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec])
++          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] ||
++            objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec])
+         } else {
+           assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec])
+         }
+@@ -788,12 +790,12 @@ class DataFrameAggregateSuite extends QueryTest
+       assert(sortAggPlans.isEmpty)
+ 
+       val objHashAggPlans = collect(aggPlan) {
+-        case objHashAgg: ObjectHashAggregateExec => objHashAgg
++        case agg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => agg
+       }
        assert(objHashAggPlans.nonEmpty)
  
        val exchangePlans = collect(aggPlan) {
diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff
index c4eeb50309..87e83e6f35 100644
--- a/dev/diffs/4.0.2.diff
+++ b/dev/diffs/4.0.2.diff
@@ -378,10 +378,14 @@ index 0f42502f1d9..e9ff802141f 100644
  
        withTempView("t0", "t1", "t2") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-index 9db406ff12f..245e4caa319 100644
+index 9db406ff12f..87a5a95e72d 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-@@ -30,7 +30,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
+@@ -27,10 +27,11 @@ import org.apache.spark.{SparkArithmeticException, SparkRuntimeException}
+ import org.apache.spark.sql.catalyst.plans.logical.Expand
+ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS
+ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
++import org.apache.spark.sql.comet.CometHashAggregateExec
  import org.apache.spark.sql.execution.WholeStageCodegenExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
@@ -390,7 +394,23 @@ index 9db406ff12f..245e4caa319 100644
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.internal.SQLConf
-@@ -855,7 +855,7 @@ class DataFrameAggregateSuite extends QueryTest
+@@ -800,7 +801,8 @@ class DataFrameAggregateSuite extends QueryTest
+         val objHashAggOrSortAggPlan =
+           stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan)
+         if (useObjectHashAgg) {
+-          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec])
++          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] ||
++            objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec])
+         } else {
+           assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec])
+         }
+@@ -850,12 +852,12 @@ class DataFrameAggregateSuite extends QueryTest
+       assert(sortAggPlans.isEmpty)
+ 
+       val objHashAggPlans = collect(aggPlan) {
+-        case objHashAgg: ObjectHashAggregateExec => objHashAgg
++        case agg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => agg
+       }
        assert(objHashAggPlans.nonEmpty)
  
        val exchangePlans = collect(aggPlan) {
diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff
index f54a185148..1b737e2f3a 100644
--- a/dev/diffs/4.1.2.diff
+++ b/dev/diffs/4.1.2.diff
@@ -392,10 +392,14 @@ index 0d807aeae4d..6d7744e771b 100644
  
        withTempView("t0", "t1", "t2") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-index bfe15b33768..55c23a38ccc 100644
+index bfe15b33768..8b1a1c08f20 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-@@ -31,7 +31,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
+@@ -28,10 +28,11 @@ import org.apache.spark.{SparkArithmeticException, SparkRuntimeException}
+ import org.apache.spark.sql.catalyst.plans.logical.Expand
+ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS
+ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
++import org.apache.spark.sql.comet.CometHashAggregateExec
  import org.apache.spark.sql.execution.WholeStageCodegenExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
@@ -404,7 +408,23 @@ index bfe15b33768..55c23a38ccc 100644
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.internal.SQLConf
-@@ -856,7 +856,7 @@ class DataFrameAggregateSuite extends QueryTest
+@@ -801,7 +802,8 @@ class DataFrameAggregateSuite extends QueryTest
+         val objHashAggOrSortAggPlan =
+           stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan)
+         if (useObjectHashAgg) {
+-          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec])
++          assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] ||
++            objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec])
+         } else {
+           assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec])
+         }
+@@ -851,12 +853,12 @@ class DataFrameAggregateSuite extends QueryTest
+       assert(sortAggPlans.isEmpty)
+ 
+       val objHashAggPlans = collect(aggPlan) {
+-        case objHashAgg: ObjectHashAggregateExec => objHashAgg
++        case agg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => agg
+       }
        assert(objHashAggPlans.nonEmpty)
  
        val exchangePlans = collect(aggPlan) {

From 9af91562913ec8b6e6b99cf27bce88c846601668 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Sat, 27 Jun 2026 20:42:18 +0800
Subject: [PATCH 05/10] fix preitter?

---
 docs/source/user-guide/latest/installation.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md
index 442133facc..3bd1ec27ee 100644
--- a/docs/source/user-guide/latest/installation.md
+++ b/docs/source/user-guide/latest/installation.md
@@ -92,6 +92,7 @@ Here are the direct links for downloading the Comet $COMET_VERSION jar file.
 - [Comet plugin for Spark 3.5 / Scala 2.13](https://repo1.maven.org/maven2/org/apache/datafusion/comet-spark-spark3.5_2.13/$COMET_VERSION/comet-spark-spark3.5_2.13-$COMET_VERSION.jar)
 - [Comet plugin for Spark 4.0 / Scala 2.13](https://repo1.maven.org/maven2/org/apache/datafusion/comet-spark-spark4.0_2.13/$COMET_VERSION/comet-spark-spark4.0_2.13-$COMET_VERSION.jar)
 - [Comet plugin for Spark 4.1 / Scala 2.13](https://repo1.maven.org/maven2/org/apache/datafusion/comet-spark-spark4.1_2.13/$COMET_VERSION/comet-spark-spark4.1_2.13-$COMET_VERSION.jar)
+
 <!-- ENDIF -->
 
 ## Building from source

From 0d10df1f632ef039ee40904c90f44c65b2108093 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Mon, 29 Jun 2026 02:15:42 +0800
Subject: [PATCH 06/10] fallback collect partial merge when child is spark
 buffer

---
 .../apache/spark/sql/comet/operators.scala    | 59 ++++++++++++++++++-
 .../comet/exec/CometAggregateSuite.scala      | 31 +++++++++-
 2 files changed, 86 insertions(+), 4 deletions(-)

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 60f8656a90..e52e9ff693 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
@@ -43,7 +43,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat
 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}
-import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, TimestampNTZType, TimestampType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.SerializableConfiguration
 import org.apache.spark.util.io.ChunkedByteBuffer
@@ -1523,13 +1523,18 @@ trait CometBaseAggregate {
     val modes = aggregate.aggregateExpressions.map(_.mode).distinct
     val modeSet = modes.toSet
     val hasPartialMerge = modeSet.contains(PartialMerge)
+    val cometPartialAgg = findCometPartialAgg(aggregate.child)
     // In distinct aggregates there can be a combination of modes.
     // We support {Partial, PartialMerge} mix; other combinations are rejected.
     val multiMode = modes.size > 1 && modeSet != Set(Partial, PartialMerge)
     // For a final mode HashAggregate, we only need to transform the HashAggregate
     // if there is Comet partial aggregation, unless all aggregates have compatible
     // intermediate buffer formats (safe for mixed Spark/Comet execution).
-    val sparkFinalMode = modes.contains(Final) && findCometPartialAgg(aggregate.child).isEmpty
+    val sparkFinalMode = modes.contains(Final) && cometPartialAgg.isEmpty
+    // For PartialMerge, the child aggregate may still be Spark/JVM even when this node's direct
+    // child is a Comet shuffle. In that case Spark's intermediate buffers must be compatible with
+    // the native merge accumulator state expected by Comet.
+    val sparkPartialMergeMode = hasPartialMerge && cometPartialAgg.isEmpty
 
     if (multiMode) {
       withFallbackReason(
@@ -1547,6 +1552,24 @@ trait CometBaseAggregate {
       return None
     }
 
+    if (partialMergeCollectHasNonNativeState(aggregate)) {
+      withFallbackReason(
+        aggregate,
+        "PartialMerge collect aggregate requires native array intermediate state")
+      return None
+    }
+
+    if (sparkPartialMergeMode) {
+      val partialMergeAggs = aggregate.aggregateExpressions.filter(_.mode == PartialMerge)
+      if (!QueryPlanSerde.allAggsSupportMixedExecution(partialMergeAggs)) {
+        withFallbackReason(
+          aggregate,
+          "Spark PartialMerge aggregate without Comet Partial requires compatible " +
+            "intermediate buffer formats")
+        return None
+      }
+    }
+
     // Check if this aggregate has been tagged as unsafe for mixed execution
     // (Comet partial + Spark final with incompatible intermediate buffers)
     val unsafeReason = aggregate.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL)
@@ -1775,6 +1798,38 @@ trait CometBaseAggregate {
         .build())
   }
 
+  /**
+   * Comet's collect_list/collect_set partial state is an ArrayType, while Spark's JVM
+   * TypedImperativeAggregate state is BinaryType. For PartialMerge, the native planner binds
+   * state input columns directly from the child schema using initialInputBufferOffset, so do not
+   * convert a collect PartialMerge when the incoming state is still Spark's binary buffer.
+   */
+  private def partialMergeCollectHasNonNativeState(aggregate: BaseAggregateExec): Boolean = {
+    var stateOffset = aggregate.initialInputBufferOffset
+
+    aggregate.aggregateExpressions.exists { aggExpr =>
+      val bufferAttrs = aggExpr.aggregateFunction.aggBufferAttributes
+      val hasNonNativeCollectState = aggExpr.mode == PartialMerge &&
+        (aggExpr.aggregateFunction.isInstanceOf[CollectList] ||
+          aggExpr.aggregateFunction.isInstanceOf[CollectSet]) &&
+        bufferAttrs.indices.exists { i =>
+          val inputIdx = stateOffset + i
+          inputIdx >= aggregate.child.output.length ||
+          (aggregate.child.output(inputIdx).dataType match {
+            case _: ArrayType => false
+            case BinaryType => true
+            case _ => true
+          })
+        }
+
+      if (aggExpr.mode == PartialMerge) {
+        stateOffset += bufferAttrs.length
+      }
+
+      hasNonNativeCollectState
+    }
+  }
+
   /**
    * 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.
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 7dc5338ac8..06cf9afef3 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -23,12 +23,12 @@ import scala.util.Random
 
 import org.apache.hadoop.fs.Path
 import org.apache.spark.SparkConf
-import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
 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.functions.{avg, col, count_distinct, sum}
+import org.apache.spark.sql.functions.{avg, col, collect_list, collect_set, count_distinct, sort_array, sum}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
 
@@ -71,6 +71,33 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test(
+    "collect_list/collect_set with distinct aggregate falls back on Spark PartialMerge state") {
+    // This mirrors Spark's DataFrameAggregateSuite SPARK-17616 test. With LocalTableScan disabled,
+    // the lower collect partial aggregate stays on Spark and produces a BinaryType JVM buffer.
+    // Native PartialMerge must not try to merge that buffer as DataFusion's list-typed state. Run
+    // the same hash-map configurations used by the inherited Spark aggregate suites that failed in
+    // CI.
+    for ((twoLevelAggMap, vectorizedHashMap) <- Seq(
+        (false, false),
+        (true, false),
+        (true, true))) {
+      withSQLConf(
+        CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "false",
+        SQLConf.CODEGEN_FALLBACK.key -> "false",
+        SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelAggMap.toString,
+        SQLConf.ENABLE_VECTORIZED_HASH_MAP.key -> vectorizedHashMap.toString) {
+        val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d"))
+          .toDF("x", "y", "z")
+        for (collect <- Seq(collect_list(_: Column), collect_set(_: Column))) {
+          checkSparkAnswerAndFallbackReason(
+            df.groupBy($"x").agg(count_distinct($"y"), sort_array(collect($"z"))),
+            "PartialMerge collect aggregate requires native array intermediate state")
+        }
+      }
+    }
+  }
+
   test("min/max floating point with negative zero") {
     val r = new Random(42)
     val schema = StructType(

From 41a2a80d87445a86cc731ba2f42d48a466cb8394 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Mon, 29 Jun 2026 04:42:26 +0800
Subject: [PATCH 07/10] Decode Spark collect partial merge state

---
 native/core/src/execution/merge_as_partial.rs |  38 +-
 native/core/src/execution/mod.rs              |   1 +
 .../src/execution/spark_aggregate_state.rs    | 428 ++++++++++++++++++
 .../apache/spark/sql/comet/operators.scala    |  48 +-
 .../comet/exec/CometAggregateSuite.scala      |  20 +-
 5 files changed, 479 insertions(+), 56 deletions(-)
 create mode 100644 native/core/src/execution/spark_aggregate_state.rs

diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs
index de4502b39f..25471d3acf 100644
--- a/native/core/src/execution/merge_as_partial.rs
+++ b/native/core/src/execution/merge_as_partial.rs
@@ -41,6 +41,8 @@ use datafusion::logical_expr::{
 use datafusion::physical_expr::aggregate::AggregateFunctionExpr;
 use datafusion::scalar::ScalarValue;
 
+use crate::execution::spark_aggregate_state::PartialMergeStateDecoder;
+
 /// An AggregateUDF wrapper that gives merge semantics in Partial mode.
 ///
 /// When DataFusion runs an AggregateExec in Partial mode, it calls `update_batch`
@@ -62,6 +64,8 @@ pub struct MergeAsPartialUDF {
     signature: Signature,
     /// Override argument fields used only when constructing the inner accumulator.
     accumulator_expr_fields: Option<Vec<FieldRef>>,
+    /// Decoder for Spark JVM state, or pass-through for native-compatible state.
+    state_decoder: PartialMergeStateDecoder,
     /// Name for this wrapper.
     name: String,
 }
@@ -87,6 +91,7 @@ impl MergeAsPartialUDF {
         let cached_state_fields = inner_expr.state_fields()?;
         let accumulator_expr_fields =
             Self::accumulator_expr_fields(inner_expr, &cached_state_fields);
+        let state_decoder = PartialMergeStateDecoder::new(inner_expr, &cached_state_fields);
 
         // Use a permissive signature since we accept state field types which
         // vary per aggregate function.
@@ -98,6 +103,7 @@ impl MergeAsPartialUDF {
             cached_state_fields,
             signature,
             accumulator_expr_fields,
+            state_decoder,
             name,
         })
     }
@@ -161,7 +167,10 @@ impl AggregateUDFImpl for MergeAsPartialUDF {
         let inner_acc = self
             .inner_udf
             .accumulator(self.inner_accumulator_args(args))?;
-        Ok(Box::new(MergeAsPartialAccumulator { inner: inner_acc }))
+        Ok(Box::new(MergeAsPartialAccumulator {
+            inner: inner_acc,
+            state_decoder: self.state_decoder.clone(),
+        }))
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
@@ -178,6 +187,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF {
             .create_groups_accumulator(self.inner_accumulator_args(args))?;
         Ok(Box::new(MergeAsPartialGroupsAccumulator {
             inner: inner_acc,
+            state_decoder: self.state_decoder.clone(),
         }))
     }
 
@@ -197,6 +207,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF {
 /// Accumulator wrapper that redirects update_batch to merge_batch.
 struct MergeAsPartialAccumulator {
     inner: Box<dyn Accumulator>,
+    state_decoder: PartialMergeStateDecoder,
 }
 
 impl Debug for MergeAsPartialAccumulator {
@@ -208,11 +219,13 @@ impl Debug for MergeAsPartialAccumulator {
 impl Accumulator for MergeAsPartialAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         // Redirect update to merge — this is the key trick.
-        self.inner.merge_batch(values)
+        let decoded = self.state_decoder.decode(values)?;
+        self.inner.merge_batch(decoded.arrays())
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
-        self.inner.merge_batch(states)
+        let decoded = self.state_decoder.decode(states)?;
+        self.inner.merge_batch(decoded.arrays())
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
@@ -231,6 +244,7 @@ impl Accumulator for MergeAsPartialAccumulator {
 /// GroupsAccumulator wrapper that redirects update_batch to merge_batch.
 struct MergeAsPartialGroupsAccumulator {
     inner: Box<dyn GroupsAccumulator>,
+    state_decoder: PartialMergeStateDecoder,
 }
 
 impl Debug for MergeAsPartialGroupsAccumulator {
@@ -248,8 +262,13 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator {
         total_num_groups: usize,
     ) -> Result<()> {
         // Redirect update to merge — this is the key trick.
-        self.inner
-            .merge_batch(values, group_indices, opt_filter, total_num_groups)
+        let decoded = self.state_decoder.decode(values)?;
+        self.inner.merge_batch(
+            decoded.arrays(),
+            group_indices,
+            opt_filter,
+            total_num_groups,
+        )
     }
 
     fn merge_batch(
@@ -259,8 +278,13 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator {
         opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
-        self.inner
-            .merge_batch(values, group_indices, opt_filter, total_num_groups)
+        let decoded = self.state_decoder.decode(values)?;
+        self.inner.merge_batch(
+            decoded.arrays(),
+            group_indices,
+            opt_filter,
+            total_num_groups,
+        )
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs
index ec247f72b7..5145843970 100644
--- a/native/core/src/execution/mod.rs
+++ b/native/core/src/execution/mod.rs
@@ -26,6 +26,7 @@ pub(crate) mod planner;
 pub mod serde;
 pub use datafusion_comet_shuffle as shuffle;
 pub(crate) mod sort;
+pub(crate) mod spark_aggregate_state;
 pub(crate) mod spark_plan;
 pub use datafusion_comet_spark_expr::timezone;
 mod memory_pools;
diff --git a/native/core/src/execution/spark_aggregate_state.rs b/native/core/src/execution/spark_aggregate_state.rs
new file mode 100644
index 0000000000..217a6785e5
--- /dev/null
+++ b/native/core/src/execution/spark_aggregate_state.rs
@@ -0,0 +1,428 @@
+// 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.
+
+//! Decoders for Spark JVM aggregate state consumed by native PartialMerge.
+
+use std::sync::Arc;
+
+use arrow::array::{
+    builder::{make_builder, ListBuilder},
+    Array, ArrayRef, BinaryArray, LargeBinaryArray,
+};
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::{DataFusionError, Result};
+use datafusion::physical_expr::aggregate::AggregateFunctionExpr;
+use datafusion_comet_shuffle::spark_unsafe::list::{append_to_builder, SparkUnsafeArray};
+
+/// State arrays after optional Spark-JVM-state decoding.
+pub(crate) enum DecodedState<'a> {
+    Borrowed(&'a [ArrayRef]),
+    Owned(Vec<ArrayRef>),
+}
+
+impl DecodedState<'_> {
+    pub(crate) fn arrays(&self) -> &[ArrayRef] {
+        match self {
+            Self::Borrowed(values) => values,
+            Self::Owned(values) => values,
+        }
+    }
+}
+
+/// Decoder used by `MergeAsPartial` before forwarding state to the inner accumulator.
+#[derive(Clone, Debug)]
+pub(crate) enum PartialMergeStateDecoder {
+    PassThrough,
+    SparkCollect(SparkCollectStateDecoder),
+}
+
+impl PartialMergeStateDecoder {
+    pub(crate) fn new(inner_expr: &AggregateFunctionExpr, state_fields: &[FieldRef]) -> Self {
+        SparkCollectStateDecoder::try_new(inner_expr, state_fields)
+            .map(Self::SparkCollect)
+            .unwrap_or(Self::PassThrough)
+    }
+
+    pub(crate) fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<DecodedState<'a>> {
+        match self {
+            Self::PassThrough => Ok(DecodedState::Borrowed(values)),
+            Self::SparkCollect(decoder) => decoder.decode(values),
+        }
+    }
+}
+
+/// Decodes Spark JVM collect aggregate buffers into DataFusion collect state.
+///
+/// Spark's `CollectList` / `CollectSet` are `TypedImperativeAggregate`s. When Spark runs the
+/// lower Partial aggregate, each buffer is serialized as a `BinaryType` value containing a
+/// single-field `UnsafeRow`; field 0 is the `UnsafeArrayData` with the collected elements.
+/// DataFusion's collect accumulators expect the merge input to be a list-typed state column, so
+/// mixed Spark-Partial -> Comet-PartialMerge plans must materialize those unsafe bytes into an
+/// Arrow `ListArray` before calling the inner accumulator's `merge_batch`.
+#[derive(Clone, Debug)]
+pub(crate) struct SparkCollectStateDecoder {
+    state_field: FieldRef,
+}
+
+impl SparkCollectStateDecoder {
+    fn try_new(inner_expr: &AggregateFunctionExpr, state_fields: &[FieldRef]) -> Option<Self> {
+        match (inner_expr.fun().name(), state_fields) {
+            ("collect_list" | "collect_set", [state_field])
+                if matches!(state_field.data_type(), DataType::List(_)) =>
+            {
+                Some(Self {
+                    state_field: Arc::clone(state_field),
+                })
+            }
+            _ => None,
+        }
+    }
+
+    fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<DecodedState<'a>> {
+        if values.len() != 1 {
+            return Err(DataFusionError::Internal(format!(
+                "Spark collect state decoder expected one state column, got {}",
+                values.len()
+            )));
+        }
+
+        match values[0].data_type() {
+            DataType::Binary => {
+                let array = values[0]
+                    .as_any()
+                    .downcast_ref::<BinaryArray>()
+                    .ok_or_else(|| {
+                        Self::decode_error("expected BinaryArray for Binary collect state")
+                    })?;
+                Ok(DecodedState::Owned(vec![self.decode_binary_array(array)?]))
+            }
+            DataType::LargeBinary => {
+                let array = values[0]
+                    .as_any()
+                    .downcast_ref::<LargeBinaryArray>()
+                    .ok_or_else(|| {
+                        Self::decode_error(
+                            "expected LargeBinaryArray for LargeBinary collect state",
+                        )
+                    })?;
+                Ok(DecodedState::Owned(vec![
+                    self.decode_large_binary_array(array)?
+                ]))
+            }
+            _ => Ok(DecodedState::Borrowed(values)),
+        }
+    }
+
+    fn decode_binary_array(&self, array: &BinaryArray) -> Result<ArrayRef> {
+        let item_field = self.item_field()?;
+        let mut builder = self.new_list_builder(item_field, array.len());
+
+        for row_idx in 0..array.len() {
+            if array.is_null(row_idx) {
+                builder.append_null();
+            } else {
+                self.append_unsafe_row_array(array.value(row_idx), item_field, &mut builder)?;
+            }
+        }
+
+        Ok(Arc::new(builder.finish()))
+    }
+
+    fn decode_large_binary_array(&self, array: &LargeBinaryArray) -> Result<ArrayRef> {
+        let item_field = self.item_field()?;
+        let mut builder = self.new_list_builder(item_field, array.len());
+
+        for row_idx in 0..array.len() {
+            if array.is_null(row_idx) {
+                builder.append_null();
+            } else {
+                self.append_unsafe_row_array(array.value(row_idx), item_field, &mut builder)?;
+            }
+        }
+
+        Ok(Arc::new(builder.finish()))
+    }
+
+    fn item_field(&self) -> Result<&FieldRef> {
+        match self.state_field.data_type() {
+            DataType::List(item_field) => Ok(item_field),
+            other => Err(Self::decode_error(format!(
+                "collect state field must be List, got {other:?}"
+            ))),
+        }
+    }
+
+    fn new_list_builder(
+        &self,
+        item_field: &FieldRef,
+        capacity: usize,
+    ) -> ListBuilder<Box<dyn arrow::array::builder::ArrayBuilder>> {
+        let value_builder = make_builder(item_field.data_type(), capacity);
+        ListBuilder::with_capacity(value_builder, capacity).with_field(Arc::clone(item_field))
+    }
+
+    fn append_unsafe_row_array(
+        &self,
+        row_bytes: &[u8],
+        item_field: &FieldRef,
+        builder: &mut ListBuilder<Box<dyn arrow::array::builder::ArrayBuilder>>,
+    ) -> Result<()> {
+        match Self::spark_array_from_single_field_unsafe_row(row_bytes)? {
+            Some(array) => {
+                append_to_builder::<true>(item_field.data_type(), builder.values(), &array)
+                    .map_err(|e| Self::decode_error(e.to_string()))?;
+                builder.append(true);
+            }
+            None => builder.append_null(),
+        }
+        Ok(())
+    }
+
+    fn spark_array_from_single_field_unsafe_row(
+        row_bytes: &[u8],
+    ) -> Result<Option<SparkUnsafeArray>> {
+        const BITSET_WIDTH: usize = 8;
+        const FIXED_FIELD_WIDTH: usize = 8;
+        const ARRAY_FIELD_INDEX: usize = 0;
+        const MIN_ROW_WIDTH: usize = BITSET_WIDTH + FIXED_FIELD_WIDTH;
+
+        if row_bytes.len() < MIN_ROW_WIDTH {
+            return Err(Self::decode_error(format!(
+                "UnsafeRow collect buffer is too small: {} bytes",
+                row_bytes.len()
+            )));
+        }
+
+        let null_bits = i64::from_le_bytes(
+            row_bytes[0..BITSET_WIDTH]
+                .try_into()
+                .expect("slice length checked"),
+        );
+        if (null_bits & (1_i64 << ARRAY_FIELD_INDEX)) != 0 {
+            return Ok(None);
+        }
+
+        let offset_and_size = i64::from_le_bytes(
+            row_bytes[BITSET_WIDTH..MIN_ROW_WIDTH]
+                .try_into()
+                .expect("slice length checked"),
+        );
+        let offset = (offset_and_size >> 32) as i32;
+        let size = offset_and_size as i32;
+
+        if offset < MIN_ROW_WIDTH as i32 || size < 0 {
+            return Err(Self::decode_error(format!(
+                "Invalid UnsafeRow array field offset/size: offset={offset}, size={size}, row_size={}",
+                row_bytes.len()
+            )));
+        }
+
+        let offset = offset as usize;
+        let size = size as usize;
+        let end = offset.checked_add(size).ok_or_else(|| {
+            Self::decode_error(format!(
+                "UnsafeRow array field range overflows: offset={offset}, size={size}"
+            ))
+        })?;
+        if end > row_bytes.len() {
+            return Err(Self::decode_error(format!(
+                "UnsafeRow array field range is out of bounds: offset={offset}, size={size}, row_size={}",
+                row_bytes.len()
+            )));
+        }
+        if size < 8 {
+            return Err(Self::decode_error(format!(
+                "UnsafeArrayData collect buffer is too small: {size} bytes"
+            )));
+        }
+
+        let array_addr = unsafe { row_bytes.as_ptr().add(offset) } as i64;
+        Ok(Some(SparkUnsafeArray::new(array_addr)))
+    }
+
+    fn decode_error(message: impl Into<String>) -> DataFusionError {
+        DataFusionError::Execution(format!(
+            "Failed to decode Spark UnsafeRow collect aggregate buffer: {}",
+            message.into()
+        ))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Int32Array, ListArray, StringArray};
+    use arrow::datatypes::Field;
+
+    fn collect_state_decoder(element_type: DataType) -> SparkCollectStateDecoder {
+        SparkCollectStateDecoder {
+            state_field: Arc::new(Field::new_list(
+                "collect_state",
+                Field::new_list_field(element_type, true),
+                true,
+            )),
+        }
+    }
+
+    fn unsafe_row_with_array(array: Vec<u8>) -> Vec<u8> {
+        const ARRAY_OFFSET: usize = 16;
+        let mut row = vec![0_u8; ARRAY_OFFSET];
+        let offset_and_size = ((ARRAY_OFFSET as i64) << 32) | array.len() as i64;
+        row[8..16].copy_from_slice(&offset_and_size.to_le_bytes());
+        row.extend_from_slice(&array);
+        row
+    }
+
+    fn unsafe_array_i32(values: &[Option<i32>]) -> Vec<u8> {
+        let num_elements = values.len();
+        let bitset_words = num_elements.div_ceil(64);
+        let header_len = 8 + bitset_words * 8;
+        let mut bytes = vec![0_u8; header_len + num_elements * std::mem::size_of::<i32>()];
+        bytes[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes());
+
+        let mut null_bits = 0_i64;
+        for (idx, value) in values.iter().enumerate() {
+            let value_offset = header_len + idx * std::mem::size_of::<i32>();
+            match value {
+                Some(value) => {
+                    bytes[value_offset..value_offset + 4].copy_from_slice(&value.to_le_bytes())
+                }
+                None => null_bits |= 1_i64 << idx,
+            }
+        }
+        if bitset_words > 0 {
+            bytes[8..16].copy_from_slice(&null_bits.to_le_bytes());
+        }
+
+        bytes
+    }
+
+    fn unsafe_array_utf8(values: &[Option<&str>]) -> Vec<u8> {
+        let num_elements = values.len();
+        let bitset_words = num_elements.div_ceil(64);
+        let header_len = 8 + bitset_words * 8;
+        let fixed_len = num_elements * 8;
+        let mut bytes = vec![0_u8; header_len + fixed_len];
+        bytes[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes());
+
+        let mut null_bits = 0_i64;
+        for (idx, value) in values.iter().enumerate() {
+            let slot_offset = header_len + idx * 8;
+            match value {
+                Some(value) => {
+                    let value_offset = bytes.len();
+                    let value_bytes = value.as_bytes();
+                    bytes.extend_from_slice(value_bytes);
+                    let padded_len = value_bytes.len().next_multiple_of(8);
+                    bytes.resize(value_offset + padded_len, 0);
+
+                    let offset_and_size = ((value_offset as i64) << 32) | value_bytes.len() as i64;
+                    bytes[slot_offset..slot_offset + 8]
+                        .copy_from_slice(&offset_and_size.to_le_bytes());
+                }
+                None => null_bits |= 1_i64 << idx,
+            }
+        }
+        if bitset_words > 0 {
+            bytes[8..16].copy_from_slice(&null_bits.to_le_bytes());
+        }
+
+        bytes
+    }
+
+    #[test]
+    fn decodes_spark_collect_binary_int_state() {
+        let first = unsafe_row_with_array(unsafe_array_i32(&[Some(1)]));
+        // The first row is intentionally 36 bytes, so the second row starts at an unaligned
+        // address in the Arrow Binary values buffer. The decoder must not rely on UnsafeRow
+        // alignment once Spark bytes have been materialized into Arrow BinaryArray storage.
+        assert_eq!(first.len(), 36);
+        let second = unsafe_row_with_array(unsafe_array_i32(&[Some(2), None, Some(4)]));
+        let third = unsafe_row_with_array(unsafe_array_i32(&[]));
+
+        let binary = BinaryArray::from_iter([
+            Some(first.as_slice()),
+            Some(second.as_slice()),
+            None,
+            Some(third.as_slice()),
+        ]);
+        let input = [Arc::new(binary) as ArrayRef];
+        let decoder = collect_state_decoder(DataType::Int32);
+        let decoded = decoder.decode(&input).unwrap();
+        let list = decoded.arrays()[0]
+            .as_any()
+            .downcast_ref::<ListArray>()
+            .unwrap();
+
+        assert_eq!(list.len(), 4);
+        assert!(!list.is_null(0));
+        assert!(!list.is_null(1));
+        assert!(list.is_null(2));
+        assert!(!list.is_null(3));
+
+        let first_values = list
+            .value(0)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap()
+            .values()
+            .to_vec();
+        assert_eq!(first_values, vec![1]);
+
+        let second_values = list
+            .value(1)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap()
+            .clone();
+        assert_eq!(second_values.values(), &[2, 0, 4]);
+        assert!(!second_values.is_null(0));
+        assert!(second_values.is_null(1));
+        assert!(!second_values.is_null(2));
+
+        assert_eq!(list.value(3).len(), 0);
+    }
+
+    #[test]
+    fn decodes_spark_collect_binary_string_state() {
+        let row = unsafe_row_with_array(unsafe_array_utf8(&[
+            Some("alpha"),
+            Some("βeta"),
+            None,
+            Some("spark"),
+        ]));
+        let binary = BinaryArray::from_iter([Some(row.as_slice())]);
+        let input = [Arc::new(binary) as ArrayRef];
+        let decoder = collect_state_decoder(DataType::Utf8);
+        let decoded = decoder.decode(&input).unwrap();
+        let list = decoded.arrays()[0]
+            .as_any()
+            .downcast_ref::<ListArray>()
+            .unwrap();
+        let values = list
+            .value(0)
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap()
+            .clone();
+
+        assert_eq!(values.value(0), "alpha");
+        assert_eq!(values.value(1), "βeta");
+        assert!(values.is_null(2));
+        assert_eq!(values.value(3), "spark");
+    }
+}
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 e52e9ff693..722afcdda2 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
@@ -43,7 +43,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat
 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}
-import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, TimestampNTZType, TimestampType}
+import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, TimestampNTZType, TimestampType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.SerializableConfiguration
 import org.apache.spark.util.io.ChunkedByteBuffer
@@ -1552,16 +1552,14 @@ trait CometBaseAggregate {
       return None
     }
 
-    if (partialMergeCollectHasNonNativeState(aggregate)) {
-      withFallbackReason(
-        aggregate,
-        "PartialMerge collect aggregate requires native array intermediate state")
-      return None
-    }
-
     if (sparkPartialMergeMode) {
       val partialMergeAggs = aggregate.aggregateExpressions.filter(_.mode == PartialMerge)
-      if (!QueryPlanSerde.allAggsSupportMixedExecution(partialMergeAggs)) {
+      val unsupportedAggs = partialMergeAggs.filterNot { aggExpr =>
+        QueryPlanSerde.allAggsSupportMixedExecution(Seq(aggExpr)) ||
+        aggExpr.aggregateFunction.isInstanceOf[CollectList] ||
+        aggExpr.aggregateFunction.isInstanceOf[CollectSet]
+      }
+      if (unsupportedAggs.nonEmpty) {
         withFallbackReason(
           aggregate,
           "Spark PartialMerge aggregate without Comet Partial requires compatible " +
@@ -1798,38 +1796,6 @@ trait CometBaseAggregate {
         .build())
   }
 
-  /**
-   * Comet's collect_list/collect_set partial state is an ArrayType, while Spark's JVM
-   * TypedImperativeAggregate state is BinaryType. For PartialMerge, the native planner binds
-   * state input columns directly from the child schema using initialInputBufferOffset, so do not
-   * convert a collect PartialMerge when the incoming state is still Spark's binary buffer.
-   */
-  private def partialMergeCollectHasNonNativeState(aggregate: BaseAggregateExec): Boolean = {
-    var stateOffset = aggregate.initialInputBufferOffset
-
-    aggregate.aggregateExpressions.exists { aggExpr =>
-      val bufferAttrs = aggExpr.aggregateFunction.aggBufferAttributes
-      val hasNonNativeCollectState = aggExpr.mode == PartialMerge &&
-        (aggExpr.aggregateFunction.isInstanceOf[CollectList] ||
-          aggExpr.aggregateFunction.isInstanceOf[CollectSet]) &&
-        bufferAttrs.indices.exists { i =>
-          val inputIdx = stateOffset + i
-          inputIdx >= aggregate.child.output.length ||
-          (aggregate.child.output(inputIdx).dataType match {
-            case _: ArrayType => false
-            case BinaryType => true
-            case _ => true
-          })
-        }
-
-      if (aggExpr.mode == PartialMerge) {
-        stateOffset += bufferAttrs.length
-      }
-
-      hasNonNativeCollectState
-    }
-  }
-
   /**
    * 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.
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 06cf9afef3..a215314561 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -71,13 +71,12 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
-  test(
-    "collect_list/collect_set with distinct aggregate falls back on Spark PartialMerge state") {
+  test("collect_list/collect_set with distinct aggregate decodes Spark PartialMerge state") {
     // This mirrors Spark's DataFrameAggregateSuite SPARK-17616 test. With LocalTableScan disabled,
     // the lower collect partial aggregate stays on Spark and produces a BinaryType JVM buffer.
-    // Native PartialMerge must not try to merge that buffer as DataFusion's list-typed state. Run
-    // the same hash-map configurations used by the inherited Spark aggregate suites that failed in
-    // CI.
+    // Native PartialMerge decodes that UnsafeRow/UnsafeArray buffer into DataFusion's list-typed
+    // state before merging. Run the same hash-map configurations used by the inherited Spark
+    // aggregate suites that failed in CI.
     for ((twoLevelAggMap, vectorizedHashMap) <- Seq(
         (false, false),
         (true, false),
@@ -89,10 +88,15 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         SQLConf.ENABLE_VECTORIZED_HASH_MAP.key -> vectorizedHashMap.toString) {
         val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d"))
           .toDF("x", "y", "z")
-        for (collect <- Seq(collect_list(_: Column), collect_set(_: Column))) {
-          checkSparkAnswerAndFallbackReason(
+        for ((name, collect) <- Seq(
+            "collect_list" -> collect_list(_: Column),
+            "collect_set" -> collect_set(_: Column))) {
+          val (_, cometPlan) = checkSparkAnswerAndOperator(
             df.groupBy($"x").agg(count_distinct($"y"), sort_array(collect($"z"))),
-            "PartialMerge collect aggregate requires native array intermediate state")
+            Seq(classOf[CometHashAggregateExec]))
+          assert(
+            cometPlan.toString.contains(s"merge_$name"),
+            s"Expected $name PartialMerge stage to run natively; plan:\n$cometPlan")
         }
       }
     }

From 165325be9eab2c29acfbc7753b4a2e3ff5a244e6 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Mon, 29 Jun 2026 05:01:36 +0800
Subject: [PATCH 08/10] fix lint

---
 .../scala/org/apache/comet/exec/CometAggregateSuite.scala   | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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 a215314561..43f3c6c580 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -88,9 +88,9 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         SQLConf.ENABLE_VECTORIZED_HASH_MAP.key -> vectorizedHashMap.toString) {
         val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d"))
           .toDF("x", "y", "z")
-        for ((name, collect) <- Seq(
-            "collect_list" -> collect_list(_: Column),
-            "collect_set" -> collect_set(_: Column))) {
+        for ((name, collect) <- Seq[(String, Column => Column)](
+            "collect_list" -> (col => collect_list(col)),
+            "collect_set" -> (col => collect_set(col)))) {
           val (_, cometPlan) = checkSparkAnswerAndOperator(
             df.groupBy($"x").agg(count_distinct($"y"), sort_array(collect($"z"))),
             Seq(classOf[CometHashAggregateExec]))

From d15be9602f5de564db72db84fdee3e83e7b79178 Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Mon, 29 Jun 2026 12:29:54 +0800
Subject: [PATCH 09/10] fix build

---
 .../scala/org/apache/comet/exec/CometAggregateSuite.scala   | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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 43f3c6c580..b89cfa70d2 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -27,7 +27,9 @@ import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
 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.LocalTableScanExec
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec
 import org.apache.spark.sql.functions.{avg, col, collect_list, collect_set, count_distinct, sort_array, sum}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
@@ -93,7 +95,9 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
             "collect_set" -> (col => collect_set(col)))) {
           val (_, cometPlan) = checkSparkAnswerAndOperator(
             df.groupBy($"x").agg(count_distinct($"y"), sort_array(collect($"z"))),
-            Seq(classOf[CometHashAggregateExec]))
+            Seq(classOf[CometHashAggregateExec]),
+            classOf[ObjectHashAggregateExec],
+            classOf[LocalTableScanExec])
           assert(
             cometPlan.toString.contains(s"merge_$name"),
             s"Expected $name PartialMerge stage to run natively; plan:\n$cometPlan")

From 36e928d46290c52741fcdad13da2b6f26db81a1e Mon Sep 17 00:00:00 2001
From: peterxcli <peterxcli@gmail.com>
Date: Wed, 1 Jul 2026 16:45:23 +0800
Subject: [PATCH 10/10] simplify

---
 native/core/src/execution/merge_as_partial.rs |   8 +-
 .../src/execution/spark_aggregate_state.rs    | 111 +++++-------------
 .../apache/spark/sql/comet/operators.scala    |  12 +-
 3 files changed, 41 insertions(+), 90 deletions(-)

diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs
index 25471d3acf..a0c1753145 100644
--- a/native/core/src/execution/merge_as_partial.rs
+++ b/native/core/src/execution/merge_as_partial.rs
@@ -220,12 +220,12 @@ impl Accumulator for MergeAsPartialAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         // Redirect update to merge — this is the key trick.
         let decoded = self.state_decoder.decode(values)?;
-        self.inner.merge_batch(decoded.arrays())
+        self.inner.merge_batch(decoded.as_ref())
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
         let decoded = self.state_decoder.decode(states)?;
-        self.inner.merge_batch(decoded.arrays())
+        self.inner.merge_batch(decoded.as_ref())
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
@@ -264,7 +264,7 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator {
         // Redirect update to merge — this is the key trick.
         let decoded = self.state_decoder.decode(values)?;
         self.inner.merge_batch(
-            decoded.arrays(),
+            decoded.as_ref(),
             group_indices,
             opt_filter,
             total_num_groups,
@@ -280,7 +280,7 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator {
     ) -> Result<()> {
         let decoded = self.state_decoder.decode(values)?;
         self.inner.merge_batch(
-            decoded.arrays(),
+            decoded.as_ref(),
             group_indices,
             opt_filter,
             total_num_groups,
diff --git a/native/core/src/execution/spark_aggregate_state.rs b/native/core/src/execution/spark_aggregate_state.rs
index 217a6785e5..222581ca2b 100644
--- a/native/core/src/execution/spark_aggregate_state.rs
+++ b/native/core/src/execution/spark_aggregate_state.rs
@@ -17,32 +17,17 @@
 
 //! Decoders for Spark JVM aggregate state consumed by native PartialMerge.
 
-use std::sync::Arc;
+use std::{borrow::Cow, sync::Arc};
 
 use arrow::array::{
-    builder::{make_builder, ListBuilder},
-    Array, ArrayRef, BinaryArray, LargeBinaryArray,
+    builder::{make_builder, ArrayBuilder, ListBuilder},
+    Array, ArrayRef, BinaryArray, GenericByteArray, LargeBinaryArray, OffsetSizeTrait,
 };
-use arrow::datatypes::{DataType, FieldRef};
+use arrow::datatypes::{DataType, FieldRef, GenericBinaryType};
 use datafusion::common::{DataFusionError, Result};
 use datafusion::physical_expr::aggregate::AggregateFunctionExpr;
 use datafusion_comet_shuffle::spark_unsafe::list::{append_to_builder, SparkUnsafeArray};
 
-/// State arrays after optional Spark-JVM-state decoding.
-pub(crate) enum DecodedState<'a> {
-    Borrowed(&'a [ArrayRef]),
-    Owned(Vec<ArrayRef>),
-}
-
-impl DecodedState<'_> {
-    pub(crate) fn arrays(&self) -> &[ArrayRef] {
-        match self {
-            Self::Borrowed(values) => values,
-            Self::Owned(values) => values,
-        }
-    }
-}
-
 /// Decoder used by `MergeAsPartial` before forwarding state to the inner accumulator.
 #[derive(Clone, Debug)]
 pub(crate) enum PartialMergeStateDecoder {
@@ -57,9 +42,9 @@ impl PartialMergeStateDecoder {
             .unwrap_or(Self::PassThrough)
     }
 
-    pub(crate) fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<DecodedState<'a>> {
+    pub(crate) fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<Cow<'a, [ArrayRef]>> {
         match self {
-            Self::PassThrough => Ok(DecodedState::Borrowed(values)),
+            Self::PassThrough => Ok(Cow::Borrowed(values)),
             Self::SparkCollect(decoder) => decoder.decode(values),
         }
     }
@@ -75,24 +60,23 @@ impl PartialMergeStateDecoder {
 /// Arrow `ListArray` before calling the inner accumulator's `merge_batch`.
 #[derive(Clone, Debug)]
 pub(crate) struct SparkCollectStateDecoder {
-    state_field: FieldRef,
+    item_field: FieldRef,
 }
 
 impl SparkCollectStateDecoder {
     fn try_new(inner_expr: &AggregateFunctionExpr, state_fields: &[FieldRef]) -> Option<Self> {
         match (inner_expr.fun().name(), state_fields) {
-            ("collect_list" | "collect_set", [state_field])
-                if matches!(state_field.data_type(), DataType::List(_)) =>
-            {
-                Some(Self {
-                    state_field: Arc::clone(state_field),
-                })
-            }
+            ("collect_list" | "collect_set", [state_field]) => match state_field.data_type() {
+                DataType::List(item_field) => Some(Self {
+                    item_field: Arc::clone(item_field),
+                }),
+                _ => None,
+            },
             _ => None,
         }
     }
 
-    fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<DecodedState<'a>> {
+    fn decode<'a>(&self, values: &'a [ArrayRef]) -> Result<Cow<'a, [ArrayRef]>> {
         if values.len() != 1 {
             return Err(DataFusionError::Internal(format!(
                 "Spark collect state decoder expected one state column, got {}",
@@ -108,7 +92,7 @@ impl SparkCollectStateDecoder {
                     .ok_or_else(|| {
                         Self::decode_error("expected BinaryArray for Binary collect state")
                     })?;
-                Ok(DecodedState::Owned(vec![self.decode_binary_array(array)?]))
+                Ok(Cow::Owned(vec![self.decode_binary_array(array)?]))
             }
             DataType::LargeBinary => {
                 let array = values[0]
@@ -119,71 +103,42 @@ impl SparkCollectStateDecoder {
                             "expected LargeBinaryArray for LargeBinary collect state",
                         )
                     })?;
-                Ok(DecodedState::Owned(vec![
-                    self.decode_large_binary_array(array)?
-                ]))
-            }
-            _ => Ok(DecodedState::Borrowed(values)),
-        }
-    }
-
-    fn decode_binary_array(&self, array: &BinaryArray) -> Result<ArrayRef> {
-        let item_field = self.item_field()?;
-        let mut builder = self.new_list_builder(item_field, array.len());
-
-        for row_idx in 0..array.len() {
-            if array.is_null(row_idx) {
-                builder.append_null();
-            } else {
-                self.append_unsafe_row_array(array.value(row_idx), item_field, &mut builder)?;
+                Ok(Cow::Owned(vec![self.decode_binary_array(array)?]))
             }
+            _ => Ok(Cow::Borrowed(values)),
         }
-
-        Ok(Arc::new(builder.finish()))
     }
 
-    fn decode_large_binary_array(&self, array: &LargeBinaryArray) -> Result<ArrayRef> {
-        let item_field = self.item_field()?;
-        let mut builder = self.new_list_builder(item_field, array.len());
+    fn decode_binary_array<O: OffsetSizeTrait>(
+        &self,
+        array: &GenericByteArray<GenericBinaryType<O>>,
+    ) -> Result<ArrayRef> {
+        let mut builder = self.new_list_builder(array.len());
 
         for row_idx in 0..array.len() {
             if array.is_null(row_idx) {
                 builder.append_null();
             } else {
-                self.append_unsafe_row_array(array.value(row_idx), item_field, &mut builder)?;
+                self.append_unsafe_row_array(array.value(row_idx), &mut builder)?;
             }
         }
 
         Ok(Arc::new(builder.finish()))
     }
 
-    fn item_field(&self) -> Result<&FieldRef> {
-        match self.state_field.data_type() {
-            DataType::List(item_field) => Ok(item_field),
-            other => Err(Self::decode_error(format!(
-                "collect state field must be List, got {other:?}"
-            ))),
-        }
-    }
-
-    fn new_list_builder(
-        &self,
-        item_field: &FieldRef,
-        capacity: usize,
-    ) -> ListBuilder<Box<dyn arrow::array::builder::ArrayBuilder>> {
-        let value_builder = make_builder(item_field.data_type(), capacity);
-        ListBuilder::with_capacity(value_builder, capacity).with_field(Arc::clone(item_field))
+    fn new_list_builder(&self, capacity: usize) -> ListBuilder<Box<dyn ArrayBuilder>> {
+        let value_builder = make_builder(self.item_field.data_type(), capacity);
+        ListBuilder::with_capacity(value_builder, capacity).with_field(Arc::clone(&self.item_field))
     }
 
     fn append_unsafe_row_array(
         &self,
         row_bytes: &[u8],
-        item_field: &FieldRef,
-        builder: &mut ListBuilder<Box<dyn arrow::array::builder::ArrayBuilder>>,
+        builder: &mut ListBuilder<Box<dyn ArrayBuilder>>,
     ) -> Result<()> {
         match Self::spark_array_from_single_field_unsafe_row(row_bytes)? {
             Some(array) => {
-                append_to_builder::<true>(item_field.data_type(), builder.values(), &array)
+                append_to_builder::<true>(self.item_field.data_type(), builder.values(), &array)
                     .map_err(|e| Self::decode_error(e.to_string()))?;
                 builder.append(true);
             }
@@ -270,11 +225,7 @@ mod tests {
 
     fn collect_state_decoder(element_type: DataType) -> SparkCollectStateDecoder {
         SparkCollectStateDecoder {
-            state_field: Arc::new(Field::new_list(
-                "collect_state",
-                Field::new_list_field(element_type, true),
-                true,
-            )),
+            item_field: Arc::new(Field::new_list_field(element_type, true)),
         }
     }
 
@@ -363,7 +314,7 @@ mod tests {
         let input = [Arc::new(binary) as ArrayRef];
         let decoder = collect_state_decoder(DataType::Int32);
         let decoded = decoder.decode(&input).unwrap();
-        let list = decoded.arrays()[0]
+        let list = decoded.as_ref()[0]
             .as_any()
             .downcast_ref::<ListArray>()
             .unwrap();
@@ -409,7 +360,7 @@ mod tests {
         let input = [Arc::new(binary) as ArrayRef];
         let decoder = collect_state_decoder(DataType::Utf8);
         let decoded = decoder.decode(&input).unwrap();
-        let list = decoded.arrays()[0]
+        let list = decoded.as_ref()[0]
             .as_any()
             .downcast_ref::<ListArray>()
             .unwrap();
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 722afcdda2..4f332bb056 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
@@ -1553,13 +1553,13 @@ trait CometBaseAggregate {
     }
 
     if (sparkPartialMergeMode) {
-      val partialMergeAggs = aggregate.aggregateExpressions.filter(_.mode == PartialMerge)
-      val unsupportedAggs = partialMergeAggs.filterNot { aggExpr =>
-        QueryPlanSerde.allAggsSupportMixedExecution(Seq(aggExpr)) ||
-        aggExpr.aggregateFunction.isInstanceOf[CollectList] ||
-        aggExpr.aggregateFunction.isInstanceOf[CollectSet]
+      val hasUnsupportedAgg = aggregate.aggregateExpressions.exists { aggExpr =>
+        aggExpr.mode == PartialMerge &&
+        !QueryPlanSerde.allAggsSupportMixedExecution(Seq(aggExpr)) &&
+        !aggExpr.aggregateFunction.isInstanceOf[CollectList] &&
+        !aggExpr.aggregateFunction.isInstanceOf[CollectSet]
       }
-      if (unsupportedAggs.nonEmpty) {
+      if (hasUnsupportedAgg) {
         withFallbackReason(
           aggregate,
           "Spark PartialMerge aggregate without Comet Partial requires compatible " +