Introduce expr_fields to AccumulatorArgs to hold input argument fields (#18100)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #16997
- Part of #11725
- Supersedes #17085

## Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->

When reviewing #17085 I was very confused by the fix suggested, and
tried to understand why `AccumulatorArgs` didn't have easy access to
`Field`s of its input expressions, as compared to scalar/window
functions which do. Introducing this new field should make it easier for
users to grab datatype, metadata, nullability of their input expressions
for aggregate functions.

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

Add a slice of `FieldRef` to `AccumulatorArgs` so users don't need to
compute the input expression fields themselves via using schema. This
addresses #16997 as it was confusing to have only the schema available
as there are valid (?) cases where the schema is empty (such as literal
only input).

This fix differs from #17085 in that it doesn't special case for when
there is literal only input; it leaves the physical `schema` provided to
`AccumulatorArgs` untouched but provides a more ergonomic (and less
confusing) API for users to retrieve `Field`s of their input arguments.

- I'm still not sure if the schema being empty for literal only inputs
is correct or not, so this might be considered a side step. If we could
remove `schema` entirely from `AccumulatorArgs` maybe we wouldn't need
to worry about this, but see my comment for why that wasn't done in this
PR

## Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

Existing unit tests.

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

Yes, new field to `AccumulatorArgs` which is publicly exposed (with all
it's fields).

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: Jefffrey

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -954,13 +954,7 @@ impl AggregateUDFImpl for MetadataBasedAggregateUdf {
     }
 
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
-        let input_expr = acc_args
-            .exprs
-            .first()
-            .ok_or(exec_datafusion_err!("Expected one argument"))?;
-        let input_field = input_expr.return_field(acc_args.schema)?;
-
-        let double_output = input_field
+        let double_output = acc_args.expr_fields[0]
             .metadata()
             .get("modify_values")
             .map(|v| v == "double_output")

datafusion/ffi/src/udaf/accumulator_args.rs

@@ -97,6 +97,7 @@ impl TryFrom<AccumulatorArgs<'_>> for FFI_AccumulatorArgs {
 pub struct ForeignAccumulatorArgs {
     pub return_field: FieldRef,
     pub schema: Schema,
+    pub expr_fields: Vec<FieldRef>,
     pub ignore_nulls: bool,
     pub order_bys: Vec<PhysicalSortExpr>,
     pub is_reversed: bool,
@@ -132,9 +133,15 @@ impl TryFrom<FFI_AccumulatorArgs> for ForeignAccumulatorArgs {
 
         let exprs = parse_physical_exprs(&proto_def.expr, &task_ctx, &schema, &codex)?;
 
+        let expr_fields = exprs
+            .iter()
+            .map(|e| e.return_field(&schema))
+            .collect::<Result<Vec<_>, _>>()?;
+
         Ok(Self {
             return_field,
             schema,
+            expr_fields,
             ignore_nulls: proto_def.ignore_nulls,
             order_bys,
             is_reversed: value.is_reversed,
@@ -150,6 +157,7 @@ impl<'a> From<&'a ForeignAccumulatorArgs> for AccumulatorArgs<'a> {
         Self {
             return_field: Arc::clone(&value.return_field),
             schema: &value.schema,
+            expr_fields: &value.expr_fields,
             ignore_nulls: value.ignore_nulls,
             order_bys: &value.order_bys,
             is_reversed: value.is_reversed,
@@ -175,6 +183,7 @@ mod tests {
         let orig_args = AccumulatorArgs {
             return_field: Field::new("f", DataType::Float64, true).into(),
             schema: &schema,
+            expr_fields: &[Field::new("a", DataType::Int32, true).into()],
             ignore_nulls: false,
             order_bys: &[PhysicalSortExpr::new_default(col("a", &schema)?)],
             is_reversed: false,

datafusion/ffi/src/udaf/mod.rs

@@ -705,6 +705,7 @@ mod tests {
         let acc_args = AccumulatorArgs {
             return_field: Field::new("f", DataType::Float64, true).into(),
             schema: &schema,
+            expr_fields: &[Field::new("a", DataType::Float64, true).into()],
             ignore_nulls: true,
             order_bys: &[PhysicalSortExpr::new_default(col("a", &schema)?)],
             is_reversed: false,
@@ -782,6 +783,7 @@ mod tests {
         let acc_args = AccumulatorArgs {
             return_field: Field::new("f", DataType::Float64, true).into(),
             schema: &schema,
+            expr_fields: &[Field::new("a", DataType::Float64, true).into()],
             ignore_nulls: true,
             order_bys: &[PhysicalSortExpr::new_default(col("a", &schema)?)],
             is_reversed: false,

datafusion/functions-aggregate-common/src/accumulator.rs

@@ -30,7 +30,8 @@ pub struct AccumulatorArgs<'a> {
     /// The return field of the aggregate function.
     pub return_field: FieldRef,
 
-    /// The schema of the input arguments
+    /// Input schema to the aggregate function. If you need to check data type, nullability
+    /// or metadata of input arguments then you should use `expr_fields` below instead.
     pub schema: &'a Schema,
 
     /// Whether to ignore nulls.
@@ -67,6 +68,9 @@ pub struct AccumulatorArgs<'a> {
 
     /// The physical expression of arguments the aggregate function takes.
     pub exprs: &'a [Arc<dyn PhysicalExpr>],
+
+    /// Fields corresponding to each expr (same order & length).
+    pub expr_fields: &'a [FieldRef],
 }
 
 impl AccumulatorArgs<'_> {

datafusion/functions-aggregate/benches/count.rs

@@ -33,15 +33,17 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
 fn prepare_group_accumulator() -> Box<dyn GroupsAccumulator> {
     let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Int32, true)]));
+    let expr = col("f", &schema).unwrap();
     let accumulator_args = AccumulatorArgs {
         return_field: Field::new("f", DataType::Int64, true).into(),
         schema: &schema,
+        expr_fields: &[expr.return_field(&schema).unwrap()],
         ignore_nulls: false,
         order_bys: &[],
         is_reversed: false,
         name: "COUNT(f)",
         is_distinct: false,
-        exprs: &[col("f", &schema).unwrap()],
+        exprs: &[expr],
     };
     let count_fn = Count::new();
 
@@ -56,15 +58,17 @@ fn prepare_accumulator() -> Box<dyn Accumulator> {
         DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
         true,
     )]));
+    let expr = col("f", &schema).unwrap();
     let accumulator_args = AccumulatorArgs {
         return_field: Arc::new(Field::new_list_field(DataType::Int64, true)),
         schema: &schema,
+        expr_fields: &[expr.return_field(&schema).unwrap()],
         ignore_nulls: false,
         order_bys: &[],
         is_reversed: false,
         name: "COUNT(f)",
         is_distinct: true,
-        exprs: &[col("f", &schema).unwrap()],
+        exprs: &[expr],
     };
     let count_fn = Count::new();
 

datafusion/functions-aggregate/benches/min_max_bytes.rs

@@ -44,6 +44,7 @@ fn create_max_bytes_accumulator() -> Box<dyn GroupsAccumulator> {
     max.create_groups_accumulator(AccumulatorArgs {
         return_field: Arc::new(Field::new("value", DataType::Utf8, true)),
         schema: &input_schema,
+        expr_fields: &[Field::new("value", DataType::Utf8, true).into()],
         ignore_nulls: true,
         order_bys: &[],
         is_reversed: false,

datafusion/functions-aggregate/benches/sum.rs

@@ -31,8 +31,9 @@ fn prepare_accumulator(data_type: &DataType) -> Box<dyn GroupsAccumulator> {
     let field = Field::new("f", data_type.clone(), true).into();
     let schema = Arc::new(Schema::new(vec![Arc::clone(&field)]));
     let accumulator_args = AccumulatorArgs {
-        return_field: field,
+        return_field: Arc::clone(&field),
         schema: &schema,
+        expr_fields: &[field],
         ignore_nulls: false,
         order_bys: &[],
         is_reversed: false,

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -361,7 +361,7 @@ impl AggregateUDFImpl for ApproxDistinct {
     }
 
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
-        let data_type = acc_args.exprs[0].data_type(acc_args.schema)?;
+        let data_type = acc_args.expr_fields[0].data_type();
 
         let accumulator: Box<dyn Accumulator> = match data_type {
             // TODO u8, i8, u16, i16 shall really be done using bitmap, not HLL

datafusion/functions-aggregate/src/approx_median.rs

@@ -134,7 +134,7 @@ impl AggregateUDFImpl for ApproxMedian {
 
         Ok(Box::new(ApproxPercentileAccumulator::new(
             0.5_f64,
-            acc_args.exprs[0].data_type(acc_args.schema)?,
+            acc_args.expr_fields[0].data_type().clone(),
         )))
     }
 

datafusion/functions-aggregate/src/approx_percentile_cont.rs

@@ -187,9 +187,9 @@ impl ApproxPercentileCont {
             None
         };
 
-        let data_type = args.exprs[0].data_type(args.schema)?;
+        let data_type = args.expr_fields[0].data_type();
         let accumulator: ApproxPercentileAccumulator = match data_type {
-            t @ (DataType::UInt8
+            DataType::UInt8
             | DataType::UInt16
             | DataType::UInt32
             | DataType::UInt64
@@ -198,12 +198,11 @@ impl ApproxPercentileCont {
             | DataType::Int32
             | DataType::Int64
             | DataType::Float32
-            | DataType::Float64) => {
+            | DataType::Float64 => {
                 if let Some(max_size) = tdigest_max_size {
-                    ApproxPercentileAccumulator::new_with_max_size(percentile, t, max_size)
-                }else{
-                    ApproxPercentileAccumulator::new(percentile, t)
-
+                    ApproxPercentileAccumulator::new_with_max_size(percentile, data_type.clone(), max_size)
+                } else {
+                    ApproxPercentileAccumulator::new(percentile, data_type.clone())
                 }
             }
             other => {