Shift from Field to FieldRef for all user defined functions (#16122)

* Initial commit of moving from Field to FieldRef for user defined functions

* Documentation updates

* Add include in example

* More documentation updates

* Updates required after rebasing on main

* Change from Field to FieldRef in upgrade documentation

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: timsaucer

datafusion-examples/examples/advanced_udaf.rs

@@ -25,6 +25,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{ArrowNativeTypeOp, ArrowPrimitiveType, Float64Type, UInt32Type};
 use arrow::record_batch::RecordBatch;
+use arrow_schema::FieldRef;
 use datafusion::common::{cast::as_float64_array, ScalarValue};
 use datafusion::error::Result;
 use datafusion::logical_expr::{
@@ -92,10 +93,10 @@ 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>> {
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new("prod", args.return_type().clone(), true),
-            Field::new("n", DataType::UInt32, true),
+            Field::new("prod", args.return_type().clone(), true).into(),
+            Field::new("n", DataType::UInt32, true).into(),
         ])
     }
 
@@ -401,7 +402,7 @@ impl AggregateUDFImpl for SimplifiedGeoMeanUdaf {
         unimplemented!("should not be invoked")
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         unimplemented!("should not be invoked")
     }
 

datafusion-examples/examples/advanced_udwf.rs

@@ -23,6 +23,7 @@ use arrow::{
     array::{ArrayRef, AsArray, Float64Array},
     datatypes::Float64Type,
 };
+use arrow_schema::FieldRef;
 use datafusion::common::ScalarValue;
 use datafusion::error::Result;
 use datafusion::functions_aggregate::average::avg_udaf;
@@ -87,8 +88,8 @@ impl WindowUDFImpl for SmoothItUdf {
         Ok(Box::new(MyPartitionEvaluator::new()))
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-        Ok(Field::new(field_args.name(), DataType::Float64, true))
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+        Ok(Field::new(field_args.name(), DataType::Float64, true).into())
     }
 }
 
@@ -205,8 +206,8 @@ impl WindowUDFImpl for SimplifySmoothItUdf {
         Some(Box::new(simplify))
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-        Ok(Field::new(field_args.name(), DataType::Float64, true))
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+        Ok(Field::new(field_args.name(), DataType::Float64, true).into())
     }
 }
 

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -128,7 +128,7 @@ fn test_update_matching_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            Field::new("f", DataType::Int32, true),
+            Field::new("f", DataType::Int32, true).into(),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 2))),
@@ -193,7 +193,7 @@ fn test_update_matching_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            Field::new("f", DataType::Int32, true),
+            Field::new("f", DataType::Int32, true).into(),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 3))),
@@ -261,7 +261,7 @@ fn test_update_projected_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            Field::new("f", DataType::Int32, true),
+            Field::new("f", DataType::Int32, true).into(),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 2))),
@@ -326,7 +326,7 @@ fn test_update_projected_exprs() -> Result<()> {
                     Arc::new(Column::new("b_new", 1)),
                 )),
             ],
-            Field::new("f", DataType::Int32, true),
+            Field::new("f", DataType::Int32, true).into(),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d_new", 3))),

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -32,6 +32,7 @@ use arrow::array::{
     StringArray, StructArray, UInt64Array,
 };
 use arrow::datatypes::{Fields, Schema};
+use arrow_schema::FieldRef;
 use datafusion::common::test_util::batches_to_string;
 use datafusion::dataframe::DataFrame;
 use datafusion::datasource::MemTable;
