This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 5c5dda8374 perf: give collect_list and collect_set a native
GroupsAccumulator (#5803)
5c5dda8374 is described below
commit 5c5dda8374d0b8b93b877281ed4b5fdca1eed324
Author: Andy Grove <[email protected]>
AuthorDate: Thu Sep 10 21:50:36 2026 -0600
perf: give collect_list and collect_set a native GroupsAccumulator (#5803)
---
.../expression-audits/agg_funcs.md | 6 +-
native/core/src/execution/planner.rs | 54 +-
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/collect.rs | 258 ++++
native/spark-expr/src/agg_funcs/collect.rs | 1255 ++++++++++++++++++++
native/spark-expr/src/agg_funcs/mod.rs | 2 +
.../apache/comet/exec/CometAggregateSuite.scala | 23 +
7 files changed, 1584 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 b2d2efae6b..8d0ff7be68 100644
--- a/docs/source/contributor-guide/expression-audits/agg_funcs.md
+++ b/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 dif [...]
+- 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 [...]
- 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 [...]
+- 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 [...]
- 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)`. D [...]
+- 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
[...]
- `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`.
## max_by
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index eeaee34478..d66638ba32 100644
--- a/native/core/src/execution/planner.rs
+++ b/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;
@@ -3174,13 +3173,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) => {
@@ -3841,6 +3840,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
@@ -5082,8 +5084,8 @@ 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;
@@ -5091,8 +5093,8 @@ mod tests {
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;
@@ -7284,15 +7286,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),
@@ -7311,8 +7314,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
@@ -7325,6 +7334,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));
@@ -7335,13 +7348,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!(
@@ -7350,15 +7363,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}"
+ );
}
}
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 98e5a99904..987d2f5935 100644
--- a/native/spark-expr/Cargo.toml
+++ b/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
diff --git a/native/spark-expr/benches/collect.rs
b/native/spark-expr/benches/collect.rs
new file mode 100644
index 0000000000..67bfa54b94
--- /dev/null
+++ b/native/spark-expr/benches/collect.rs
@@ -0,0 +1,258 @@
+// 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.
+
+//! Grouped `collect_list` / `collect_set` aggregation over a high-cardinality
string key,
+//! the shape reported in
<https://github.com/apache/datafusion-comet/issues/5797>.
+
+use arrow::array::builder::{Int64Builder, StringBuilder};
+use arrow::array::{ArrayRef, RecordBatch, StructArray};
+use arrow::datatypes::{DataType, Field, Fields, Schema};
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::AggregateUDF;
+use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode,
PhysicalGroupBy};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::{
+ datasource::memory::MemorySourceConfig, datasource::source::DataSourceExec,
+ physical_plan::collect,
+};
+use datafusion_comet_spark_expr::{CometCollectList, CometCollectSet};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::runtime::Runtime;
+
+const BATCH_SIZE: usize = 8192;
+const NUM_BATCHES: usize = 16;
+
+/// Element column under aggregation.
+#[derive(Clone, Copy)]
+enum Element {
+ Int64,
+ Utf8,
+ Struct,
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let rt = Runtime::new().unwrap();
+ let mut group = c.benchmark_group("collect");
+
+ // (label, element type, distinct group keys, distinct element values,
null ratio, keys
+ // clustered into runs rather than cycling)
+ let shapes = [
+ ("int64/high_card", Element::Int64, 16384, 1 << 20, 0, false),
+ ("utf8/high_card", Element::Utf8, 16384, 1 << 20, 0, false),
+ ("utf8/low_card", Element::Utf8, 64, 1 << 20, 0, false),
+ ("utf8/clustered", Element::Utf8, 64, 1 << 20, 0, true),
+ ("utf8/dup_elements", Element::Utf8, 16384, 64, 0, false),
+ ("utf8/nulls", Element::Utf8, 16384, 1 << 20, 3, false),
+ (
+ "struct/high_card",
+ Element::Struct,
+ 16384,
+ 1 << 20,
+ 0,
+ false,
+ ),
+ ];
+
+ for (label, element, num_groups, num_values, null_every, clustered) in
shapes {
+ let partitions = vec![build_batches(
+ element, num_groups, num_values, null_every, clustered,
+ )];
+ for (fn_name, udf) in [
+ (
+ "collect_list",
+ Arc::new(AggregateUDF::new_from_impl(CometCollectList::new())),
+ ),
+ (
+ "collect_set",
+ Arc::new(AggregateUDF::new_from_impl(CometCollectSet::new())),
+ ),
+ ] {
+ for two_stage in [false, true] {
+ let mode = if two_stage { "two_stage" } else { "partial" };
+ // Build the plan once: only execution should be measured.
+ let plan = aggregate_plan(&partitions, Arc::clone(&udf),
two_stage);
+ group.bench_function(format!("{fn_name}/{label}/{mode}"), |b| {
+ b.iter(|| {
+ rt.block_on(async {
+ let batches =
+ collect(Arc::clone(&plan),
Arc::new(TaskContext::default()))
+ .await
+ .unwrap();
+ assert!(!batches.is_empty());
+ })
+ })
+ });
+ }
+ }
+ }
+
+ group.finish();
+}
+
+fn aggregate_plan(
+ partitions: &[Vec<RecordBatch>],
+ udf: Arc<AggregateUDF>,
+ two_stage: bool,
+) -> Arc<dyn ExecutionPlan> {
+ let schema = &partitions[0][0].schema();
+ let scan: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
+ MemorySourceConfig::try_new(partitions, Arc::clone(schema),
None).unwrap(),
+ )));
+ let key: Arc<dyn PhysicalExpr> = Arc::new(Column::new("key", 0));
+ let value: Arc<dyn PhysicalExpr> = Arc::new(Column::new("value", 1));
+ let aggr_expr = Arc::new(
+ AggregateExprBuilder::new(udf, vec![value])
+ .schema(Arc::clone(schema))
+ .alias("agg")
+ .build()
+ .unwrap(),
+ );
+
+ let partial: Arc<dyn ExecutionPlan> = Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Partial,
+ PhysicalGroupBy::new_single(vec![(Arc::clone(&key),
"key".to_string())]),
+ vec![Arc::clone(&aggr_expr)],
+ vec![None],
+ scan,
+ Arc::clone(schema),
+ )
+ .unwrap(),
+ );
+
+ if !two_stage {
+ return partial;
+ }
+
+ Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Final,
+ PhysicalGroupBy::new_single(vec![(key, "key".to_string())]),
+ vec![aggr_expr],
+ vec![None],
+ partial,
+ Arc::clone(schema),
+ )
+ .unwrap(),
+ )
+}
+
+fn build_batches(
+ element: Element,
+ num_groups: usize,
+ num_values: usize,
+ null_every: usize,
+ clustered: bool,
+) -> Vec<RecordBatch> {
+ let value_field = match element {
+ Element::Int64 => Field::new("value", DataType::Int64, true),
+ Element::Utf8 => Field::new("value", DataType::Utf8, true),
+ Element::Struct => Field::new("value",
DataType::Struct(struct_fields()), true),
+ };
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Utf8, false),
+ value_field,
+ ]));
+
+ (0..NUM_BATCHES)
+ .map(|batch| {
+ let start = batch * BATCH_SIZE;
+ let rows = start..start + BATCH_SIZE;
+
+ let mut keys = StringBuilder::with_capacity(BATCH_SIZE, BATCH_SIZE
* 16);
+ for i in rows.clone() {
+ // Clustered keys hold for a run of rows, which is what lets
the accumulator
+ // coalesce them into one range; cycling keys give a fresh
range per row.
+ let group = if clustered {
+ (i / BATCH_SIZE.div_ceil(num_groups)) % num_groups
+ } else {
+ i % num_groups
+ };
+ keys.append_value(format!("key_{group}"));
+ }
+ let keys: ArrayRef = Arc::new(keys.finish());
+
+ let is_null = |i: usize| null_every != 0 &&
i.is_multiple_of(null_every);
+ let values: ArrayRef = match element {
+ Element::Int64 => {
+ let mut b = Int64Builder::with_capacity(BATCH_SIZE);
+ for i in rows {
+ if is_null(i) {
+ b.append_null();
+ } else {
+ b.append_value((i % num_values) as i64);
+ }
+ }
+ Arc::new(b.finish())
+ }
+ Element::Utf8 => {
+ let mut b = StringBuilder::with_capacity(BATCH_SIZE,
BATCH_SIZE * 24);
+ for i in rows {
+ if is_null(i) {
+ b.append_null();
+ } else {
+ b.append_value(format!("this is value #{}", i %
num_values));
+ }
+ }
+ Arc::new(b.finish())
+ }
+ Element::Struct => {
+ let mut a = StringBuilder::with_capacity(BATCH_SIZE,
BATCH_SIZE * 24);
+ let mut b = Int64Builder::with_capacity(BATCH_SIZE);
+ for i in rows {
+ a.append_value(format!("this is value #{}", i %
num_values));
+ b.append_value((i % num_values) as i64);
+ }
+ Arc::new(StructArray::new(
+ struct_fields(),
+ vec![Arc::new(a.finish()) as ArrayRef,
Arc::new(b.finish())],
+ None,
+ ))
+ }
+ };
+
+ RecordBatch::try_new(Arc::clone(&schema), vec![keys,
values]).unwrap()
+ })
+ .collect()
+}
+
+fn struct_fields() -> Fields {
+ vec![
+ Arc::new(Field::new("a", DataType::Utf8, true)),
+ Arc::new(Field::new("b", DataType::Int64, true)),
+ ]
+ .into()
+}
+
+fn config() -> Criterion {
+ Criterion::default()
+ .measurement_time(Duration::from_millis(2000))
+ .warm_up_time(Duration::from_millis(500))
+ .sample_size(10)
+}
+
+criterion_group! {
+ name = benches;
+ config = config();
+ targets = criterion_benchmark
+}
+criterion_main!(benches);
diff --git a/native/spark-expr/src/agg_funcs/collect.rs
b/native/spark-expr/src/agg_funcs/collect.rs
new file mode 100644
index 0000000000..00b2759c25
--- /dev/null
+++ b/native/spark-expr/src/agg_funcs/collect.rs
@@ -0,0 +1,1255 @@
+// 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.
+
+//! Spark `collect_list` / `collect_set`, wrapping `datafusion_spark`'s
implementations with a
+//! `GroupsAccumulator` so grouped aggregation does not go through DataFusion's
+//! `GroupsAccumulatorAdapter`.
+//!
+//! The adapter keeps one boxed `Accumulator` per group and slices every input
batch into a
+//! per-group `update_batch` call, which dominates the runtime once a query
groups by a
+//! high-cardinality key. See
<https://github.com/apache/datafusion-comet/issues/5797>.
+
+use arrow::array::{
+ new_empty_array, Array, ArrayRef, AsArray, BooleanArray, ListArray,
UInt32Array,
+};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::compute::{cast, concat, interleave, take};
+use arrow::datatypes::{DataType, Field, FieldRef};
+use arrow::row::{RowConverter, SortField};
+use datafusion::common::{
+ internal_datafusion_err, internal_err, DataFusionError, Result as
DFResult, ScalarValue,
+};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
+};
+use datafusion_spark::function::aggregate::collect::{SparkCollectList,
SparkCollectSet};
+use std::mem::{size_of, take as take_field};
+use std::ops::Range;
+use std::sync::Arc;
+use twox_hash::XxHash64;
+
+/// Spark's `collect_list`.
+///
+/// Everything except grouped aggregation is delegated to
[`SparkCollectList`]; this type only
+/// adds the [`CollectListGroupsAccumulator`].
+#[derive(Debug, PartialEq, Eq, Hash, Default)]
+pub struct CometCollectList {
+ inner: SparkCollectList,
+}
+
+impl CometCollectList {
+ pub fn new() -> Self {
+ Self::default()
+ }
+}
+
+impl AggregateUDFImpl for CometCollectList {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn signature(&self) -> &Signature {
+ self.inner.signature()
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> DFResult<DataType> {
+ self.inner.return_type(arg_types)
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ self.inner.state_fields(args)
+ }
+
+ fn accumulator(&self, acc_args: AccumulatorArgs) -> DFResult<Box<dyn
Accumulator>> {
+ self.inner.accumulator(acc_args)
+ }
+
+ fn default_value(&self, data_type: &DataType) -> DFResult<ScalarValue> {
+ self.inner.default_value(data_type)
+ }
+
+ fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
+ !args.is_distinct && args.order_bys.is_empty()
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(CollectListGroupsAccumulator::new(
+ args.expr_fields[0].data_type().clone(),
+ )))
+ }
+}
+
+/// Spark's `collect_set`.
+///
+/// Everything except grouped aggregation is delegated to [`SparkCollectSet`].
+#[derive(Debug, PartialEq, Eq, Hash, Default)]
+pub struct CometCollectSet {
+ inner: SparkCollectSet,
+}
+
+impl CometCollectSet {
+ pub fn new() -> Self {
+ Self::default()
+ }
+}
+
+impl AggregateUDFImpl for CometCollectSet {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn signature(&self) -> &Signature {
+ self.inner.signature()
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> DFResult<DataType> {
+ self.inner.return_type(arg_types)
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
+ self.inner.state_fields(args)
+ }
+
+ fn accumulator(&self, acc_args: AccumulatorArgs) -> DFResult<Box<dyn
Accumulator>> {
+ self.inner.accumulator(acc_args)
+ }
+
+ fn default_value(&self, data_type: &DataType) -> DFResult<ScalarValue> {
+ self.inner.default_value(data_type)
+ }
+
+ fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
+ !args.is_distinct && args.order_bys.is_empty()
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ args: AccumulatorArgs,
+ ) -> DFResult<Box<dyn GroupsAccumulator>> {
+ Ok(Box::new(CollectSetGroupsAccumulator::try_new(
+ args.expr_fields[0].data_type().clone(),
+ )?))
+ }
+}
+
+/// Casts `array` to `datatype` when the two have drifted apart, which a
dictionary-encoded input
+/// or nested field nullability surviving a shuffle round trip can both cause.
The cast is
+/// metadata-only when only nullability differs.
+fn coerce(array: &ArrayRef, datatype: &DataType) -> DFResult<ArrayRef> {
+ if array.data_type() == datatype {
+ Ok(Arc::clone(array))
+ } else {
+ Ok(cast(array.as_ref(), datatype)?)
+ }
+}
+
+/// The list field a `collect_list` / `collect_set` state or result is built
from. Spark's element
+/// type is always nullable, and the group's own list entry is never null.
+fn list_field(element_type: &DataType) -> FieldRef {
+ Arc::new(Field::new_list_field(element_type.clone(), true))
+}
+
+/// `GroupsAccumulator` for Spark's `collect_list`.
+///
+/// Input arrays are retained rather than copied per group; each one is paired
with the
+/// `(group, row range)` contributions it makes. `evaluate` then rearranges
those contributions
+/// into group order with a counting sort and gathers the values in a single
pass, so the cost is
+/// proportional to the data rather than to the number of groups.
+#[derive(Debug)]
+struct CollectListGroupsAccumulator {
+ /// The element type as declared by the plan.
+ datatype: DataType,
+ /// Source arrays: input arrays from `update_batch`, list values arrays
from `merge_batch`.
+ batches: Vec<ArrayRef>,
+ /// Contributions of each retained batch, in the order they were seen.
+ ranges: Vec<Vec<RowRange>>,
+ /// Running total of what `batches` and `ranges` hold. `size()` is called
once per input
+ /// batch, so it cannot afford to walk the retained batches to work this
out.
+ retained_bytes: usize,
+ num_groups: usize,
+}
+
+/// A run of consecutive rows of one batch that all belong to `group`.
+#[derive(Debug, Clone, Copy, Default)]
+struct RowRange {
+ group: u32,
+ start: u32,
+ len: u32,
+}
+
+impl RowRange {
+ fn end(&self) -> u32 {
+ self.start + self.len
+ }
+}
+
+/// Below this many rows per range, gathering with `interleave` beats slicing
and concatenating.
+/// Scattered single rows (a high-cardinality `update_batch`) sit at 1 range
per row; merging
+/// partial states back together sits at one range per group per batch, which
is far longer.
+const CONCAT_MIN_RANGE_LEN: usize = 4;
+
+impl CollectListGroupsAccumulator {
+ fn new(datatype: DataType) -> Self {
+ Self {
+ datatype,
+ batches: Vec::new(),
+ ranges: Vec::new(),
+ retained_bytes: 0,
+ num_groups: 0,
+ }
+ }
+
+ /// Appends `row` of the batch being accumulated into `ranges`, extending
the previous range
+ /// when the row continues it.
+ fn push_row(ranges: &mut Vec<RowRange>, group: u32, row: u32) {
+ match ranges.last_mut() {
+ Some(last) if last.group == group && last.end() == row => last.len
+= 1,
+ _ => ranges.push(RowRange {
+ group,
+ start: row,
+ len: 1,
+ }),
+ }
+ }
+
+ fn retain(&mut self, batch: ArrayRef, mut ranges: Vec<RowRange>) {
+ if ranges.is_empty() {
+ return;
+ }
+ // `ranges` is sized for one range per row; run coalescing can leave
most of that unused,
+ // and the excess would be retained for as long as the batch is.
+ if ranges.len() * 2 <= ranges.capacity() {
+ ranges.shrink_to_fit();
+ }
+ self.retained_bytes +=
+ batch.get_array_memory_size() + ranges.capacity() *
size_of::<RowRange>();
+ self.batches.push(batch);
+ self.ranges.push(ranges);
+ }
+
+ fn clear(&mut self) {
+ // `size()` reports capacity rather than length, so release the
buffers outright.
+ self.batches = Vec::new();
+ self.ranges = Vec::new();
+ self.retained_bytes = 0;
+ self.num_groups = 0;
+ }
+
+ /// Drops the contributions of groups `0..emit_groups` and renumbers what
is left to start at
+ /// 0. Batches that no longer contribute anything are released, and
batches only partly
+ /// retained are compacted so their emitted rows stop being pinned.
+ fn compact(&mut self, emit_groups: u32) -> DFResult<()> {
+ let batches = take_field(&mut self.batches);
+ let ranges = take_field(&mut self.ranges);
+ self.retained_bytes = 0;
+ self.num_groups -= emit_groups as usize;
+
+ for (batch, ranges) in batches.into_iter().zip(ranges) {
+ if ranges.iter().all(|r| r.group >= emit_groups) {
+ // Nothing was emitted from this batch, so only the group
numbering changes and
+ // the batch itself can be kept as it is.
+ let mut ranges = ranges;
+ for range in &mut ranges {
+ range.group -= emit_groups;
+ }
+ self.retain(batch, ranges);
+ continue;
+ }
+ let mut retained: Vec<RowRange> = Vec::new();
+ let mut retained_rows: Vec<u32> = Vec::new();
+ for range in ranges {
+ if range.group < emit_groups {
+ continue;
+ }
+ retained.push(RowRange {
+ group: range.group - emit_groups,
+ start: retained_rows.len() as u32,
+ len: range.len,
+ });
+ retained_rows.extend(range.start..range.end());
+ }
+ if retained.is_empty() {
+ continue;
+ }
+ // Compact the batch so the rows this emit released stop being
pinned.
+ let batch = take(batch.as_ref(),
&UInt32Array::from(retained_rows), None)?;
+ self.retain(batch, retained);
+ }
+ Ok(())
+ }
+}
+
+impl GroupsAccumulator for CollectListGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ self.num_groups = self.num_groups.max(total_num_groups);
+ let input = coerce(&values[0], &self.datatype)?;
+ // Spark's collect_list drops null inputs.
+ let nulls = input.logical_nulls();
+ let mut ranges = Vec::with_capacity(group_indices.len());
+ for (row_idx, &group_idx) in group_indices.iter().enumerate() {
+ if let Some(filter) = opt_filter {
+ if filter.is_null(row_idx) || !filter.value(row_idx) {
+ continue;
+ }
+ }
+ if nulls.as_ref().is_some_and(|n| n.is_null(row_idx)) {
+ continue;
+ }
+ Self::push_row(&mut ranges, group_idx as u32, row_idx as u32);
+ }
+ self.retain(input, ranges);
+ Ok(())
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ self.num_groups = self.num_groups.max(total_num_groups);
+ let Some(lists) = values[0].as_list_opt::<i32>() else {
+ return internal_err!(
+ "collect_list expected a List state, got {:?}",
+ values[0].data_type()
+ );
+ };
+ let elements = coerce(lists.values(), &self.datatype)?;
+ let offsets = lists.offsets();
+ let mut ranges = Vec::with_capacity(group_indices.len());
+ for (row_idx, &group_idx) in group_indices.iter().enumerate() {
+ let (start, end) = (offsets[row_idx], offsets[row_idx + 1]);
+ if lists.is_null(row_idx) || end == start {
+ continue;
+ }
+ ranges.push(RowRange {
+ group: group_idx as u32,
+ start: start as u32,
+ len: (end - start) as u32,
+ });
+ }
+ self.retain(elements, ranges);
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ let emit_groups = match emit_to {
+ EmitTo::All => self.num_groups,
+ EmitTo::First(n) => n,
+ };
+
+ // Counting sort of the retained ranges into group order: count,
prefix sum, scatter.
+ // Rows and ranges are counted separately: rows give the output list
offsets, ranges
+ // decide how the values are gathered.
+ let mut range_counts = vec![0u32; emit_groups];
+ let mut row_counts = vec![0u32; emit_groups];
+ let mut num_ranges = 0usize;
+ let mut num_rows = 0usize;
+ for ranges in &self.ranges {
+ for range in ranges {
+ if (range.group as usize) < emit_groups {
+ range_counts[range.group as usize] += 1;
+ row_counts[range.group as usize] += range.len;
+ num_ranges += 1;
+ num_rows += range.len as usize;
+ }
+ }
+ }
+ let offsets =
OffsetBuffer::<i32>::from_lengths(row_counts.iter().map(|c| *c as usize));
+
+ let values = if num_rows == 0 {
+ new_empty_array(&self.datatype)
+ } else if num_rows >= num_ranges * CONCAT_MIN_RANGE_LEN {
+ // Long runs: scatter the ranges themselves into group order, then
copy each one
+ // whole. `concat` preallocates and does a single copy per run.
+ let mut cursors = prefix_sum(&range_counts);
+ let mut ordered = vec![(0u32, RowRange::default()); num_ranges];
+ for (batch_idx, ranges) in self.ranges.iter().enumerate() {
+ for range in ranges {
+ let group = range.group as usize;
+ if group < emit_groups {
+ ordered[cursors[group] as usize] = (batch_idx as u32,
*range);
+ cursors[group] += 1;
+ }
+ }
+ }
+ let slices: Vec<ArrayRef> = ordered
+ .iter()
+ .map(|(batch_idx, range)| {
+ self.batches[*batch_idx as usize]
+ .slice(range.start as usize, range.len as usize)
+ })
+ .collect();
+ let slices: Vec<&dyn Array> = slices.iter().map(|s|
s.as_ref()).collect();
+ concat(&slices)?
+ } else {
+ // Scattered rows: scatter straight into the gather indices and
take them all in one
+ // `interleave` call.
+ let mut cursors = prefix_sum(&row_counts);
+ let mut indices = vec![(0usize, 0usize); num_rows];
+ for (batch_idx, ranges) in self.ranges.iter().enumerate() {
+ for range in ranges {
+ let group = range.group as usize;
+ if group < emit_groups {
+ let mut cursor = cursors[group] as usize;
+ for row in range.start..range.end() {
+ indices[cursor] = (batch_idx, row as usize);
+ cursor += 1;
+ }
+ cursors[group] = cursor as u32;
+ }
+ }
+ }
+ let sources: Vec<&dyn Array> = self.batches.iter().map(|b|
b.as_ref()).collect();
+ interleave(&sources, &indices)?
+ };
+
+ match emit_to {
+ EmitTo::All => self.clear(),
+ EmitTo::First(n) => self.compact(n as u32)?,
+ }
+
+ // A group that collected nothing gets an empty list, not a null one,
matching Spark.
+ Ok(Arc::new(ListArray::new(
+ list_field(&self.datatype),
+ offsets,
+ values,
+ None,
+ )))
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ Ok(vec![self.evaluate(emit_to)?])
+ }
+
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> DFResult<Vec<ArrayRef>> {
+ single_row_lists(&values[0], opt_filter, &self.datatype)
+ }
+
+ fn size(&self) -> usize {
+ self.retained_bytes
+ + self.batches.capacity() * size_of::<ArrayRef>()
+ + self.ranges.capacity() * size_of::<Vec<RowRange>>()
+ + self.datatype.size()
+ }
+}
+
+/// Exclusive prefix sum of `counts`, giving each group the position its first
item is written to.
+fn prefix_sum(counts: &[u32]) -> Vec<u32> {
+ let mut cursors = Vec::with_capacity(counts.len());
+ let mut total = 0u32;
+ for count in counts {
+ cursors.push(total);
+ total += count;
+ }
+ cursors
+}
+
+/// A `NullBuffer` that is null wherever `filter` does not select the row, so
that a filtered row
+/// can be represented as a null list entry. Mirrors DataFusion's
`filter_to_nulls`, which lives
+/// in `datafusion-functions-aggregate-common` and is not re-exported by the
`datafusion` facade.
+fn filter_to_nulls(filter: &BooleanArray) -> NullBuffer {
+ let selected = match filter.nulls() {
+ Some(nulls) => filter.values() & nulls.inner(),
+ None => filter.values().clone(),
+ };
+ NullBuffer::new(selected)
+}
+
+/// Turns every row of `input` into its own single-element list, which is what
+/// `GroupsAccumulator::convert_to_state` has to produce for both collect
functions. Null and
+/// filtered-out rows become null list entries, which `merge_batch` skips, so
the input array
+/// itself can back the lists without being copied.
+fn single_row_lists(
+ input: &ArrayRef,
+ opt_filter: Option<&BooleanArray>,
+ datatype: &DataType,
+) -> DFResult<Vec<ArrayRef>> {
+ let filter_nulls = opt_filter.map(filter_to_nulls);
+ let nulls = NullBuffer::union(filter_nulls.as_ref(),
input.logical_nulls().as_ref());
+ Ok(vec![Arc::new(ListArray::new(
+ list_field(datatype),
+ OffsetBuffer::from_repeated_length(1, input.len()),
+ Arc::clone(input),
+ nulls,
+ ))])
+}
+
+/// One distinct `(group, value)` pair. `start` and `len` address the value's
row-encoded bytes
+/// in [`CollectSetGroupsAccumulator::arena`], and `hash` is of those bytes
alone, so it survives
+/// the group renumbering [`CollectSetGroupsAccumulator::compact`] does.
+///
+/// `start` is 64-bit because the arena holds every group's bytes: a large
execution-memory
+/// budget can push it past 4 GiB while every group's own result, and the
combined `List<i32>`
+/// element count, stay valid, since that limit counts elements rather than
encoded bytes. `len`
+/// stays 32-bit and is range-checked on insert, which keeps the struct at
three machine words.
+/// Field order is chosen so the wider `start` costs nothing:
`size_of::<SetEntry>()` is pinned
+/// by [`tests::set_entry_stays_three_words`].
+#[derive(Debug, Clone, Copy)]
+struct SetEntry {
+ hash: u64,
+ start: u64,
+ group: u32,
+ len: u32,
+}
+
+/// Largest row-encoded value a [`SetEntry`] can address. A single Spark value
cannot approach
+/// this, but the conversion is checked rather than truncating.
+const MAX_ENTRY_LEN: usize = u32::MAX as usize;
+
+/// Narrows an encoded value's length to the 32 bits a [`SetEntry`] stores,
erroring rather than
+/// truncating a length whose low bits would address a shorter, wrong slice.
+fn checked_entry_len(len: usize) -> DFResult<u32> {
+ u32::try_from(len).map_err(|_| {
+ DataFusionError::Execution(format!(
+ "collect_set value of {len} encoded bytes exceeds the
{MAX_ENTRY_LEN} byte limit"
+ ))
+ })
+}
+
+/// Byte range of `entry` within the arena, computed in 64-bit and widened
once.
+///
+/// Split out from [`CollectSetGroupsAccumulator::entry_bytes`] so the
arithmetic can be tested at
+/// the 4 GiB boundary without allocating an arena that large.
+#[inline]
+fn entry_range(start: u64, len: u32) -> Range<usize> {
+ let start = usize::try_from(start).expect("arena offset exceeds the
address space");
+ let end = start
+ .checked_add(len as usize)
+ .expect("arena entry end exceeds the address space");
+ start..end
+}
+
+/// `GroupsAccumulator` for Spark's `collect_set`.
+///
+/// Values are kept row-encoded (`arrow::row`) in one shared arena,
deduplicated on insert against
+/// an open-addressed index keyed by `(group, encoded value)`. Encoding is
what the per-group
+/// `Accumulator` did as well, so which values count as duplicates is
unchanged; what changes is
+/// that a batch is encoded once for all of its groups instead of once per
group, and that the
+/// state is flat rather than one hash table plus one `IndexSet` per group.
+#[derive(Debug)]
+struct CollectSetGroupsAccumulator {
+ /// The element type as declared by the plan.
+ datatype: DataType,
+ converter: RowConverter,
+ /// Row-encoded bytes of every distinct `(group, value)` pair, in
insertion order.
+ arena: Vec<u8>,
+ entries: Vec<SetEntry>,
+ index: RowIndex,
+ num_groups: usize,
+}
+
+impl CollectSetGroupsAccumulator {
+ fn try_new(datatype: DataType) -> DFResult<Self> {
+ let converter =
RowConverter::new(vec![SortField::new(datatype.clone())])?;
+ Ok(Self {
+ datatype,
+ converter,
+ arena: Vec::new(),
+ entries: Vec::new(),
+ index: RowIndex::default(),
+ num_groups: 0,
+ })
+ }
+
+ /// Records `(group, encoded)` unless the group already holds that value.
`hash` is of
+ /// `encoded` alone, so a caller that already has it (`compact`) need not
recompute it.
+ fn insert(&mut self, group: u32, hash: u64, encoded: &[u8]) ->
DFResult<()> {
+ let len = checked_entry_len(encoded.len())?;
+ let Self {
+ arena,
+ entries,
+ index,
+ ..
+ } = self;
+ let key = key_hash(group, hash);
+ let Some(slot) = index.vacant_slot(key, |entry_idx| {
+ let entry = &entries[entry_idx as usize];
+ entry.group == group
+ && entry.hash == hash
+ && &arena[entry_range(entry.start, entry.len)] == encoded
+ }) else {
+ return Ok(());
+ };
+ let entry_idx = entries.len() as u32;
+ entries.push(SetEntry {
+ hash,
+ start: arena.len() as u64,
+ group,
+ len,
+ });
+ arena.extend_from_slice(encoded);
+ index.fill(slot, entry_idx);
+ if index.is_full(entries.len()) {
+ index.grow(entries.iter().map(|e| key_hash(e.group, e.hash)));
+ }
+ Ok(())
+ }
+
+ fn entry_bytes(&self, entry: &SetEntry) -> &[u8] {
+ &self.arena[entry_range(entry.start, entry.len)]
+ }
+
+ /// Drops the entries of groups `0..emit_groups` and renumbers what is
left to start at 0.
+ fn compact(&mut self, emit_groups: u32) -> DFResult<()> {
+ let old_entries = take_field(&mut self.entries);
+ let old_arena = take_field(&mut self.arena);
+ self.index = RowIndex::default();
+ self.num_groups -= emit_groups as usize;
+ for entry in old_entries {
+ if entry.group < emit_groups {
+ continue;
+ }
+ let bytes = &old_arena[entry_range(entry.start, entry.len)];
+ self.insert(entry.group - emit_groups, entry.hash, bytes)?;
+ }
+ Ok(())
+ }
+
+ fn clear(&mut self) {
+ // `size()` reports capacity rather than length, so release the
buffers outright.
+ self.arena = Vec::new();
+ self.entries = Vec::new();
+ self.index = RowIndex::default();
+ self.num_groups = 0;
+ }
+}
+
+impl GroupsAccumulator for CollectSetGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ self.num_groups = self.num_groups.max(total_num_groups);
+ let input = coerce(&values[0], &self.datatype)?;
+ let rows = self.converter.convert_columns(&[Arc::clone(&input)])?;
+ // Spark's collect_set drops null inputs.
+ let nulls = input.logical_nulls();
+ for (row_idx, &group_idx) in group_indices.iter().enumerate() {
+ if let Some(filter) = opt_filter {
+ if filter.is_null(row_idx) || !filter.value(row_idx) {
+ continue;
+ }
+ }
+ if nulls.as_ref().is_some_and(|n| n.is_null(row_idx)) {
+ continue;
+ }
+ let row = rows.row(row_idx);
+ self.insert(group_idx as u32, value_hash(row.as_ref()),
row.as_ref())?;
+ }
+ Ok(())
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> DFResult<()> {
+ self.num_groups = self.num_groups.max(total_num_groups);
+ let Some(lists) = values[0].as_list_opt::<i32>() else {
+ return internal_err!(
+ "collect_set expected a List state, got {:?}",
+ values[0].data_type()
+ );
+ };
+ let elements = coerce(lists.values(), &self.datatype)?;
+ let rows = self.converter.convert_columns(&[elements])?;
+ let offsets = lists.offsets();
+ for (row_idx, &group_idx) in group_indices.iter().enumerate() {
+ if lists.is_null(row_idx) {
+ continue;
+ }
+ for pos in offsets[row_idx] as usize..offsets[row_idx + 1] as
usize {
+ let row = rows.row(pos);
+ self.insert(group_idx as u32, value_hash(row.as_ref()),
row.as_ref())?;
+ }
+ }
+ Ok(())
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> DFResult<ArrayRef> {
+ let emit_groups = match emit_to {
+ EmitTo::All => self.num_groups,
+ EmitTo::First(n) => n,
+ };
+
+ // Counting sort of the entries into group order: count, prefix sum,
scatter.
+ let mut counts = vec![0u32; emit_groups];
+ let mut total = 0u32;
+ for entry in &self.entries {
+ if (entry.group as usize) < emit_groups {
+ counts[entry.group as usize] += 1;
+ total += 1;
+ }
+ }
+ let offsets = OffsetBuffer::<i32>::from_lengths(counts.iter().map(|c|
*c as usize));
+ let mut cursors = prefix_sum(&counts);
+
+ let values = if total == 0 {
+ new_empty_array(&self.datatype)
+ } else {
+ let mut ordered = vec![0u32; total as usize];
+ for (entry_idx, entry) in self.entries.iter().enumerate() {
+ let group = entry.group as usize;
+ if group < emit_groups {
+ ordered[cursors[group] as usize] = entry_idx as u32;
+ cursors[group] += 1;
+ }
+ }
+ let parser = self.converter.parser();
+ let decoded = self.converter.convert_rows(
+ ordered
+ .iter()
+ .map(|&i| parser.parse(self.entry_bytes(&self.entries[i as
usize]))),
+ )?;
+ let decoded = decoded
+ .into_iter()
+ .next()
+ .ok_or_else(|| internal_datafusion_err!("collect_set decoded
no columns"))?;
+ // `RowConverter` always decodes to the physical type, so a
dictionary element type
+ // has to be restored.
+ coerce(&decoded, &self.datatype)?
+ };
+
+ match emit_to {
+ EmitTo::All => self.clear(),
+ EmitTo::First(n) => self.compact(n as u32)?,
+ }
+
+ Ok(Arc::new(ListArray::new(
+ list_field(&self.datatype),
+ offsets,
+ values,
+ None,
+ )))
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> DFResult<Vec<ArrayRef>> {
+ Ok(vec![self.evaluate(emit_to)?])
+ }
+
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> DFResult<Vec<ArrayRef>> {
+ // A single-element list is trivially deduplicated; `merge_batch` does
the real work.
+ single_row_lists(&values[0], opt_filter, &self.datatype)
+ }
+
+ fn size(&self) -> usize {
+ self.arena.capacity()
+ + self.entries.capacity() * size_of::<SetEntry>()
+ + self.index.size()
+ + self.datatype.size()
+ }
+}
+
+/// Hash of a row-encoded value, independent of the group that holds it.
+fn value_hash(encoded: &[u8]) -> u64 {
+ XxHash64::oneshot(0, encoded)
+}
+
+/// Mixes a group index into a value hash. Keeping the two separate lets an
entry's value hash be
+/// reused when [`CollectSetGroupsAccumulator::compact`] renumbers its group.
+fn key_hash(group: u32, value_hash: u64) -> u64 {
+ value_hash ^ (group as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
+}
+
+/// Open-addressed index from a hash to an entry position in
+/// [`CollectSetGroupsAccumulator::entries`], with linear probing and a load
factor of 0.5.
+///
+/// Slots hold `entry index + 1` so that `0` means empty.
+#[derive(Debug, Default)]
+struct RowIndex {
+ slots: Vec<u32>,
+}
+
+impl RowIndex {
+ const INITIAL_SLOTS: usize = 1024;
+
+ /// The slot `hash` belongs in, or `None` when `eq` matched an entry
already indexed there.
+ fn vacant_slot(&mut self, hash: u64, eq: impl Fn(u32) -> bool) ->
Option<usize> {
+ if self.slots.is_empty() {
+ self.slots = vec![0; Self::INITIAL_SLOTS];
+ }
+ let mask = self.slots.len() - 1;
+ let mut slot = hash as usize & mask;
+ loop {
+ match self.slots[slot] {
+ 0 => return Some(slot),
+ v if eq(v - 1) => return None,
+ _ => slot = (slot + 1) & mask,
+ }
+ }
+ }
+
+ fn fill(&mut self, slot: usize, entry_idx: u32) {
+ self.slots[slot] = entry_idx + 1;
+ }
+
+ fn is_full(&self, num_entries: usize) -> bool {
+ num_entries * 2 >= self.slots.len()
+ }
+
+ fn grow(&mut self, hashes: impl Iterator<Item = u64>) {
+ let mask = self.slots.len() * 2 - 1;
+ let mut slots = vec![0u32; mask + 1];
+ for (entry_idx, hash) in hashes.enumerate() {
+ let mut slot = hash as usize & mask;
+ while slots[slot] != 0 {
+ slot = (slot + 1) & mask;
+ }
+ slots[slot] = entry_idx as u32 + 1;
+ }
+ self.slots = slots;
+ }
+
+ fn size(&self) -> usize {
+ self.slots.capacity() * size_of::<u32>()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{
+ Int32Array, Int32Builder, ListBuilder, StringArray, StringBuilder,
StructArray,
+ };
+ use arrow::datatypes::{Fields, Int32Type};
+
+ /// The collected values of every group, which are never null lists.
+ fn int_groups(array: &ArrayRef) -> Vec<Vec<Option<i32>>> {
+ let list = array.as_list::<i32>();
+ assert_eq!(list.null_count(), 0, "collected lists are never null");
+ (0..list.len())
+ .map(|i| {
+ let values = list.value(i);
+ let values = values.as_primitive::<Int32Type>();
+ (0..values.len())
+ .map(|j| (!values.is_null(j)).then(|| values.value(j)))
+ .collect()
+ })
+ .collect()
+ }
+
+ fn string_groups(array: &ArrayRef) -> Vec<Vec<String>> {
+ let list = array.as_list::<i32>();
+ (0..list.len())
+ .map(|i| {
+ let values = list.value(i);
+ let values = values.as_string::<i32>();
+ (0..values.len())
+ .map(|j| values.value(j).to_string())
+ .collect()
+ })
+ .collect()
+ }
+
+ fn ints(values: Vec<Option<i32>>) -> ArrayRef {
+ Arc::new(Int32Array::from(values))
+ }
+
+ /// A `List<Int32>` state array, one list per row.
+ fn int_lists(rows: Vec<Option<Vec<i32>>>) -> ArrayRef {
+ let mut builder = ListBuilder::new(Int32Builder::new())
+ .with_field(Field::new_list_field(DataType::Int32, true));
+ for row in rows {
+ match row {
+ Some(values) => {
+ for value in values {
+ builder.values().append_value(value);
+ }
+ builder.append(true);
+ }
+ None => builder.append(false),
+ }
+ }
+ Arc::new(builder.finish())
+ }
+
+ fn collect_list() -> CollectListGroupsAccumulator {
+ CollectListGroupsAccumulator::new(DataType::Int32)
+ }
+
+ fn collect_set() -> CollectSetGroupsAccumulator {
+ CollectSetGroupsAccumulator::try_new(DataType::Int32).unwrap()
+ }
+
+ #[test]
+ fn collect_list_keeps_input_order_per_group_and_drops_nulls() {
+ let mut acc = collect_list();
+ acc.update_batch(
+ &[ints(vec![Some(1), None, Some(2), Some(3), Some(4)])],
+ &[0, 1, 0, 2, 1],
+ None,
+ 3,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(1), Some(2)], vec![Some(4)], vec![Some(3)],]
+ );
+ }
+
+ #[test]
+ fn collect_list_emits_empty_lists_for_groups_that_collected_nothing() {
+ let mut acc = collect_list();
+ // Group 0 only ever sees a null, group 2 is never seen at all.
+ acc.update_batch(&[ints(vec![None, Some(7)])], &[0, 1], None, 3)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(int_groups(&emitted), vec![vec![], vec![Some(7)], vec![]]);
+ }
+
+ #[test]
+ fn collect_list_honours_the_filter() {
+ let mut acc = collect_list();
+ let filter = BooleanArray::from(vec![Some(true), Some(false), None,
Some(true)]);
+ acc.update_batch(
+ &[ints(vec![Some(1), Some(2), Some(3), Some(4)])],
+ &[0, 0, 0, 1],
+ Some(&filter),
+ 2,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(int_groups(&emitted), vec![vec![Some(1)], vec![Some(4)]]);
+ }
+
+ #[test]
+ fn collect_list_merges_partial_states() {
+ let mut acc = collect_list();
+ acc.merge_batch(
+ &[int_lists(vec![Some(vec![1, 2]), Some(vec![3])])],
+ &[0, 1],
+ 2,
+ )
+ .unwrap();
+ acc.merge_batch(
+ &[int_lists(vec![Some(vec![4]), None, Some(vec![])])],
+ &[1, 0, 2],
+ 3,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(1), Some(2)], vec![Some(3), Some(4)], vec![]]
+ );
+ }
+
+ /// `evaluate` picks its gather by average run length: two runs of 32 take
the `concat`
+ /// path, 64 alternating rows the `interleave` one. Both have to collect
the right values.
+ #[test]
+ fn collect_list_gathers_runs_and_scattered_rows_alike() {
+ let values = || ints((0..64).map(Some).collect());
+
+ let mut clustered = collect_list();
+ let groups: Vec<usize> = (0..64).map(|i| i / 32).collect();
+ clustered
+ .update_batch(&[values()], &groups, None, 2)
+ .unwrap();
+ assert_eq!(
+ int_groups(&clustered.evaluate(EmitTo::All).unwrap()),
+ vec![
+ (0..32).map(Some).collect::<Vec<_>>(),
+ (32..64).map(Some).collect::<Vec<_>>(),
+ ]
+ );
+
+ let mut scattered = collect_list();
+ let groups: Vec<usize> = (0..64).map(|i| i % 2).collect();
+ scattered
+ .update_batch(&[values()], &groups, None, 2)
+ .unwrap();
+ assert_eq!(
+ int_groups(&scattered.evaluate(EmitTo::All).unwrap()),
+ vec![
+ (0..64).step_by(2).map(Some).collect::<Vec<_>>(),
+ (1..64).step_by(2).map(Some).collect::<Vec<_>>(),
+ ]
+ );
+ }
+
+ #[test]
+ fn collect_list_emit_first_keeps_the_remaining_groups() {
+ let mut acc = collect_list();
+ acc.update_batch(
+ &[ints(vec![Some(1), Some(2), Some(3), Some(4)])],
+ &[0, 1, 2, 0],
+ None,
+ 3,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::First(2)).unwrap();
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(1), Some(4)], vec![Some(2)]]
+ );
+
+ // Group 2 is renumbered to 0, and a further batch lands on the
renumbered groups.
+ acc.update_batch(&[ints(vec![Some(5), Some(6)])], &[0, 1], None, 2)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(3), Some(5)], vec![Some(6)]]
+ );
+ }
+
+ #[test]
+ fn collect_list_converts_rows_to_single_element_states() {
+ let acc = collect_list();
+ let filter = BooleanArray::from(vec![true, true, false]);
+ let states = acc
+ .convert_to_state(&[ints(vec![Some(1), None, Some(3)])],
Some(&filter))
+ .unwrap();
+ let mut merged = collect_list();
+ merged.merge_batch(&states, &[0, 0, 0], 1).unwrap();
+ assert_eq!(
+ int_groups(&merged.evaluate(EmitTo::All).unwrap()),
+ vec![vec![Some(1)]]
+ );
+ }
+
+ #[test]
+ fn collect_set_deduplicates_per_group_and_drops_nulls() {
+ let mut acc = collect_set();
+ acc.update_batch(
+ &[ints(vec![Some(1), Some(1), None, Some(2), Some(1)])],
+ &[0, 0, 0, 0, 1],
+ None,
+ 2,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ // The same value in a different group is a different entry.
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(1), Some(2)], vec![Some(1)]]
+ );
+ }
+
+ #[test]
+ fn collect_set_emits_empty_lists_for_groups_that_collected_nothing() {
+ let mut acc = collect_set();
+ acc.update_batch(&[ints(vec![None, Some(7)])], &[0, 1], None, 3)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(int_groups(&emitted), vec![vec![], vec![Some(7)], vec![]]);
+ }
+
+ #[test]
+ fn collect_set_deduplicates_across_merged_states() {
+ let mut acc = collect_set();
+ acc.merge_batch(
+ &[int_lists(vec![Some(vec![1, 2]), Some(vec![3])])],
+ &[0, 1],
+ 2,
+ )
+ .unwrap();
+ acc.merge_batch(&[int_lists(vec![Some(vec![2, 4]), None])], &[0, 1], 2)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(
+ int_groups(&emitted),
+ vec![vec![Some(1), Some(2), Some(4)], vec![Some(3)]]
+ );
+ }
+
+ #[test]
+ fn collect_set_emit_first_keeps_the_remaining_groups() {
+ let mut acc = collect_set();
+ acc.update_batch(
+ &[ints(vec![Some(1), Some(2), Some(3), Some(1)])],
+ &[0, 1, 2, 0],
+ None,
+ 3,
+ )
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::First(2)).unwrap();
+ assert_eq!(int_groups(&emitted), vec![vec![Some(1)], vec![Some(2)]]);
+
+ acc.update_batch(&[ints(vec![Some(3), Some(5)])], &[0, 0], None, 1)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ // Group 2 became group 0 and already held 3, so the repeat is dropped.
+ assert_eq!(int_groups(&emitted), vec![vec![Some(3), Some(5)]]);
+ }
+
+ #[test]
+ fn collect_set_grows_its_index_past_the_initial_capacity() {
+ let mut acc = collect_set();
+ let values: Vec<Option<i32>> = (0..4096).map(Some).collect();
+ let groups: Vec<usize> = (0..4096).map(|i| i % 4).collect();
+ acc.update_batch(&[ints(values)], &groups, None, 4).unwrap();
+ // Every value is distinct, so nothing is deduplicated away.
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ let collected = int_groups(&emitted);
+ assert_eq!(collected.iter().map(Vec::len).sum::<usize>(), 4096);
+ assert_eq!(collected[0][..2], [Some(0), Some(4)]);
+ }
+
+ #[test]
+ fn collect_set_deduplicates_strings() {
+ let mut acc =
CollectSetGroupsAccumulator::try_new(DataType::Utf8).unwrap();
+ let values: ArrayRef = Arc::new(StringArray::from(vec![
+ Some("a"),
+ Some("b"),
+ Some("a"),
+ None,
+ Some("b"),
+ ]));
+ acc.update_batch(&[values], &[0, 0, 0, 0, 1], None, 2)
+ .unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ assert_eq!(
+ string_groups(&emitted),
+ vec![
+ vec!["a".to_string(), "b".to_string()],
+ vec!["b".to_string()]
+ ]
+ );
+ }
+
+ #[test]
+ fn collect_set_deduplicates_structs() {
+ let fields: Fields = vec![
+ Arc::new(Field::new("a", DataType::Utf8, true)),
+ Arc::new(Field::new("b", DataType::Int32, true)),
+ ]
+ .into();
+ let mut names = StringBuilder::new();
+ let mut numbers = Int32Builder::new();
+ for (name, number) in [("x", 1), ("x", 1), ("x", 2)] {
+ names.append_value(name);
+ numbers.append_value(number);
+ }
+ let values: ArrayRef = Arc::new(StructArray::new(
+ fields.clone(),
+ vec![Arc::new(names.finish()), Arc::new(numbers.finish())],
+ None,
+ ));
+
+ let mut acc =
+
CollectSetGroupsAccumulator::try_new(DataType::Struct(fields.clone())).unwrap();
+ acc.update_batch(&[values], &[0, 0, 0], None, 1).unwrap();
+ let emitted = acc.evaluate(EmitTo::All).unwrap();
+ let list = emitted.as_list::<i32>();
+ assert_eq!(list.len(), 1);
+ // `{x, 1}` collapses into one entry, `{x, 2}` stays separate.
+ assert_eq!(list.value(0).len(), 2);
+ assert_eq!(list.value(0).data_type(), &DataType::Struct(fields));
+ }
+
+ /// The two UDFs forward the rest of `AggregateUDFImpl` to
`datafusion_spark`, and the
+ /// planner and the `GroupsAccumulator` both depend on what those forwards
report.
+ /// The arena addresses every group's bytes at once, so its offsets have
to stay valid past
+ /// 4 GiB. The old 32-bit `start`/`len` pair wrapped there: `0xFFFF_FFFB +
8` came back as 3,
+ /// and the resulting slice either panicked or read another entry's bytes.
Nothing about a
+ /// group's own result or the combined `List<i32>` element count bounds
the arena, because that
+ /// limit counts elements rather than encoded bytes.
+ ///
+ /// Checked here on the range arithmetic rather than through a real
accumulator, so the
+ /// boundary is covered without allocating a 4 GiB arena.
+ #[test]
+ fn entry_range_stays_valid_past_four_gibibytes() {
+ const FOUR_GIB: u64 = u32::MAX as u64 + 1;
+
+ // Straddling the boundary: the wrapping version produced 3..3 for
this pair.
+ assert_eq!(
+ entry_range(u32::MAX as u64 - 4, 8),
+ 4_294_967_291..4_294_967_299
+ );
+ // Wholly above it.
+ assert_eq!(entry_range(FOUR_GIB, 8), 4_294_967_296..4_294_967_304);
+ // A later start is no longer truncated back into the first 4 GiB.
+ assert_eq!(entry_range(FOUR_GIB + 16, 0).start, 4_294_967_312);
+ // Ordinary offsets are unaffected.
+ assert_eq!(entry_range(0, 0), 0..0);
+ assert_eq!(entry_range(24, 8), 24..32);
+ }
+
+ /// Widening `start` to 64 bits is only free because of the field order.
If someone reorders
+ /// `SetEntry`, padding appears and every distinct value in the aggregate
pays 8 more bytes,
+ /// which `size()` reports to the memory pool.
+ #[test]
+ fn set_entry_stays_three_words() {
+ assert_eq!(size_of::<SetEntry>(), 3 * size_of::<u64>());
+ }
+
+ /// `len` stays 32-bit, so it is the one narrowing conversion left.
Reaching it through the
+ /// accumulator would need a single row-encoded value over 4 GiB, so the
conversion is checked
+ /// on its own: it must error rather than truncate a length into a shorter
slice that would
+ /// silently return part of a value.
+ #[test]
+ fn oversized_entry_len_is_rejected_rather_than_truncated() {
+ assert_eq!(checked_entry_len(0).unwrap(), 0);
+ assert_eq!(checked_entry_len(24).unwrap(), 24);
+ assert_eq!(checked_entry_len(MAX_ENTRY_LEN).unwrap(), u32::MAX);
+ let err = checked_entry_len(MAX_ENTRY_LEN + 1).unwrap_err();
+ assert!(err.to_string().contains("exceeds"), "{err}");
+ }
+
+ /// The ordinary path still round-trips through the widened representation.
+ #[test]
+ fn entry_bytes_round_trip_through_the_arena() {
+ let mut acc =
CollectSetGroupsAccumulator::try_new(DataType::Int32).unwrap();
+ acc.insert(0, 1, &[1, 2, 3]).unwrap();
+ acc.insert(1, 2, &[4, 5]).unwrap();
+ acc.insert(0, 1, &[1, 2, 3]).unwrap(); // duplicate, not stored again
+ assert_eq!(acc.entries.len(), 2);
+ assert_eq!(acc.entry_bytes(&acc.entries[0]), &[1, 2, 3]);
+ assert_eq!(acc.entry_bytes(&acc.entries[1]), &[4, 5]);
+ assert_eq!(acc.entries[1].start, 3);
+ }
+
+ #[test]
+ fn collect_functions_forward_their_shape_to_datafusion_spark() {
+ for udf in [
+ Box::new(CometCollectList::new()) as Box<dyn AggregateUDFImpl>,
+ Box::new(CometCollectSet::new()),
+ ] {
+ let name = udf.name().to_string();
+ assert!(name == "collect_list" || name == "collect_set", "{name}");
+ assert_eq!(
+ udf.return_type(&[DataType::Int32]).unwrap(),
+ DataType::List(list_field(&DataType::Int32))
+ );
+ }
+ }
+}
diff --git a/native/spark-expr/src/agg_funcs/mod.rs
b/native/spark-expr/src/agg_funcs/mod.rs
index d0d8083195..fcbe465c78 100644
--- a/native/spark-expr/src/agg_funcs/mod.rs
+++ b/native/spark-expr/src/agg_funcs/mod.rs
@@ -18,6 +18,7 @@
mod approx_percentile;
mod avg;
mod avg_decimal;
+mod collect;
mod correlation;
mod covariance;
mod hll_plus_plus;
@@ -36,6 +37,7 @@ mod welford;
pub use approx_percentile::ApproxPercentile;
pub use avg::Avg;
pub use avg_decimal::AvgDecimal;
+pub use collect::{CometCollectList, CometCollectSet};
pub use correlation::Correlation;
pub use covariance::Covariance;
pub use hll_plus_plus::{hllpp_precision, HllPlusPlus};
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 3592b3b362..23c6b53509 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -185,6 +185,29 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
classOf[LocalTableScanExec])
}
+ test("grouped collect_list/collect_set over nulls, duplicates and several
batches") {
+ // Grouped collect_list/collect_set are served by a native
GroupsAccumulator rather than one
+ // boxed accumulator per group, so the group's identity travels with each
row instead of being
+ // implied by which accumulator was called. This covers the cases where
that bookkeeping shows:
+ // a group whose every input is NULL (Spark returns an empty array, not
NULL), repeated values
+ // (collect_list keeps every copy, collect_set one), and enough rows that
a group is fed by
+ // several batches and, after the exchange, by several partial states.
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "128") {
+ val data = (0 until 2000).map { i =>
+ val group = i % 37
+ (group, if (group == 5) None else Some(i % 11))
+ }
+ withParquetTable(data, "tbl") {
+ checkSparkAnswerAndOperator(sql("""
+ SELECT _1,
+ sort_array(collect_list(_2)),
+ sort_array(collect_set(_2)),
+ size(collect_list(_2))
+ FROM tbl GROUP BY _1"""))
+ }
+ }
+ }
+
test("collect_list/collect_set combined with distinct aggregate falls back
safely") {
// SPARK-17616: a distinct aggregate combined with
collect_list/collect_set produces a
// multi-stage plan where the buffer-producing Partial may run in Spark
(e.g. over a
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
