Add field trait method to WindowUDFImpl, remove return_type/nullable

Cargo.toml

@@ -31,6 +31,7 @@ members = [
     "datafusion/functions-aggregate-common",
     "datafusion/functions-nested",
     "datafusion/functions-window",
+    "datafusion/functions-window-common",
     "datafusion/optimizer",
     "datafusion/physical-expr",
     "datafusion/physical-expr-common",
@@ -103,6 +104,7 @@ datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", vers
 datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "41.0.0" }
 datafusion-functions-nested = { path = "datafusion/functions-nested", version = "41.0.0" }
 datafusion-functions-window = { path = "datafusion/functions-window", version = "41.0.0" }
+datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "41.0.0" }
 datafusion-optimizer = { path = "datafusion/optimizer", version = "41.0.0", default-features = false }
 datafusion-physical-expr = { path = "datafusion/physical-expr", version = "41.0.0", default-features = false }
 datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "41.0.0", default-features = false }

datafusion-examples/examples/advanced_udwf.rs

@@ -22,9 +22,11 @@ use arrow::{
     array::{ArrayRef, AsArray, Float64Array},
     datatypes::Float64Type,
 };
+use arrow_schema::Field;
 use datafusion::error::Result;
 use datafusion::prelude::*;
 use datafusion_common::ScalarValue;
+use datafusion_expr::function::WindowUDFFieldArgs;
 use datafusion_expr::{
     PartitionEvaluator, Signature, WindowFrame, WindowUDF, WindowUDFImpl,
 };
@@ -80,6 +82,10 @@ impl WindowUDFImpl for SmoothItUdf {
     fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         Ok(Box::new(MyPartitionEvaluator::new()))
     }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        Ok(Field::new(field_args.name(), DataType::Float64, true))
+    }
 }
 
 /// This implements the lowest level evaluation for a window function

datafusion-examples/examples/simplify_udwf_expression.rs

@@ -17,12 +17,12 @@
 
 use std::any::Any;
 
-use arrow_schema::DataType;
+use arrow_schema::{DataType, Field};
 
 use datafusion::execution::context::SessionContext;
 use datafusion::functions_aggregate::average::avg_udaf;
 use datafusion::{error::Result, execution::options::CsvReadOptions};
-use datafusion_expr::function::WindowFunctionSimplification;
+use datafusion_expr::function::{WindowFunctionSimplification, WindowUDFFieldArgs};
 use datafusion_expr::{
     expr::WindowFunction, simplify::SimplifyInfo, Expr, PartitionEvaluator, Signature,
     Volatility, WindowUDF, WindowUDFImpl,
@@ -84,6 +84,10 @@ 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))
+    }
 }
 
 // create local execution context with `cars.csv` registered as a table named `cars`

datafusion/core/Cargo.toml

@@ -145,6 +145,7 @@ bigdecimal = { workspace = true }
 criterion = { version = "0.5", features = ["async_tokio"] }
 csv = "1.1.6"
 ctor = { workspace = true }
+datafusion-functions-window-common = { workspace = true }
 doc-comment = { workspace = true }
 env_logger = { workspace = true }
 half = { workspace = true, default-features = true }

datafusion/core/tests/user_defined/user_defined_window_functions.rs

@@ -29,12 +29,13 @@ use std::{
 
 use arrow::array::AsArray;
 use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
-use arrow_schema::DataType;
+use arrow_schema::{DataType, Field};
 use datafusion::{assert_batches_eq, prelude::SessionContext};
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::{
     PartitionEvaluator, Signature, Volatility, WindowUDF, WindowUDFImpl,
 };
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 
 /// A query with a window function evaluated over the entire partition
 const UNBOUNDED_WINDOW_QUERY: &str = "SELECT x, y, val, \
@@ -565,6 +566,14 @@ impl OddCounter {
             fn aliases(&self) -> &[String] {
                 &self.aliases
             }
+
+            fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+                Ok(Field::new(
+                    field_args.name(),
+                    self.return_type.clone(),
+                    true,
+                ))
+            }
         }
 
         ctx.register_udwf(WindowUDF::from(SimpleWindowUDF::new(test_state)))

