init

Author: coderfender

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

@@ -18,9 +18,11 @@
 mod bytes;
 mod dict;
 mod native;
+mod groups;
 
 pub use bytes::BytesDistinctCountAccumulator;
 pub use bytes::BytesViewDistinctCountAccumulator;
 pub use dict::DictionaryCountAccumulator;
 pub use native::FloatDistinctCountAccumulator;
 pub use native::PrimitiveDistinctCountAccumulator;
+pub use groups::PrimitiveDistinctCountGroupsAccumulator;

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

@@ -0,0 +1,51 @@
+use std::hash::Hash;
+use arrow::array::{ArrayRef, BooleanArray, Int64Array};
+use arrow::datatypes::{ArrowPrimitiveType, DataType};
+use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
+
+pub struct PrimitiveDistinctCountGroupsAccumulator<T: ArrowPrimitiveType>
+    where
+    T::Native: Eq + Hash,
+    {
+    /// Count distinct per group.
+    values: Vec<T>, 
+        data_type : DataType
+}
+
+impl<T: ArrowPrimitiveType> PrimitiveDistinctCountGroupsAccumulator<T>
+where
+    T::Native: Eq + Hash,
+{
+    pub fn new(data_type: DataType) -> Self {
+        Self {
+            values: Vec::new(),
+            data_type,
+        }
+    }
+}
+
+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<()> {
+        todo!()
+    }
+
+    fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
+        todo!()
+    }
+
+    fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
+        todo!()
+    }
+
+    fn merge_batch(&mut self, values: &[ArrayRef], group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> datafusion_common::Result<()> {
+        todo!()
+    }
+
+    fn size(&self) -> usize {
+        todo!()
+    }
+}

datafusion/functions-aggregate/src/count.rs

@@ -59,6 +59,7 @@ use std::{
     ops::BitAnd,
     sync::Arc,
 };
+use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator;
 
 make_udaf_expr_and_func!(
     Count,
@@ -336,18 +337,27 @@ 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 {
-            return false;
-        }
         args.exprs.len() == 1
     }
 
     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(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),
+            }
+        }
         // instantiate specialized accumulator
         Ok(Box::new(CountGroupsAccumulator::new()))
     }