perf: give collect_list and collect_set a native GroupsAccumulator

docs/source/contributor-guide/expression-audits/agg_funcs.md

@@ -57,8 +57,9 @@
 - 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 (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (the same pattern applies to [collect_set](#collect-set)).
+- Comet implementation: `CometCollectList` (`native/spark-expr/src/agg_funcs/collect.rs`) delegates the ungrouped path 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 (matching Spark's `defaultResult`); grouped aggregation uses Comet's own `GroupsAccumulator`, which retains the input arrays and gathers them into group order on emit. The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (the same pattern applies to [collect_set](#collect-set)).
 - Buffer shape (applies equally to `collect_set`): both are `TypedImperativeAggregate`s, so Spark's `aggBufferAttributes` declares the intermediate buffer as `BinaryType` (serialized state) while the native accumulator's `state_fields` is a `List`. `CometBaseAggregate.adjustOutputForNativeState` rewrites the Comet-side Partial output to the list shape. Neither collector can split a Partial and Final across Comet and Spark, and a multi-stage distinct rewrite (which inserts a `PartialMerge` stage) forces the whole chain back to Spark ([#4724](https://github.com/apache/datafusion-comet/issues/4724)).
+- Performance (tuned 2026-09-09, PR [#5803](https://github.com/apache/datafusion-comet/pull/5803)): grouped `collect_list` used to run through DataFusion's `GroupsAccumulatorAdapter`, which keeps one boxed `Accumulator` per group and slices every batch into a per-group `update_batch` call. `CollectListGroupsAccumulator` instead records a `(group, row range)` contribution per batch and rearranges them with a counting sort on emit, gathering long runs with `concat` and scattered rows with `interleave`. 16-99% faster ([#5797](https://github.com/apache/datafusion-comet/issues/5797)). Benchmark: `benches/collect.rs`.
 - Spark 4.2 (preview): `CollectList` and `CollectSet` gain an `ignoreNulls` field (default `true`); `RESPECT NULLS` sets it to `false` and keeps null elements. The native path always drops nulls, so `CometCollectShim` reads the field per Spark version (always `true` on 3.4-4.1) and `CometCollectList` / `CometCollectSet` report `Unsupported` when it is `false`, falling back to Spark.
 
 ## collect_set
@@ -67,7 +68,8 @@
 - Spark 3.5.8 (audited 2026-07-27): identical to 3.4.3.
 - Spark 4.0.1 (audited 2026-07-27): adds `with UnaryLike[Expression]` to the case class, and `checkInputDataTypes` additionally requires `UnsafeRowUtils.isBinaryStable(child.dataType)`, so non-default-collation strings are rejected along with maps. Deduplication semantics unchanged.
 - Spark 4.1.1 (audited 2026-07-27): identical to 4.0.1.
-- Comet implementation: the native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps `DistinctArrayAggAccumulator` with `ignore_nulls = true` in a `NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty array. The `containsNull` mismatch against Spark's declared output type, and its rationale, are identical to [collect_list](#collect-list).
+- Comet implementation: `CometCollectSet` (`native/spark-expr/src/agg_funcs/collect.rs`) delegates the ungrouped path to `datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps `DistinctArrayAggAccumulator` with `ignore_nulls = true` in a `NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty array; grouped aggregation uses Comet's own `GroupsAccumulator`, which keeps the distinct values row-encoded in one arena keyed by `(group, value)`. Deduplication is still `arrow::row` encoded-byte equality, so which values collapse together is unchanged; the emitted order is now insertion order rather than hash-table order. The `containsNull` mismatch against Spark's declared output type, and its rationale, are identical to [collect_list](#collect-list).
+- Performance (tuned 2026-09-09, PR [#5803](https://github.com/apache/datafusion-comet/pull/5803)): as for [collect_list](#collect-list), grouped `collect_set` no longer goes through `GroupsAccumulatorAdapter`. `CollectSetGroupsAccumulator` encodes each batch once for all of its groups and deduplicates against an open-addressed index over a shared arena, replacing one hash table plus one owned `Row` per distinct value per group. 60-93% faster ([#5797](https://github.com/apache/datafusion-comet/issues/5797)). Benchmark: `benches/collect.rs`.
 - `CometCollectSet` reports `Incompatible` for float and double input when `spark.comet.exec.strictFloatingPoint=true`, because the native distinct comparison treats `NaN == NaN` and collapses repeated `NaN`s into a single element while Spark keeps each one. The native path for floating-point input is then opt-in via `spark.comet.expression.CollectSet.allowIncompatible=true`. All other input types are `Compatible`.
 
 ## median

native/core/src/execution/planner.rs

@@ -90,10 +90,9 @@ use datafusion::{
 };
 use datafusion_comet_spark_expr::{
     create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
-    BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SparkArraysZipFunc,
-    SparkBloomFilterVersion, SparkPercentile, SumInteger, ToCsv,
+    BloomFilterAgg, BloomFilterMightContain, CometCollectList, CometCollectSet, CsvWriteOptions,
+    EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SparkPercentile, SumInteger, ToCsv,
 };
-use datafusion_spark::function::aggregate::collect::{SparkCollectList, SparkCollectSet};
 use iceberg::expr::Bind;
 
 use crate::execution::operators::ExecutionError::GeneralError;
@@ -3172,13 +3171,13 @@ impl PhysicalPlanner {
             AggExprStruct::CollectSet(expr) => {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?;
                 let child = Self::coerce_collect_child_nullability(child, &schema)?;
-                let func = AggregateUDF::new_from_impl(SparkCollectSet::new());
+                let func = AggregateUDF::new_from_impl(CometCollectSet::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 child = Self::coerce_collect_child_nullability(child, &schema)?;
-                let func = AggregateUDF::new_from_impl(SparkCollectList::new());
+                let func = AggregateUDF::new_from_impl(CometCollectList::new());
                 Self::create_aggr_func_expr("collect_list", schema, vec![child], func)
             }
             AggExprStruct::Hllpp(expr) => {
@@ -3818,6 +3817,9 @@ impl PhysicalPlanner {
     /// schema types") — `RecordBatch::try_new` compares with `DataType::equals_datatype`, which
     /// does compare nested nullability.
     ///
+    /// The grouped path normalizes its own inputs, so this only still matters for the ungrouped
+    /// `Accumulator` (a global aggregate, or a window frame).
+    ///
     /// Casting unconditionally (rather than only when the declared type has a non-nullable
     /// nested field) makes this a normalization barrier: the accumulator is guaranteed to see
     /// arrays of exactly the type the plan declared, whichever direction the drift goes. The
@@ -5059,17 +5061,17 @@ mod tests {
         FileGroup, FileScanConfigBuilder, FileSource, ParquetSource,
     };
     use datafusion::error::DataFusionError;
-    use datafusion::logical_expr::AggregateUDF;
     use datafusion::logical_expr::ScalarUDF;
+    use datafusion::logical_expr::{AggregateUDF, EmitTo};
     use datafusion::physical_expr::expressions::Column;
     use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
     use datafusion::physical_plan::sorts::sort::SortExec;
     use datafusion::physical_plan::windows::get_ordered_partition_by_indices;
     use datafusion::physical_plan::ExecutionPlan;
     use datafusion::prelude::SessionConfig;
     use datafusion::{assert_batches_eq, physical_plan::common::collect, prelude::SessionContext};
+    use datafusion_comet_spark_expr::{CometCollectList, CometCollectSet};
     use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
-    use datafusion_spark::function::aggregate::collect::{SparkCollectList, SparkCollectSet};
     use parquet::variant::VariantType;
     use tempfile::TempDir;
     use tokio::sync::mpsc;
@@ -7261,15 +7263,16 @@ mod tests {
     }
 
     /// Builds `func` over `child` against `plan_schema`, feeds it a batch built from
-    /// `batch_schema`, and returns (type declared by `state_fields()`, type of the emitted
-    /// state array). This is the pair `RecordBatch::try_new` compares inside
+    /// `batch_schema`, and returns (type declared by `state_fields()`, type of the state the
+    /// ungrouped `Accumulator` emits, type of the state the `GroupsAccumulator` emits). The
+    /// declared type is compared against the emitted one by `RecordBatch::try_new` inside
     /// `GroupedHashAggregateStream::emit`.
     fn collect_agg_declared_vs_produced(
         child: &Arc<dyn PhysicalExpr>,
         plan_schema: &SchemaRef,
         batch_schema: &SchemaRef,
         func: AggregateUDF,
-    ) -> (DataType, DataType) {
+    ) -> (DataType, DataType, DataType) {
         let agg = PhysicalPlanner::create_aggr_func_expr(
             "collect",
             Arc::clone(plan_schema),
@@ -7288,8 +7291,14 @@ mod tests {
             .into_array(batch.num_rows())
             .unwrap();
         let mut acc = agg.create_accumulator().unwrap();
-        acc.update_batch(&[arg]).unwrap();
-        (declared, acc.state().unwrap()[0].data_type())
+        acc.update_batch(&[Arc::clone(&arg)]).unwrap();
+        let produced = acc.state().unwrap()[0].data_type();
+
+        let mut groups = agg.create_groups_accumulator().unwrap();
+        groups.update_batch(&[arg], &[0, 0], None, 1).unwrap();
+        let grouped = groups.state(EmitTo::All).unwrap()[0].data_type().clone();
+
+        (declared, produced, grouped)
     }
 
     /// `collect_list` / `collect_set` derive their accumulator state type from the declared
@@ -7302,6 +7311,10 @@ mod tests {
     /// rebuilds its state list honouring `data_type`. `collect_list` still rebuilds from the
     /// runtime array and drifts, so the coercion remains necessary. This test pins that
     /// `collect_list` still drifts without the coercion, and that the coercion normalizes both.
+    ///
+    /// The grouped path never drifts: `CollectListGroupsAccumulator` /
+    /// `CollectSetGroupsAccumulator` normalize every array they take in to the declared element
+    /// type, so their state matches the declaration with or without the coercion.
     #[test]
     fn test_collect_agg_absorbs_nested_nullability_drift() {
         let plan_schema = collect_agg_schema(collect_agg_struct_type(true));
@@ -7312,13 +7325,13 @@ mod tests {
                 .unwrap();
 
         for func in [
-            AggregateUDF::new_from_impl(SparkCollectSet::new()),
-            AggregateUDF::new_from_impl(SparkCollectList::new()),
+            AggregateUDF::new_from_impl(CometCollectSet::new()),
+            AggregateUDF::new_from_impl(CometCollectList::new()),
         ] {
             // Without the coercion, collect_list still drifts (declared nullable leaves vs
             // produced non-null ones). collect_set no longer drifts after DataFusion 55's new_list
             // fix, so only assert the drift for the func that still exhibits it.
-            let (declared, produced) =
+            let (declared, produced, grouped) =
                 collect_agg_declared_vs_produced(&raw, &plan_schema, &batch_schema, func.clone());
             if func.name() == "collect_list" {
                 assert!(
@@ -7327,15 +7340,24 @@ mod tests {
                     func.name()
                 );
             }
+            assert!(
+                declared.equals_datatype(&grouped),
+                "grouped state drifted for {}: declared {declared}, produced {grouped}",
+                func.name()
+            );
 
             // With the coercion, the argument is normalized to the declared type before it
             // reaches the accumulator, so both funcs emit a state matching the declared type.
-            let (declared, produced) =
+            let (declared, produced, grouped) =
                 collect_agg_declared_vs_produced(&coerced, &plan_schema, &batch_schema, func);
             assert!(
                 declared.equals_datatype(&produced),
                 "declared {declared} != produced {produced}"
             );
+            assert!(
+                declared.equals_datatype(&grouped),
+                "declared {declared} != grouped {grouped}"
+            );
         }
     }
 

native/spark-expr/Cargo.toml

@@ -163,6 +163,10 @@ harness = false
 name = "bloom_filter_agg"
 harness = false
 
+[[bench]]
+name = "collect"
+harness = false
+
 [[bench]]
 name = "padding"
 harness = false