datafusion/expr/Cargo.toml

@@ -46,6 +46,7 @@ chrono = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-expr-common = { workspace = true }
 datafusion-functions-aggregate-common = { workspace = true }
+datafusion-functions-window-common = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 paste = "^1.0"
 serde_json = { workspace = true }

datafusion/expr/src/expr.rs

@@ -714,7 +714,12 @@ impl WindowFunctionDefinition {
             WindowFunctionDefinition::AggregateUDF(fun) => {
                 fun.return_type(input_expr_types)
             }
-            WindowFunctionDefinition::WindowUDF(fun) => fun.return_type(input_expr_types),
+            WindowFunctionDefinition::WindowUDF(_) => {
+                // To get the return data type of the result from
+                // evaluating the user-defined window function instead
+                // use the `WindowUDF::field` trait method.
+                unreachable!()
+            }
         }
     }
 

datafusion/expr/src/expr_fn.rs

@@ -38,6 +38,7 @@ use arrow::compute::kernels::cast_utils::{
 };
 use arrow::datatypes::{DataType, Field};
 use datafusion_common::{plan_err, Column, Result, ScalarValue, TableReference};
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use sqlparser::ast::NullTreatment;
 use std::any::Any;
 use std::fmt::Debug;
@@ -665,6 +666,14 @@ impl WindowUDFImpl for SimpleWindowUDF {
     fn partition_evaluator(&self) -> Result<Box<dyn crate::PartitionEvaluator>> {
         (self.partition_evaluator_factory)()
     }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        Ok(Field::new(
+            field_args.name(),
+            self.return_type.clone(),
+            true,
+        ))
+    }
 }
 
 pub fn interval_year_month_lit(value: &str) -> Expr {

datafusion/expr/src/expr_schema.rs

@@ -31,6 +31,7 @@ use datafusion_common::{
     not_impl_err, plan_datafusion_err, plan_err, Column, ExprSchema, Result,
     TableReference,
 };
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -204,7 +205,12 @@ impl ExprSchemable for Expr {
                                     )
                                 )
                             })?;
-                        Ok(fun.return_type(&new_types, &nullability)?)
+                        let function_name = self.schema_name().to_string();
+                        let field_args =
+                            WindowUDFFieldArgs::new(&new_types, &function_name);
+
+                        udwf.field(field_args)
+                            .map(|field| field.data_type().clone())
                     }
                     _ => fun.return_type(&data_types, &nullability),
                 }
