perf: Implement groups accumulator count distinct primitive types

datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs

@@ -17,10 +17,12 @@
 
 mod bytes;
 mod dict;
+mod groups;
 mod native;
 
 pub use bytes::BytesDistinctCountAccumulator;
 pub use bytes::BytesViewDistinctCountAccumulator;
 pub use dict::DictionaryCountAccumulator;
+pub use groups::PrimitiveDistinctCountGroupsAccumulator;
 pub use native::FloatDistinctCountAccumulator;
 pub use native::PrimitiveDistinctCountAccumulator;

datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs

@@ -0,0 +1,182 @@
+// 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.
+
+use arrow::array::{
+    ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray,
+};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::{ArrowPrimitiveType, Field};
+use datafusion_common::HashSet;
+use datafusion_common::hash_utils::RandomState;
+use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
+use std::hash::Hash;
+use std::mem::size_of;
+use std::sync::Arc;
+
+use crate::aggregate::groups_accumulator::accumulate::accumulate;
+
+pub struct PrimitiveDistinctCountGroupsAccumulator<T: ArrowPrimitiveType>
+where
+    T::Native: Eq + Hash,
+{
+    seen: HashSet<(usize, T::Native), RandomState>,
+    num_groups: usize,
+}
+
+impl<T: ArrowPrimitiveType> PrimitiveDistinctCountGroupsAccumulator<T>
+where
+    T::Native: Eq + Hash,
+{
+    pub fn new() -> Self {
+        Self {
+            seen: HashSet::default(),
+            num_groups: 0,
+        }
+    }
+}
+
+impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
+where
+    T::Native: Eq + Hash,
+{
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: ArrowPrimitiveType + Send + std::fmt::Debug> GroupsAccumulator
+    for PrimitiveDistinctCountGroupsAccumulator<T>
+where
+    T::Native: Eq + Hash,
+{
+    fn update_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion_common::Result<()> {
+        debug_assert_eq!(values.len(), 1);
+        self.num_groups = self.num_groups.max(total_num_groups);
+        let arr = values[0].as_primitive::<T>();
+        accumulate(group_indices, arr, opt_filter, |group_idx, value| {
+            self.seen.insert((group_idx, value));
+        });
+        Ok(())
+    }
+
+    fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
+        let num_emitted = match emit_to {
+            EmitTo::All => self.num_groups,
+            EmitTo::First(n) => n,
+        };
+
+        let mut counts = vec![0i64; num_emitted];
+
+        if matches!(emit_to, EmitTo::All) {
+            for &(group_idx, _) in self.seen.iter() {
+                counts[group_idx] += 1;
+            }
+            self.seen.clear();
+            self.num_groups = 0;
+        } else {
+            let mut remaining = HashSet::default();
+            for (group_idx, value) in self.seen.drain() {
+                if group_idx < num_emitted {
+                    counts[group_idx] += 1;
+                } else {
+                    remaining.insert((group_idx - num_emitted, value));
+                }
+            }
+            self.seen = remaining;
+            self.num_groups = self.num_groups.saturating_sub(num_emitted);
+        }
+
+        Ok(Arc::new(Int64Array::from(counts)))
+    }
+
+    fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
+        let num_emitted = match emit_to {
+            EmitTo::All => self.num_groups,
+            EmitTo::First(n) => n,
+        };
+
+        let mut group_values: Vec<Vec<T::Native>> = vec![Vec::new(); num_emitted];
+
+        if matches!(emit_to, EmitTo::All) {
+            for (group_idx, value) in self.seen.drain() {
+                group_values[group_idx].push(value);
+            }
+            self.num_groups = 0;
+        } else {
+            let mut remaining = HashSet::default();
+            for (group_idx, value) in self.seen.drain() {
+                if group_idx < num_emitted {
+                    group_values[group_idx].push(value);
+                } else {
+                    remaining.insert((group_idx - num_emitted, value));
+                }
+            }
+            self.seen = remaining;
+            self.num_groups = self.num_groups.saturating_sub(num_emitted);
+        }
+
+        let mut offsets = vec![0i32];
+        let mut all_values = Vec::new();
+        for values in &group_values {
+            all_values.extend(values.iter().copied());
+            offsets.push(all_values.len() as i32);
+        }
+
+        let values_array = Arc::new(PrimitiveArray::<T>::from_iter_values(all_values));
+        let list_array = ListArray::new(
+            Arc::new(Field::new_list_field(T::DATA_TYPE, true)),
+            OffsetBuffer::new(offsets.into()),
+            values_array,
+            None,
+        );
+
+        Ok(vec![Arc::new(list_array)])
+    }
+
+    fn merge_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        _opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion_common::Result<()> {
+        debug_assert_eq!(values.len(), 1);
+        self.num_groups = self.num_groups.max(total_num_groups);
+        let list_array = values[0].as_list::<i32>();
+
+        for (row_idx, group_idx) in group_indices.iter().enumerate() {
+            let inner = list_array.value(row_idx);
+            let inner_arr = inner.as_primitive::<T>();
+            for value in inner_arr.values().iter() {
+                self.seen.insert((*group_idx, *value));
+            }
+        }
+
+        Ok(())
+    }
+
+    fn size(&self) -> usize {
+        size_of::<Self>()
+            + self.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::<u64>())
+    }
+}

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -23,7 +23,7 @@ use arrow::array::{
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{Criterion, criterion_group, criterion_main};
 use datafusion_expr::function::AccumulatorArgs;
-use datafusion_expr::{Accumulator, AggregateUDFImpl};
+use datafusion_expr::{Accumulator, AggregateUDFImpl, EmitTo};
 use datafusion_functions_aggregate::count::Count;
 use datafusion_physical_expr::expressions::col;
 use rand::rngs::StdRng;
@@ -87,6 +87,37 @@ fn create_i16_array(n_distinct: usize) -> Int16Array {
         .collect()
 }
 
+fn create_group_indices(num_groups: usize) -> Vec<usize> {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| rng.random_range(0..num_groups))
+        .collect()
+}
+
+fn prepare_args(data_type: DataType) -> (Arc<Schema>, AccumulatorArgs<'static>) {
+    let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)]));
+    let schema_leaked: &'static Schema = Box::leak(Box::new((*schema).clone()));
+    let expr = col("f", schema_leaked).unwrap();
+    let expr_leaked: &'static _ = Box::leak(Box::new(expr));
+    let return_field: Arc<Field> = Field::new("f", DataType::Int64, true).into();
+    let return_field_leaked: &'static _ = Box::leak(Box::new(return_field.clone()));
+    let expr_field = expr_leaked.return_field(schema_leaked).unwrap();
+    let expr_field_leaked: &'static _ = Box::leak(Box::new(expr_field));
+
+    let accumulator_args = AccumulatorArgs {
+        return_field: return_field_leaked.clone(),
+        schema: schema_leaked,
+        expr_fields: std::slice::from_ref(expr_field_leaked),
+        ignore_nulls: false,
+        order_bys: &[],
+        is_reversed: false,
+        name: "count(distinct f)",
+        is_distinct: true,
+        exprs: std::slice::from_ref(expr_leaked),
+    };
+    (schema, accumulator_args)
+}
+
 fn count_distinct_benchmark(c: &mut Criterion) {
     for pct in [80, 99] {
         let n_distinct = BATCH_SIZE * pct / 100;
@@ -150,5 +181,69 @@ fn count_distinct_benchmark(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, count_distinct_benchmark);
+fn count_distinct_groups_benchmark(c: &mut Criterion) {
+    let count_fn = Count::new();
+
+    for num_groups in [10, 100, 1000] {
+        let n_distinct = BATCH_SIZE * 80 / 100;
+        let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
+        let group_indices = create_group_indices(num_groups);
+
+        let (_schema, args) = prepare_args(DataType::Int64);
+
+        if count_fn.groups_accumulator_supported(args.clone()) {
+            c.bench_function(
+                &format!("count_distinct_groups i64 {num_groups} groups"),
+                |b| {
+                    b.iter(|| {
+                        let (_schema, args) = prepare_args(DataType::Int64);
+                        let mut acc = count_fn.create_groups_accumulator(args).unwrap();
+                        acc.update_batch(
+                            std::slice::from_ref(&values),
+                            &group_indices,
+                            None,
+                            num_groups,
+                        )
+                        .unwrap();
+                        acc.evaluate(EmitTo::All).unwrap()
+                    })
+                },
+            );
+        } else {
+            c.bench_function(
+                &format!("count_distinct_groups i64 {num_groups} groups"),
+                |b| {
+                    b.iter(|| {
+                        let mut accumulators: Vec<_> = (0..num_groups)
+                            .map(|_| prepare_accumulator(DataType::Int64))
+                            .collect();
+
+                        let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
+                        for (idx, group_idx) in group_indices.iter().enumerate() {
+                            if let Some(val) = arr.value(idx).into() {
+                                let single_val =
+                                    Arc::new(Int64Array::from(vec![Some(val)]))
+                                        as ArrayRef;
+                                accumulators[*group_idx]
+                                    .update_batch(std::slice::from_ref(&single_val))
+                                    .unwrap();
+                            }
+                        }
+
+                        let _results: Vec<_> = accumulators
+                            .iter_mut()
+                            .map(|acc| acc.evaluate().unwrap())
+                            .collect();
+                    })
+                },
+            );
+        }
+    }
+}
+
+criterion_group!(
+    benches,
+    count_distinct_benchmark,
+    count_distinct_groups_benchmark
+);
 criterion_main!(benches);

datafusion/functions-aggregate/src/count.rs

@@ -41,6 +41,7 @@ use datafusion_expr::{
     function::{AccumulatorArgs, StateFieldsArgs},
     utils::format_state_name,
 };
+use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator;
 use datafusion_functions_aggregate_common::aggregate::{
     count_distinct::BytesDistinctCountAccumulator,
     count_distinct::BytesViewDistinctCountAccumulator,
@@ -336,18 +337,70 @@ impl AggregateUDFImpl for Count {
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
-        // groups accumulator only supports `COUNT(c1)`, not
-        // `COUNT(c1, c2)`, etc
-        if args.is_distinct {
+        if args.exprs.len() != 1 {
             return false;
         }
-        args.exprs.len() == 1
+        if args.is_distinct {
+            // Only support primitive integer types for now
+            matches!(
+                args.expr_fields[0].data_type(),
+                DataType::Int8
+                    | DataType::Int16
+                    | DataType::Int32
+                    | DataType::Int64
+                    | DataType::UInt8
+                    | DataType::UInt16
+                    | DataType::UInt32
+                    | DataType::UInt64
+            )
+        } else {
+            true
+        }
     }
 
     fn create_groups_accumulator(
         &self,
-        _args: AccumulatorArgs,
+        args: AccumulatorArgs,
     ) -> Result<Box<dyn GroupsAccumulator>> {
+        if args.is_distinct {
+            let data_type = args.expr_fields[0].data_type();
+            return match data_type {
+                DataType::Int8 => Ok(Box::new(
+                    PrimitiveDistinctCountGroupsAccumulator::<Int8Type>::new(),
+                )),
+                DataType::Int16 => Ok(Box::new(
+                    PrimitiveDistinctCountGroupsAccumulator::<Int16Type>::new(),
+                )),
+                DataType::Int32 => Ok(Box::new(
+                    PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new(),
+                )),
+                DataType::Int64 => Ok(Box::new(
+                    PrimitiveDistinctCountGroupsAccumulator::<Int64Type>::new(),
+                )),
+                DataType::UInt8 => Ok(Box::new(
+                    PrimitiveDistinctCountGroupsAccumulator::<UInt8Type>::new(),
+                )),
+                DataType::UInt16 => {
+                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+                        UInt16Type,
+                    >::new()))
+                }
+                DataType::UInt32 => {
+                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+                        UInt32Type,
+                    >::new()))
+                }
+                DataType::UInt64 => {
+                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+                        UInt64Type,
+                    >::new()))
+                }
+                _ => not_impl_err!(
+                    "GroupsAccumulator not supported for COUNT(DISTINCT) with {}",
+                    data_type
+                ),
+            };
+        }
         // instantiate specialized accumulator
         Ok(Box::new(CountGroupsAccumulator::new()))
     }