Move window functions over to use Field instead of DataType

Author: timsaucer

datafusion/core/tests/user_defined/user_defined_window_functions.rs

@@ -18,11 +18,15 @@
 //! This module contains end to end tests of creating
 //! user defined window functions
 
-use arrow::array::{ArrayRef, AsArray, Int64Array, RecordBatch, StringArray};
+use arrow::array::{
+    record_batch, Array, ArrayRef, AsArray, Int64Array, RecordBatch, StringArray,
+    UInt64Array,
+};
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion::common::test_util::batches_to_string;
 use datafusion::common::{Result, ScalarValue};
 use datafusion::prelude::SessionContext;
+use datafusion_common::exec_datafusion_err;
 use datafusion_expr::{
     PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF, WindowUDFImpl,
 };
@@ -34,6 +38,7 @@ use datafusion_physical_expr::{
     expressions::{col, lit},
     PhysicalExpr,
 };
+use std::collections::HashMap;
 use std::{
     any::Any,
     ops::Range,
@@ -723,11 +728,11 @@ fn test_default_expressions() -> Result<()> {
     ];
 
     for input_exprs in &test_cases {
-        let input_types = input_exprs
+        let input_fields = input_exprs
             .iter()
-            .map(|expr: &Arc<dyn PhysicalExpr>| expr.data_type(&schema).unwrap())
+            .map(|expr: &Arc<dyn PhysicalExpr>| expr.return_field(&schema).unwrap())
             .collect::<Vec<_>>();
-        let expr_args = ExpressionArgs::new(input_exprs, &input_types);
+        let expr_args = ExpressionArgs::new(input_exprs, &input_fields);
 
         let ret_exprs = udwf.expressions(expr_args);
 
@@ -753,3 +758,148 @@ fn test_default_expressions() -> Result<()> {
     }
     Ok(())
 }
+
+#[derive(Debug)]
+struct MetadataBasedWindowUdf {
+    name: String,
+    signature: Signature,
+    metadata: HashMap<String, String>,
+}
+
+impl MetadataBasedWindowUdf {
+    fn new(metadata: HashMap<String, String>) -> Self {
+        // The name we return must be unique. Otherwise we will not call distinct
+        // instances of this UDF. This is a small hack for the unit tests to get unique
+        // names, but you could do something more elegant with the metadata.
+        let name = format!("metadata_based_udf_{}", metadata.len());
+        Self {
+            name,
+            signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+            metadata,
+        }
+    }
+}
+
+impl WindowUDFImpl for MetadataBasedWindowUdf {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn partition_evaluator(
+        &self,
+        partition_evaluator_args: PartitionEvaluatorArgs,
+    ) -> Result<Box<dyn PartitionEvaluator>> {
+        let input_field = partition_evaluator_args
+            .input_fields()
+            .first()
+            .ok_or(exec_datafusion_err!("Expected one argument"))?;
+
+        let double_output = input_field
+            .metadata()
+            .get("modify_values")
+            .map(|v| v == "double_output")
+            .unwrap_or(false);
+
+        Ok(Box::new(MetadataBasedPartitionEvaluator { double_output }))
+    }
+
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
+        Ok(Field::new(field_args.name(), DataType::UInt64, true)
+            .with_metadata(self.metadata.clone()))
+    }
+}
+
+#[derive(Debug)]
+struct MetadataBasedPartitionEvaluator {
+    double_output: bool,
+}
+
+impl PartitionEvaluator for MetadataBasedPartitionEvaluator {
+    fn evaluate_all(&mut self, values: &[ArrayRef], num_rows: usize) -> Result<ArrayRef> {
+        let values = values[0].as_any().downcast_ref::<UInt64Array>().unwrap();
+        let sum = values.iter().fold(0_u64, |acc, v| acc + v.unwrap_or(0));
+
+        let result = if self.double_output { sum * 2 } else { sum };
+
+        Ok(Arc::new(UInt64Array::from_value(result, num_rows)))
+    }
+}
+
+#[tokio::test]
+async fn test_metadata_based_window_fn() -> Result<()> {
+    let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("no_metadata", DataType::UInt64, true),
+        Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+            [("modify_values".to_string(), "double_output".to_string())]
+                .into_iter()
+                .collect(),
+        ),
+    ]));
+
+    let batch = RecordBatch::try_new(
+        schema,
+        vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+    )?;
+
+    let ctx = SessionContext::new();
+    ctx.register_batch("t", batch)?;
+    let df = ctx.table("t").await?;
+
+    let no_output_meta_udf = WindowUDF::from(MetadataBasedWindowUdf::new(HashMap::new()));
+    let with_output_meta_udf = WindowUDF::from(MetadataBasedWindowUdf::new(
+        [("output_metatype".to_string(), "custom_value".to_string())]
+            .into_iter()
+            .collect(),
+    ));
+
+    let df = df.select(vec![
+        no_output_meta_udf
+            .call(vec![datafusion_expr::col("no_metadata")])
+            .alias("meta_no_in_no_out"),
+        no_output_meta_udf
+            .call(vec![datafusion_expr::col("with_metadata")])
+            .alias("meta_with_in_no_out"),
+        with_output_meta_udf
+            .call(vec![datafusion_expr::col("no_metadata")])
+            .alias("meta_no_in_with_out"),
+        with_output_meta_udf
+            .call(vec![datafusion_expr::col("with_metadata")])
+            .alias("meta_with_in_with_out"),
+    ])?;
+
+    let actual = df.collect().await?;
+
+    // To test for output metadata handling, we set the expected values on the result
+    // To test for input metadata handling, we check the numbers returned
+    let mut output_meta = HashMap::new();
+    let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+    let expected_schema = Schema::new(vec![
+        Field::new("meta_no_in_no_out", DataType::UInt64, true),
+        Field::new("meta_with_in_no_out", DataType::UInt64, true),
+        Field::new("meta_no_in_with_out", DataType::UInt64, true)
+            .with_metadata(output_meta.clone()),
+        Field::new("meta_with_in_with_out", DataType::UInt64, true)
+            .with_metadata(output_meta.clone()),
+    ]);
+
+    let expected = record_batch!(
+        ("meta_no_in_no_out", UInt64, [50, 50, 50, 50, 50]),
+        ("meta_with_in_no_out", UInt64, [100, 100, 100, 100, 100]),
+        ("meta_no_in_with_out", UInt64, [50, 50, 50, 50, 50]),
+        ("meta_with_in_with_out", UInt64, [100, 100, 100, 100, 100])
+    )?
+    .with_schema(Arc::new(expected_schema))?;
+
+    assert_eq!(expected, actual[0]);
+
+    Ok(())
+}

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, FieldRef};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::cse::{HashNode, NormalizeEq, Normalizeable};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion,
@@ -838,19 +838,19 @@ pub enum WindowFunctionDefinition {
 
 impl WindowFunctionDefinition {
     /// Returns the datatype of the window function
-    pub fn return_type(
+    pub fn return_field(
         &self,
-        input_expr_types: &[DataType],
+        input_expr_fields: &[Field],
         _input_expr_nullable: &[bool],
         display_name: &str,
-    ) -> Result<DataType> {
+    ) -> Result<Field> {
         match self {
             WindowFunctionDefinition::AggregateUDF(fun) => {
-                fun.return_type(input_expr_types)
+                fun.return_field(input_expr_fields)
+            }
+            WindowFunctionDefinition::WindowUDF(fun) => {
+                fun.field(WindowUDFFieldArgs::new(input_expr_fields, display_name))
             }
-            WindowFunctionDefinition::WindowUDF(fun) => fun
-                .field(WindowUDFFieldArgs::new(input_expr_types, display_name))
-                .map(|field| field.data_type().clone()),
         }
     }
 

datafusion/expr/src/expr_schema.rs

@@ -647,8 +647,8 @@ impl Expr {
                     .map(|f| f.data_type())
                     .cloned()
                     .collect::<Vec<_>>();
-                let new_types =
-                    data_types_with_window_udf(&data_types, udwf).map_err(|err| {
+                let new_fields =
+                    data_types_with_window_udf(&fields, udwf).map_err(|err| {
                         plan_datafusion_err!(
                             "{} {}",
                             match err {
@@ -663,7 +663,7 @@ impl Expr {
                         )
                     })?;
                 let (_, function_name) = self.qualified_name();
-                let field_args = WindowUDFFieldArgs::new(&new_types, &function_name);
+                let field_args = WindowUDFFieldArgs::new(&new_fields, &function_name);
 
                 udwf.field(field_args)
                     .map(|field| (field.data_type().clone(), field.is_nullable()))

datafusion/expr/src/type_coercion/functions.rs

@@ -130,13 +130,13 @@ pub fn data_types_with_aggregate_udf(
 /// For more details on coercion in general, please see the
 /// [`type_coercion`](crate::type_coercion) module.
 pub fn data_types_with_window_udf(
-    current_types: &[DataType],
+    current_fields: &[Field],
     func: &WindowUDF,
-) -> Result<Vec<DataType>> {
+) -> Result<Vec<Field>> {
     let signature = func.signature();
     let type_signature = &signature.type_signature;
 
-    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
+    if current_fields.is_empty() && type_signature != &TypeSignature::UserDefined {
         if type_signature.supports_zero_argument() {
             return Ok(vec![]);
         } else if type_signature.used_to_support_zero_arguments() {
@@ -147,16 +147,28 @@ pub fn data_types_with_window_udf(
         }
     }
 
+    let current_types = current_fields
+        .iter()
+        .map(|f| f.data_type())
+        .cloned()
+        .collect::<Vec<_>>();
     let valid_types =
-        get_valid_types_with_window_udf(type_signature, current_types, func)?;
+        get_valid_types_with_window_udf(type_signature, &current_types, func)?;
     if valid_types
         .iter()
-        .any(|data_type| data_type == current_types)
+        .any(|data_type| data_type == &current_types)
     {
-        return Ok(current_types.to_vec());
+        return Ok(current_fields.to_vec());
     }
 
-    try_coerce_types(func.name(), valid_types, current_types, type_signature)
+    let updated_types =
+        try_coerce_types(func.name(), valid_types, &current_types, type_signature)?;
+
+    Ok(current_fields
+        .iter()
+        .zip(updated_types)
+        .map(|(current_field, new_type)| current_field.clone().with_data_type(new_type))
+        .collect())
 }
 
 /// Performs type coercion for function arguments.

datafusion/functions-window-common/src/expr.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion_common::arrow::datatypes::DataType;
+use datafusion_common::arrow::datatypes::Field;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use std::sync::Arc;
 
@@ -27,7 +27,7 @@ pub struct ExpressionArgs<'a> {
     input_exprs: &'a [Arc<dyn PhysicalExpr>],
     /// The corresponding data types of expressions passed as arguments
     /// to the user-defined window function.
-    input_types: &'a [DataType],
+    input_fields: &'a [Field],
 }
 
 impl<'a> ExpressionArgs<'a> {
@@ -42,11 +42,11 @@ impl<'a> ExpressionArgs<'a> {
     ///
     pub fn new(
         input_exprs: &'a [Arc<dyn PhysicalExpr>],
-        input_types: &'a [DataType],
+        input_fields: &'a [Field],
     ) -> Self {
         Self {
             input_exprs,
-            input_types,
+            input_fields,
         }
     }
 
@@ -56,9 +56,9 @@ impl<'a> ExpressionArgs<'a> {
         self.input_exprs
     }
 
-    /// Returns the [`DataType`]s corresponding to the input expressions
+    /// Returns the [`Field`]s corresponding to the input expressions
     /// to the user-defined window function.
-    pub fn input_types(&self) -> &'a [DataType] {
-        self.input_types
+    pub fn input_fields(&self) -> &'a [Field] {
+        self.input_fields
     }
 }

datafusion/functions-window-common/src/field.rs

@@ -15,14 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion_common::arrow::datatypes::DataType;
+use datafusion_common::arrow::datatypes::Field;
 
 /// Metadata for defining the result field from evaluating a
 /// user-defined window function.
 pub struct WindowUDFFieldArgs<'a> {
     /// The data types corresponding to the arguments to the
     /// user-defined window function.
-    input_types: &'a [DataType],
+    input_fields: &'a [Field],
     /// The display name of the user-defined window function.
     display_name: &'a str,
 }
@@ -32,22 +32,22 @@ impl<'a> WindowUDFFieldArgs<'a> {
     ///
     /// # Arguments
     ///
-    /// * `input_types` - The data types corresponding to the
+    /// * `input_fields` - The fields corresponding to the
     ///   arguments to the user-defined window function.
     /// * `function_name` - The qualified schema name of the
     ///   user-defined window function expression.
     ///
-    pub fn new(input_types: &'a [DataType], display_name: &'a str) -> Self {
+    pub fn new(input_fields: &'a [Field], display_name: &'a str) -> Self {
         WindowUDFFieldArgs {
-            input_types,
+            input_fields,
             display_name,
         }
     }
 
-    /// Returns the data type of input expressions passed as arguments
+    /// Returns the field of input expressions passed as arguments
     /// to the user-defined window function.
-    pub fn input_types(&self) -> &[DataType] {
-        self.input_types
+    pub fn input_fields(&self) -> &[Field] {
+        self.input_fields
     }
 
     /// Returns the name for the field of the final result of evaluating
@@ -56,9 +56,9 @@ impl<'a> WindowUDFFieldArgs<'a> {
         self.display_name
     }
 
-    /// Returns `Some(DataType)` of input expression at index, otherwise
+    /// Returns `Some(Field)` of input expression at index, otherwise
     /// returns `None` if the index is out of bounds.
-    pub fn get_input_type(&self, index: usize) -> Option<DataType> {
-        self.input_types.get(index).cloned()
+    pub fn get_input_field(&self, index: usize) -> Option<Field> {
+        self.input_fields.get(index).cloned()
     }
 }