@@ -340,7 +346,7 @@ impl ExprSchemable for Expr {
             Expr::AggregateFunction(AggregateFunction { func, .. }) => {
                 Ok(func.is_nullable())
             }
-            Expr::WindowFunction(WindowFunction { fun, .. }) => match fun {
+            Expr::WindowFunction(WindowFunction { fun, args, .. }) => match fun {
                 WindowFunctionDefinition::BuiltInWindowFunction(func) => {
                     if func.name() == "RANK"
                         || func.name() == "NTILE"
@@ -352,7 +358,18 @@ impl ExprSchemable for Expr {
                     }
                 }
                 WindowFunctionDefinition::AggregateUDF(func) => Ok(func.is_nullable()),
-                WindowFunctionDefinition::WindowUDF(udwf) => Ok(udwf.nullable()),
+                WindowFunctionDefinition::WindowUDF(udwf) => {
+                    let data_types = args
+                        .iter()
+                        .map(|e| e.get_type(input_schema))
+                        .collect::<Result<Vec<_>>>()?;
+                    let input_types = data_types_with_window_udf(&data_types, udwf)?;
+                    let function_name = self.schema_name().to_string();
+                    let field_args =
+                        WindowUDFFieldArgs::new(&input_types, &function_name);
+
+                    udwf.field(field_args).map(|field| field.is_nullable())
+                }
             },
             Expr::ScalarVariable(_, _)
             | Expr::TryCast { .. }

datafusion/expr/src/function.rs

@@ -27,6 +27,8 @@ pub use datafusion_functions_aggregate_common::accumulator::{
     AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs,
 };
 
+pub use datafusion_functions_window_common::field::WindowUDFFieldArgs;
+
 #[derive(Debug, Clone, Copy)]
 pub enum Hint {
     /// Indicates the argument needs to be padded if it is scalar

datafusion/expr/src/udwf.rs

@@ -25,9 +25,10 @@ use std::{
     sync::Arc,
 };
 
-use arrow::datatypes::DataType;
+use arrow::datatypes::{DataType, Field};
 
 use datafusion_common::{not_impl_err, Result};
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 
 use crate::expr::WindowFunction;
 use crate::{
@@ -162,6 +163,7 @@ impl WindowUDF {
     /// Return the type of the function given its input types
     ///
     /// See [`WindowUDFImpl::return_type`] for more details.
+    #[allow(deprecated)]
     pub fn return_type(&self, args: &[DataType]) -> Result<DataType> {
         self.inner.return_type(args)
     }
@@ -178,11 +180,11 @@ impl WindowUDF {
         self.inner.partition_evaluator()
     }
 
-    /// Returns if column values are nullable for this window function.
+    /// Returns the field of the final result of evaluating this window function.
     ///
-    /// See [`WindowUDFImpl::nullable`] for more details.
-    pub fn nullable(&self) -> bool {
-        self.inner.nullable()
+    /// See [`WindowUDFImpl::field`] for more details.
+    pub fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        self.inner.field(field_args)
     }
 
     /// Returns custom result ordering introduced by this window function
@@ -221,10 +223,11 @@ where
 /// # Basic Example
 /// ```
 /// # use std::any::Any;
-/// # use arrow::datatypes::DataType;
+/// # use arrow::datatypes::{DataType, Field};
 /// # use datafusion_common::{DataFusionError, plan_err, Result};
 /// # use datafusion_expr::{col, Signature, Volatility, PartitionEvaluator, WindowFrame, ExprFunctionExt};
 /// # use datafusion_expr::{WindowUDFImpl, WindowUDF};
+/// use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 /// #[derive(Debug, Clone)]
 /// struct SmoothIt {
 ///   signature: Signature
@@ -251,6 +254,7 @@ where
 ///    }
 ///    // The actual implementation would add one to the argument
 ///    fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { unimplemented!() }
+///    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> { unimplemented!() }
 /// }
 ///
 /// // Create a new WindowUDF from the implementation
@@ -281,6 +285,10 @@ pub trait WindowUDFImpl: Debug + Send + Sync {
 
     /// What [`DataType`] will be returned by this function, given the types of
     /// the arguments
+    #[deprecated(
+        since = "41.0.0",
+        note = "Use `field` instead to define the final result of evaluating this user-defined window function."
+    )]
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;
 
     /// Invoke the function, returning the [`PartitionEvaluator`] instance
@@ -344,14 +352,8 @@ pub trait WindowUDFImpl: Debug + Send + Sync {
         hasher.finish()
     }
 
-    /// Allows customizing nullable of column for this window UDF.
-    ///
-    /// By default, the final result of evaluating the window UDF is
-    /// allowed to have null values. But if that is not the case then
-    /// it can be customized in the window UDF implementation.
-    fn nullable(&self) -> bool {
-        true
-    }
+    /// The [`Field`] of the final result of evaluating this window function.
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field>;
 
     /// Allows the window UDF to define a custom result ordering.
     ///
@@ -419,6 +421,7 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
         self.inner.signature()
     }
 
+    #[allow(deprecated)]
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         self.inner.return_type(arg_types)
     }
@@ -450,8 +453,8 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
         hasher.finish()
     }
 
-    fn nullable(&self) -> bool {
-        self.inner.nullable()
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        self.inner.field(field_args)
     }
 
     fn sort_options(&self) -> Option<SortOptions> {
@@ -510,4 +513,14 @@ impl WindowUDFImpl for WindowUDFLegacyWrapper {
     fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         (self.partition_evaluator_factory)()
     }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        let return_type = (self.return_type)(field_args.input_types())?;
+
+        Ok(Field::new(
+            field_args.name(),
+            return_type.as_ref().clone(),
+            true,
+        ))
+    }
 }