implement_group_accumulators_count_distinct

Author: coderfender

datafusion/functions-aggregate/src/count.rs

@@ -337,7 +337,25 @@ impl AggregateUDFImpl for Count {
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
-        args.exprs.len() == 1
+        if args.exprs.len() != 1 {
+            return false;
+        }
+        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(