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,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. 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`, whose per-group
+ /// state can be orders of magnitude larger.
+ ///
+ /// An implementation whose answer is decided by the argument types and
+ /// `DISTINCT` alone should override this and have
+ /// [`Self::groups_accumulator_supported`] call it, so the two cannot
+ /// disagree. One whose answer needs more than that should leave the `None`
+ /// default.
+ ///
+ /// `None` means unanswered, not "no": a caller must not read it as either
+ /// `Some(true)` or `Some(false)`, since both readings are wrong for some
+ /// implementation that returns it.
+ 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 +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,
+ ) -> 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(