Use return_field_from_args in information schema and date_trunc (#20079)

## Which issue does this PR close?

- Closes #19870.

## Rationale for this change

Some UDFs/UDAFs implement `return_field_from_args` / `return_field`
instead of `return_type`. The information schema was calling
`return_type` directly, which fails for those functions. The default
implementation of `return_field_from_args` already delegates to
`return_type`, so switching to the newer API works for all functions.

## What changes are included in this PR?

- **`information_schema.rs`**: `get_udf_args_and_return_types` now calls
`return_field_from_args` instead of `return_type`;
`get_udaf_args_and_return_types` now calls `return_field` instead of
`return_type`. Removed stale comments referencing the old API.
- **`date_trunc.rs`**: `return_type` now returns `internal_err`, and
`return_field_from_args` is self-contained (no longer delegates to
`return_type`), following the same pattern as other UDFs like
`named_struct` and `map_from_arrays` (ref: #19275).

## Are these changes tested?

Covered by existing information_schema sqllogictests and
`datafusion-functions` unit tests.

## Are there any user-facing changes?

No.

Author: AndreaBozzo

datafusion/catalog/src/information_schema.rs

@@ -24,7 +24,7 @@ use crate::{CatalogProviderList, SchemaProvider, TableProvider};
 use arrow::array::builder::{BooleanBuilder, UInt8Builder};
 use arrow::{
     array::{StringBuilder, UInt64Builder},
-    datatypes::{DataType, Field, Schema, SchemaRef},
+    datatypes::{DataType, Field, FieldRef, Schema, SchemaRef},
     record_batch::RecordBatch,
 };
 use async_trait::async_trait;
@@ -34,7 +34,10 @@ use datafusion_common::error::Result;
 use datafusion_common::types::NativeType;
 use datafusion_execution::TaskContext;
 use datafusion_execution::runtime_env::RuntimeEnv;
-use datafusion_expr::{AggregateUDF, ScalarUDF, Signature, TypeSignature, WindowUDF};
+use datafusion_expr::function::WindowUDFFieldArgs;
+use datafusion_expr::{
+    AggregateUDF, ReturnFieldArgs, ScalarUDF, Signature, TypeSignature, WindowUDF,
+};
 use datafusion_expr::{TableType, Volatility};
 use datafusion_physical_plan::SendableRecordBatchStream;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
@@ -421,10 +424,24 @@ fn get_udf_args_and_return_types(
         Ok(arg_types
             .into_iter()
             .map(|arg_types| {
-                // only handle the function which implemented [`ScalarUDFImpl::return_type`] method
+                let arg_fields: Vec<FieldRef> = arg_types
+                    .iter()
+                    .enumerate()
+                    .map(|(i, t)| {
+                        Arc::new(Field::new(format!("arg_{i}"), t.clone(), true))
+                    })
+                    .collect();
+                let scalar_arguments = vec![None; arg_fields.len()];
                 let return_type = udf
-                    .return_type(&arg_types)
-                    .map(|t| remove_native_type_prefix(&NativeType::from(t)))
+                    .return_field_from_args(ReturnFieldArgs {
+                        arg_fields: &arg_fields,
+                        scalar_arguments: &scalar_arguments,
+                    })
+                    .map(|f| {
+                        remove_native_type_prefix(&NativeType::from(
+                            f.data_type().clone(),
+                        ))
+                    })
                     .ok();
                 let arg_types = arg_types
                     .into_iter()
@@ -447,11 +464,21 @@ fn get_udaf_args_and_return_types(
         Ok(arg_types
             .into_iter()
             .map(|arg_types| {
-                // only handle the function which implemented [`ScalarUDFImpl::return_type`] method
+                let arg_fields: Vec<FieldRef> = arg_types
+                    .iter()
+                    .enumerate()
+                    .map(|(i, t)| {
+                        Arc::new(Field::new(format!("arg_{i}"), t.clone(), true))
+                    })
+                    .collect();
                 let return_type = udaf
-                    .return_type(&arg_types)
-                    .ok()
-                    .map(|t| remove_native_type_prefix(&NativeType::from(t)));
+                    .return_field(&arg_fields)
+                    .map(|f| {
+                        remove_native_type_prefix(&NativeType::from(
+                            f.data_type().clone(),
+                        ))
+                    })
+                    .ok();
                 let arg_types = arg_types
                     .into_iter()
                     .map(|t| remove_native_type_prefix(&NativeType::from(t)))
@@ -473,12 +500,26 @@ fn get_udwf_args_and_return_types(
         Ok(arg_types
             .into_iter()
             .map(|arg_types| {
-                // only handle the function which implemented [`ScalarUDFImpl::return_type`] method
+                let arg_fields: Vec<FieldRef> = arg_types
+                    .iter()
+                    .enumerate()
+                    .map(|(i, t)| {
+                        Arc::new(Field::new(format!("arg_{i}"), t.clone(), true))
+                    })
+                    .collect();
+                let return_type = udwf
+                    .field(WindowUDFFieldArgs::new(&arg_fields, udwf.name()))
+                    .map(|f| {
+                        remove_native_type_prefix(&NativeType::from(
+                            f.data_type().clone(),
+                        ))
+                    })
+                    .ok();
                 let arg_types = arg_types
                     .into_iter()
                     .map(|t| remove_native_type_prefix(&NativeType::from(t)))
                     .collect::<Vec<_>>();
-                (arg_types, None)
+                (arg_types, return_type)
             })
             .collect::<BTreeSet<_>>())
     }

datafusion/functions/benches/date_trunc.rs

@@ -25,7 +25,7 @@ use arrow::datatypes::Field;
 use criterion::{Criterion, criterion_group, criterion_main};
 use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
-use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
+use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs};
 use datafusion_functions::datetime::date_trunc;
 use rand::Rng;
 use rand::rngs::ThreadRng;
@@ -57,10 +57,13 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
             .collect::<Vec<_>>();
 
-        let return_type = udf
-            .return_type(&args.iter().map(|arg| arg.data_type()).collect::<Vec<_>>())
+        let scalar_arguments = vec![None; arg_fields.len()];
+        let return_field = udf
+            .return_field_from_args(ReturnFieldArgs {
+                arg_fields: &arg_fields,
+                scalar_arguments: &scalar_arguments,
+            })
             .unwrap();
-        let return_field = Arc::new(Field::new("f", return_type, true));
         let config_options = Arc::new(ConfigOptions::default());
 
         b.iter(|| {

datafusion/functions/src/datetime/date_trunc.rs

@@ -38,7 +38,7 @@ use arrow::datatypes::{Field, FieldRef};
 use datafusion_common::cast::as_primitive_array;
 use datafusion_common::types::{NativeType, logical_date, logical_string};
 use datafusion_common::{
-    DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err,
+    DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, internal_err,
 };
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
@@ -223,27 +223,21 @@ impl ScalarUDFImpl for DateTruncFunc {
         &self.signature
     }
 
-    // keep return_type implementation for information schema generation
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        if arg_types[1].is_null() {
-            Ok(Timestamp(Nanosecond, None))
-        } else {
-            Ok(arg_types[1].clone())
-        }
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        internal_err!("return_field_from_args should be called instead")
     }
 
     fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
-        let data_types = args
-            .arg_fields
-            .iter()
-            .map(|f| f.data_type())
-            .cloned()
-            .collect::<Vec<_>>();
-        let return_type = self.return_type(&data_types)?;
+        let field = &args.arg_fields[1];
+        let return_type = if field.data_type().is_null() {
+            Timestamp(Nanosecond, None)
+        } else {
+            field.data_type().clone()
+        };
         Ok(Arc::new(Field::new(
             self.name(),
             return_type,
-            args.arg_fields[1].is_nullable(),
+            field.is_nullable(),
         )))
     }