Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy

datafusion/expr/src/udaf.rs

@@ -253,6 +253,16 @@ impl AggregateUDF {
         self.inner.groups_accumulator_supported(args)
     }
 
+    /// See [`AggregateUDFImpl::groups_accumulator_supported_for_types`] for more details.
+    pub fn groups_accumulator_supported_for_types(
+        &self,
+        arg_types: &[DataType],
+        is_distinct: bool,
+    ) -> bool {
+        self.inner
+            .groups_accumulator_supported_for_types(arg_types, is_distinct)
+    }
+
     /// See [`AggregateUDFImpl::create_groups_accumulator`] for more details.
     pub fn create_groups_accumulator(
         &self,
@@ -617,6 +627,32 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         false
     }
 
+    /// The same question as [`Self::groups_accumulator_supported`], asked with
+    /// only the information a logical plan carries.
+    ///
+    /// Physical planning has an [`AccumulatorArgs`] to ask with; an optimizer
+    /// rule does not, and cannot fabricate one, so this is how a logical
+    /// caller learns whether a call would get a specialized
+    /// [`GroupsAccumulator`] or fall back to one boxed [`Accumulator`] per
+    /// group in `GroupsAccumulatorAdapter`. That distinction is worth a rule
+    /// changing its mind over: the adapter's per-group state can be orders of
+    /// magnitude larger.
+    ///
+    /// The default is `false`, matching the default of
+    /// [`Self::groups_accumulator_supported`]. An implementation that
+    /// overrides that one and whose answer is decided by the argument types
+    /// and `DISTINCT` alone should override this one too, and have the
+    /// physical method call it so the two cannot disagree. An implementation
+    /// whose answer needs more than the argument types should leave this at
+    /// `false`, which claims nothing.
+    fn groups_accumulator_supported_for_types(
+        &self,
+        _arg_types: &[DataType],
+        _is_distinct: bool,
+    ) -> bool {
+        false
+    }
+
     /// Return a specialized [`GroupsAccumulator`] that manages state
     /// for all groups.
     ///
@@ -1552,6 +1588,15 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.inner.groups_accumulator_supported(args)
     }
 
+    fn groups_accumulator_supported_for_types(
+        &self,
+        arg_types: &[DataType],
+        is_distinct: bool,
+    ) -> bool {
+        self.inner
+            .groups_accumulator_supported_for_types(arg_types, is_distinct)
+    }
+
     fn create_groups_accumulator(
         &self,
         args: AccumulatorArgs,

datafusion/functions-aggregate/src/count.rs

@@ -346,14 +346,30 @@ impl AggregateUDFImpl for Count {
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
-        if args.exprs.len() != 1 {
+        // The answer depends on nothing but the argument types and `DISTINCT`,
+        // so defer to the logical form and keep one list of supported types.
+        let arg_types = args
+            .expr_fields
+            .iter()
+            .map(|field| field.data_type().clone())
+            .collect::<Vec<_>>();
+        self.groups_accumulator_supported_for_types(&arg_types, args.is_distinct)
+    }
+
+    fn groups_accumulator_supported_for_types(
+        &self,
+        arg_types: &[DataType],
+        is_distinct: bool,
+    ) -> bool {
+        if arg_types.len() != 1 {
             return false;
         }
-        if !args.is_distinct {
+        if !is_distinct {
             return true;
         }
+        // Keep in step with `create_distinct_count_groups_accumulator`.
         matches!(
-            args.expr_fields[0].data_type(),
+            arg_types[0],
             DataType::Int8
                 | DataType::Int16
                 | DataType::Int32