perf: add a GroupsAccumulator for array_agg(DISTINCT)

datafusion/core/tests/sql/aggregates/nested_nullability.rs

@@ -235,7 +235,8 @@ async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> {
 #[tokio::test]
 async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> {
     AggregateBatchesTest::new()
-        .with_num_rows(10_000)
+        // Enough groups that even the compact native accumulator exceeds the limit
+        .with_num_rows(100_000)
         // One partition keeps the native aggregate's memory pressure deterministic.
         .with_target_partitions(1)
         .without_single_distinct_to_groupby()

datafusion/functions-aggregate/benches/array_agg.rs

@@ -25,10 +25,11 @@ use arrow::array::{
 use arrow::datatypes::{DataType, Field, FieldRef, Int64Type, Schema};
 use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use datafusion_expr::function::AccumulatorArgs;
-use datafusion_expr::{Accumulator, AggregateUDFImpl};
+use datafusion_expr::{Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator};
 use datafusion_functions_aggregate::array_agg::{
     ArrayAgg, ArrayAggAccumulator, DistinctArrayAggAccumulator,
 };
+use datafusion_functions_aggregate_common::aggregate::groups_accumulator::GroupsAccumulatorAdapter;
 use datafusion_physical_expr::{PhysicalSortExpr, expressions::col};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 
@@ -323,6 +324,113 @@ fn distinct_array_agg_benchmark(c: &mut Criterion) {
     );
 }
 
+/// `ARRAY_AGG(DISTINCT x) ... GROUP BY g` over one 8 192-row Utf8 batch, with
+/// the groups accumulator against the previous per-group path
+/// (`GroupsAccumulatorAdapter` over one `DistinctArrayAggAccumulator` per group),
+/// across group counts, distinct cardinalities and partial emission.
+fn distinct_array_agg_groups_benchmark(c: &mut Criterion) {
+    let schema = Schema::new(vec![Field::new("value", DataType::Utf8, true)]);
+    let value_expr = col("value", &schema).unwrap();
+    let value_field = value_expr.return_field(&schema).unwrap();
+    let return_field: FieldRef = Field::new(
+        "array_agg",
+        DataType::List(Field::new_list_field(DataType::Utf8, true).into()),
+        true,
+    )
+    .into();
+    let args = AccumulatorArgs {
+        return_field,
+        schema: &schema,
+        expr_fields: std::slice::from_ref(&value_field),
+        ignore_nulls: false,
+        order_bys: &[],
+        is_reversed: false,
+        name: "array_agg(DISTINCT value)",
+        is_distinct: true,
+        exprs: std::slice::from_ref(&value_expr),
+    };
+    let array_agg = ArrayAgg::default();
+
+    let inputs = [
+        (
+            "low cardinality (~25 distinct)",
+            Arc::new(create_string_array_low_cardinality(8192, 0.0)) as ArrayRef,
+        ),
+        (
+            "high cardinality (~7800 distinct)",
+            Arc::new(create_string_array_high_cardinality(8192, 0.05, 0.0)) as ArrayRef,
+        ),
+    ];
+    for num_groups in [100, 1000, 8192] {
+        let mut rng = StdRng::seed_from_u64(42);
+        let group_indices = (0..8192)
+            .map(|_| rng.random_range(0..num_groups))
+            .collect::<Vec<usize>>();
+        for (input_name, values) in &inputs {
+            // Emits every group at once, or in eight `EmitTo::First` chunks, which
+            // makes the accumulator rebuild its state after each partial emission.
+            for emit_chunks in [1, 8] {
+                let run = |acc: &mut dyn GroupsAccumulator| {
+                    acc.update_batch(
+                        std::slice::from_ref(values),
+                        &group_indices,
+                        None,
+                        num_groups,
+                    )
+                    .unwrap();
+                    if emit_chunks == 1 {
+                        return vec![acc.evaluate(EmitTo::All).unwrap()];
+                    }
+                    let chunk = num_groups.div_ceil(emit_chunks);
+                    let mut remaining = num_groups;
+                    let mut emitted = Vec::with_capacity(emit_chunks);
+                    while remaining > 0 {
+                        let n = chunk.min(remaining);
+                        emitted.push(acc.evaluate(EmitTo::First(n)).unwrap());
+                        remaining -= n;
+                    }
+                    emitted
+                };
+                let emit = if emit_chunks == 1 {
+                    "emit all".to_string()
+                } else {
+                    format!("emit in {emit_chunks} chunks")
+                };
+                c.bench_function(
+                    &format!(
+                        "distinct_array_agg groups {num_groups} groups, {input_name}, {emit}"
+                    ),
+                    |b| {
+                        b.iter(|| {
+                            let mut acc =
+                                array_agg.create_groups_accumulator(args.clone()).unwrap();
+                            run(acc.as_mut())
+                        })
+                    },
+                );
+                c.bench_function(
+                    &format!(
+                        "distinct_array_agg adapter {num_groups} groups, {input_name}, {emit}"
+                    ),
+                    |b| {
+                        b.iter(|| {
+                            let factory = || {
+                                Ok(Box::new(DistinctArrayAggAccumulator::try_new(
+                                    &DataType::Utf8,
+                                    None,
+                                    false,
+                                )?) as Box<dyn Accumulator>)
+                            };
+                            let mut acc = GroupsAccumulatorAdapter::new(factory);
+                            run(&mut acc)
+                        })
+                    },
+                );
+            }
+        }
+    }
+}
+
 /// Precomputes the schema, physical expressions, sort expression, and aggregate
 /// metadata so each benchmark iteration can focus on accumulator operations.
 struct OrderedArrayAggBenchFixture {
@@ -545,6 +653,7 @@ criterion_group!(
     benches,
     array_agg_benchmark,
     distinct_array_agg_benchmark,
+    distinct_array_agg_groups_benchmark,
     ordered_array_agg_benchmark,
     ordered_array_agg_wide_utf8_benchmark
 );