@@ -572,7 +573,7 @@ impl TimeSum {
         // Returns the same type as its input
         let return_type = timestamp_type.clone();
 
-        let state_fields = vec![Field::new("sum", timestamp_type, true)];
+        let state_fields = vec![Field::new("sum", timestamp_type, true).into()];
 
         let volatility = Volatility::Immutable;
 
@@ -672,7 +673,7 @@ impl FirstSelector {
         let state_fields = state_type
             .into_iter()
             .enumerate()
-            .map(|(i, t)| Field::new(format!("{i}"), t, true))
+            .map(|(i, t)| Field::new(format!("{i}"), t, true).into())
             .collect::<Vec<_>>();
 
         // Possible input signatures
@@ -932,9 +933,10 @@ impl AggregateUDFImpl for MetadataBasedAggregateUdf {
         unimplemented!("this should never be called since return_field is implemented");
     }
 
-    fn return_field(&self, _arg_fields: &[Field]) -> Result<Field> {
+    fn return_field(&self, _arg_fields: &[FieldRef]) -> Result<FieldRef> {
         Ok(Field::new(self.name(), DataType::UInt64, true)
-            .with_metadata(self.metadata.clone()))
+            .with_metadata(self.metadata.clone())
+            .into())
     }
 
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -28,7 +28,7 @@ use arrow::array::{
 use arrow::compute::kernels::numeric::add;
 use arrow::datatypes::{DataType, Field, Schema};
 use arrow_schema::extension::{Bool8, CanonicalExtensionType, ExtensionType};
-use arrow_schema::ArrowError;
+use arrow_schema::{ArrowError, FieldRef};
 use datafusion::common::test_util::batches_to_string;
 use datafusion::execution::context::{FunctionFactory, RegisterFunction, SessionState};
 use datafusion::prelude::*;
@@ -814,7 +814,7 @@ impl ScalarUDFImpl for TakeUDF {
     ///
     /// 1. If the third argument is '0', return the type of the first argument
     /// 2. If the third argument is '1', return the type of the second argument
-    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Field> {
+    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
         if args.arg_fields.len() != 3 {
             return plan_err!("Expected 3 arguments, got {}.", args.arg_fields.len());
         }
@@ -845,7 +845,8 @@ impl ScalarUDFImpl for TakeUDF {
             self.name(),
             args.arg_fields[take_idx].data_type().to_owned(),
             true,
-        ))
+        )
+        .into())
     }
 
     // The actual implementation
@@ -1412,9 +1413,10 @@ impl ScalarUDFImpl for MetadataBasedUdf {
         );
     }
 
-    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
         Ok(Field::new(self.name(), DataType::UInt64, true)
-            .with_metadata(self.metadata.clone()))
+            .with_metadata(self.metadata.clone())
+            .into())
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -1562,14 +1564,15 @@ impl ScalarUDFImpl for ExtensionBasedUdf {
         Ok(DataType::Utf8)
     }
 
-    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
         Ok(Field::new("canonical_extension_udf", DataType::Utf8, true)
-            .with_extension_type(MyUserExtentionType {}))
+            .with_extension_type(MyUserExtentionType {})
+            .into())
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
         assert_eq!(args.arg_fields.len(), 1);
-        let input_field = args.arg_fields[0];
+        let input_field = args.arg_fields[0].as_ref();
 
         let output_as_bool = matches!(
             CanonicalExtensionType::try_from(input_field),

datafusion/core/tests/user_defined/user_defined_window_functions.rs

@@ -23,6 +23,7 @@ use arrow::array::{
     UInt64Array,
 };
 use arrow::datatypes::{DataType, Field, Schema};
+use arrow_schema::FieldRef;
 use datafusion::common::test_util::batches_to_string;
 use datafusion::common::{Result, ScalarValue};
 use datafusion::prelude::SessionContext;
@@ -564,8 +565,8 @@ impl OddCounter {
                 &self.aliases
             }
 
-            fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-                Ok(Field::new(field_args.name(), DataType::Int64, true))
+            fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+                Ok(Field::new(field_args.name(), DataType::Int64, true).into())
             }
         }
 
@@ -683,7 +684,7 @@ impl WindowUDFImpl for VariadicWindowUDF {
         unimplemented!("unnecessary for testing");
     }
 
-    fn field(&self, _: WindowUDFFieldArgs) -> Result<Field> {
+    fn field(&self, _: WindowUDFFieldArgs) -> Result<FieldRef> {
         unimplemented!("unnecessary for testing");
     }
 }
