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,
+    ) -> Option<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,38 @@ 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 `None`, which means the implementation does not answer
+    /// this question. An implementation that overrides
+    /// [`Self::groups_accumulator_supported`], 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 `None`.
+    ///
+    /// `None` is not a third answer to the question. A caller must not read it
+    /// as either `Some(true)` or `Some(false)`, because both readings are
+    /// wrong for some implementation that returns it. A caller that has to act
+    /// on an unanswered question must take the action that is safe when either
+    /// answer turns out to be the true one.
+    fn groups_accumulator_supported_for_types(
+        &self,
+        _arg_types: &[DataType],
+        _is_distinct: bool,
+    ) -> Option<bool> {
+        None
+    }
+
     /// Return a specialized [`GroupsAccumulator`] that manages state
     /// for all groups.
     ///
@@ -1552,6 +1594,15 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl {
         self.inner.groups_accumulator_supported(args)
     }
 
+    fn groups_accumulator_supported_for_types(
+        &self,
+        arg_types: &[DataType],
+        is_distinct: bool,
+    ) -> Option<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,31 @@ impl AggregateUDFImpl for Count {
     }
 
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
-        if args.exprs.len() != 1 {
-            return false;
+        // 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)
+            .unwrap_or(false)
+    }
+
+    fn groups_accumulator_supported_for_types(
+        &self,
+        arg_types: &[DataType],
+        is_distinct: bool,
+    ) -> Option<bool> {
+        if arg_types.len() != 1 {
+            return Some(false);
         }
-        if !args.is_distinct {
-            return true;
+        if !is_distinct {
+            return Some(true);
         }
-        matches!(
-            args.expr_fields[0].data_type(),
+        // Keep in step with `create_distinct_count_groups_accumulator`.
+        Some(matches!(
+            arg_types[0],
             DataType::Int8
                 | DataType::Int16
                 | DataType::Int32
@@ -362,7 +379,7 @@ impl AggregateUDFImpl for Count {
                 | DataType::UInt16
                 | DataType::UInt32
                 | DataType::UInt64
-        )
+        ))
     }
 
     fn create_groups_accumulator(