implement_group_accumulators_count_distinct

Author: coderfender

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

@@ -17,12 +17,12 @@
 
 mod bytes;
 mod dict;
-mod native;
 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;
-pub use groups::PrimitiveDistinctCountGroupsAccumulator;

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

@@ -1,51 +1,121 @@
-use std::hash::Hash;
-use arrow::array::{ArrayRef, BooleanArray, Int64Array};
-use arrow::datatypes::{ArrowPrimitiveType, DataType};
+use arrow::array::{
+    Array, ArrayRef, AsArray, BooleanArray, Int64Array, ListBuilder, PrimitiveBuilder,
+};
+use arrow::datatypes::ArrowPrimitiveType;
+use datafusion_common::HashSet;
+use datafusion_common::hash_utils::RandomState;
 use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
+use std::hash::Hash;
+use std::sync::Arc;
 
 pub struct PrimitiveDistinctCountGroupsAccumulator<T: ArrowPrimitiveType>
-    where
+where
     T::Native: Eq + Hash,
-    {
+{
     /// Count distinct per group.
-    values: Vec<T>, 
-        data_type : DataType
+    values: Vec<HashSet<T::Native, RandomState>>,
+}
+
+impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
+where
+    T::Native: Eq + Hash,
+{
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
 impl<T: ArrowPrimitiveType> PrimitiveDistinctCountGroupsAccumulator<T>
 where
     T::Native: Eq + Hash,
 {
-    pub fn new(data_type: DataType) -> Self {
-        Self {
-            values: Vec::new(),
-            data_type,
-        }
+    pub fn new() -> Self {
+        Self { values: Vec::new() }
     }
 }
 
 impl<T: ArrowPrimitiveType + Send + std::fmt::Debug> GroupsAccumulator
-for PrimitiveDistinctCountGroupsAccumulator<T>
+    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<()> {
-        todo!()
+{
+    fn update_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion_common::Result<()> {
+        self.values.resize_with(total_num_groups, HashSet::default);
+        debug_assert_eq!(values.len(), 1, "multiple arguments are not supported");
+
+        let arr = values[0].as_primitive::<T>();
+
+        for (idx, group_idx) in group_indices.iter().enumerate() {
+            if let Some(filter) = opt_filter
+                && !filter.value(idx)
+            {
+                continue;
+            }
+            if arr.is_valid(idx) {
+                let value = arr.value(idx);
+                self.values[*group_idx].insert(value);
+            }
+        }
+
+        Ok(())
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
-        todo!()
+        let counts: Vec<i64> = emit_to
+            .take_needed(&mut self.values)
+            .iter()
+            .map(|groups| groups.len() as i64)
+            .collect();
+
+        Ok(Arc::new(Int64Array::from(counts)))
     }
 
     fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
-        todo!()
+        let hash_sets = emit_to.take_needed(&mut self.values);
+        let mut builder = ListBuilder::new(PrimitiveBuilder::<T>::new());
+
+        for set in hash_sets {
+            for value in set {
+                builder.values().append_value(value);
+            }
+            builder.append(true);
+        }
+        Ok(vec![Arc::new(builder.finish())])
     }
 
-    fn merge_batch(&mut self, values: &[ArrayRef], group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> datafusion_common::Result<()> {
-        todo!()
+    fn merge_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        _opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion_common::Result<()> {
+        self.values.resize_with(total_num_groups, HashSet::default);
+        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_set = inner.as_primitive::<T>();
+            for i in 0..inner.len() {
+                self.values[*group_idx].insert(inner_set.value(i));
+            }
+        }
+        Ok(())
     }
 
     fn size(&self) -> usize {
-        todo!()
+        size_of::<Self>()
+            + self.values.capacity() * size_of::<HashSet<T::Native, RandomState>>()
+            + self
+                .values
+                .iter()
+                .map(|s| s.capacity() * size_of::<T::Native>())
+                .sum::<usize>()
     }
 }

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,
@@ -59,7 +60,6 @@ use std::{
     ops::BitAnd,
     sync::Arc,
 };
-use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator;
 
 make_udaf_expr_and_func!(
     Count,
@@ -344,19 +344,38 @@ impl AggregateUDFImpl for Count {
         &self,
         args: AccumulatorArgs,
     ) -> Result<Box<dyn GroupsAccumulator>> {
-        if(args.is_distinct){
+        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(data_type.clone()))),
-                DataType::Int16 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<Int16Type>::new(data_type.clone()))),
-                DataType::Int32 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new(data_type.clone()))),
-                DataType::Int64 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<Int64Type>::new(data_type.clone()))),
-                DataType::UInt8 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<UInt8Type>::new(data_type.clone()))),
-                DataType::UInt16 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<UInt16Type>::new(data_type.clone()))),
-                DataType::UInt32 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<UInt32Type>::new(data_type.clone()))),
-                DataType::UInt64 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<UInt64Type>::new(data_type.clone()))),
-                _ => not_impl_err!("GroupsAccumulator not supported for COUNT(DISTINCT) with {}", 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()))