expl_tuple_vec_on_demand_sort

Author: B Vadlamani

datafusion/functions-aggregate/src/count.rs

@@ -377,32 +377,30 @@ impl AggregateUDFImpl for Count {
         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
+        if !args.is_distinct {
+            return true;
         }
+        matches!(
+            args.expr_fields[0].data_type(),
+            DataType::Int8
+                | DataType::Int16
+                | DataType::Int32
+                | DataType::Int64
+                | DataType::UInt8
+                | DataType::UInt16
+                | DataType::UInt32
+                | DataType::UInt64
+        )
     }
 
     fn create_groups_accumulator(
         &self,
         args: AccumulatorArgs,
     ) -> Result<Box<dyn GroupsAccumulator>> {
-        if args.is_distinct {
-            return create_distinct_groups_accumulator(args.expr_fields[0].data_type());
+        if !args.is_distinct {
+            return Ok(Box::new(CountGroupsAccumulator::new()));
         }
-        Ok(Box::new(CountGroupsAccumulator::new()))
+        create_distinct_groups_accumulator(args.expr_fields[0].data_type())
     }
 
     fn reverse_expr(&self) -> ReversedUDAF {