Add StateFieldsArgs::return_field (#16112)

Author: alamb

datafusion-examples/examples/advanced_udaf.rs

@@ -94,7 +94,7 @@ impl AggregateUDFImpl for GeoMeanUdaf {
     /// This is the description of the state. accumulator's state() must match the types here.
     fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
         Ok(vec![
-            Field::new("prod", args.return_field.data_type().clone(), true),
+            Field::new("prod", args.return_type().clone(), true),
             Field::new("n", DataType::UInt32, true),
         ])
     }

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

@@ -100,3 +100,10 @@ pub struct StateFieldsArgs<'a> {
     /// Whether the aggregate function is distinct.
     pub is_distinct: bool,
 }
+
+impl StateFieldsArgs<'_> {
+    /// The return type of the aggregate function.
+    pub fn return_type(&self) -> &DataType {
+        self.return_field.data_type()
+    }
+}

datafusion/functions-aggregate/src/bit_and_or_xor.rs

@@ -271,7 +271,7 @@ impl AggregateUDFImpl for BitwiseOperation {
                     format!("{} distinct", self.name()).as_str(),
                 ),
                 // See COMMENTS.md to understand why nullable is set to true
-                Field::new_list_field(args.return_field.data_type().clone(), true),
+                Field::new_list_field(args.return_type().clone(), true),
                 false,
             )])
         } else {

datafusion/functions-aggregate/src/first_last.rs

@@ -172,7 +172,7 @@ impl AggregateUDFImpl for FirstValue {
     fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
         let mut fields = vec![Field::new(
             format_state_name(args.name, "first_value"),
-            args.return_field.data_type().clone(),
+            args.return_type().clone(),
             true,
         )];
         fields.extend(args.ordering_fields.to_vec());

datafusion/functions-aggregate/src/sum.rs

@@ -206,13 +206,13 @@ impl AggregateUDFImpl for Sum {
             Ok(vec![Field::new_list(
                 format_state_name(args.name, "sum distinct"),
                 // See COMMENTS.md to understand why nullable is set to true
-                Field::new_list_field(args.return_field.data_type().clone(), true),
+                Field::new_list_field(args.return_type().clone(), true),
                 false,
             )])
         } else {
             Ok(vec![Field::new(
                 format_state_name(args.name, "sum"),
-                args.return_field.data_type().clone(),
+                args.return_type().clone(),
                 true,
             )])
         }