@@ -809,9 +810,10 @@ impl WindowUDFImpl for MetadataBasedWindowUdf {
         Ok(Box::new(MetadataBasedPartitionEvaluator { double_output }))
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
         Ok(Field::new(field_args.name(), DataType::UInt64, true)
-            .with_metadata(self.metadata.clone()))
+            .with_metadata(self.metadata.clone())
+            .into())
     }
 }
 

datafusion/expr-common/src/type_coercion/aggregates.rs

@@ -17,7 +17,7 @@
 
 use crate::signature::TypeSignature;
 use arrow::datatypes::{
-    DataType, Field, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
+    DataType, FieldRef, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
     DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE,
 };
 
@@ -89,7 +89,7 @@ pub static TIMES: &[DataType] = &[
 /// number of input types.
 pub fn check_arg_count(
     func_name: &str,
-    input_fields: &[Field],
+    input_fields: &[FieldRef],
     signature: &TypeSignature,
 ) -> Result<()> {
     match signature {

datafusion/expr/src/expr.rs

@@ -28,7 +28,7 @@ use crate::logical_plan::Subquery;
 use crate::Volatility;
 use crate::{udaf, ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
 
-use arrow::datatypes::{DataType, Field, FieldRef};
+use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::cse::{HashNode, NormalizeEq, Normalizeable};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion,
@@ -846,10 +846,10 @@ impl WindowFunctionDefinition {
     /// Returns the datatype of the window function
     pub fn return_field(
         &self,
-        input_expr_fields: &[Field],
+        input_expr_fields: &[FieldRef],
         _input_expr_nullable: &[bool],
         display_name: &str,
-    ) -> Result<Field> {
+    ) -> Result<FieldRef> {
         match self {
             WindowFunctionDefinition::AggregateUDF(fun) => {
                 fun.return_field(input_expr_fields)

datafusion/expr/src/expr_fn.rs

@@ -37,7 +37,7 @@ use crate::{
 use arrow::compute::kernels::cast_utils::{
     parse_interval_day_time, parse_interval_month_day_nano, parse_interval_year_month,
 };
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::{plan_err, Column, Result, ScalarValue, Spans, TableReference};
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
@@ -492,6 +492,7 @@ pub fn create_udaf(
         .into_iter()
         .enumerate()
         .map(|(i, t)| Field::new(format!("{i}"), t, true))
+        .map(Arc::new)
         .collect::<Vec<_>>();
     AggregateUDF::from(SimpleAggregateUDF::new(
         name,
@@ -510,7 +511,7 @@ pub struct SimpleAggregateUDF {
     signature: Signature,
     return_type: DataType,
     accumulator: AccumulatorFactoryFunction,
-    state_fields: Vec<Field>,
+    state_fields: Vec<FieldRef>,
 }
 
 impl Debug for SimpleAggregateUDF {
@@ -533,7 +534,7 @@ impl SimpleAggregateUDF {
         return_type: DataType,
         volatility: Volatility,
         accumulator: AccumulatorFactoryFunction,
-        state_fields: Vec<Field>,
+        state_fields: Vec<FieldRef>,
     ) -> Self {
         let name = name.into();
         let signature = Signature::exact(input_type, volatility);
@@ -553,7 +554,7 @@ impl SimpleAggregateUDF {
         signature: Signature,
         return_type: DataType,
         accumulator: AccumulatorFactoryFunction,
-        state_fields: Vec<Field>,
+        state_fields: Vec<FieldRef>,
     ) -> Self {
         let name = name.into();
         Self {
@@ -590,7 +591,7 @@ impl AggregateUDFImpl for SimpleAggregateUDF {
         (self.accumulator)(acc_args)
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(self.state_fields.clone())
     }
 }
@@ -678,12 +679,12 @@ impl WindowUDFImpl for SimpleWindowUDF {
         (self.partition_evaluator_factory)()
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
-        Ok(Field::new(
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+        Ok(Arc::new(Field::new(
             field_args.name(),
             self.return_type.clone(),
             true,
-        ))
+        )))
     }
 }