signature and return types

Author: Jiayu Liu

datafusion/src/physical_plan/aggregates.rs

@@ -182,7 +182,7 @@ static TIMESTAMPS: &[DataType] = &[
 ];
 
 /// the signatures supported by the function `fun`.
-fn signature(fun: &AggregateFunction) -> Signature {
+pub fn signature(fun: &AggregateFunction) -> Signature {
     // note: the physical expression must accept the type returned by this function or the execution panics.
     match fun {
         AggregateFunction::Count => Signature::Any(1),

datafusion/src/physical_plan/windows.rs

@@ -20,11 +20,10 @@
 //!
 //! see also https://www.postgresql.org/docs/current/functions-window.html
 
-use super::expressions::{avg_return_type, sum_return_type};
 use super::{functions::Signature, type_coercion::data_types};
 use crate::error::{DataFusionError, Result};
-use crate::physical_plan::aggregates::AggregateFunction;
-use arrow::datatypes::{DataType, TimeUnit};
+use crate::physical_plan::{aggregates, aggregates::AggregateFunction};
+use arrow::datatypes::DataType;
 use std::{fmt, str::FromStr};
 
 /// WindowFunction
@@ -90,59 +89,32 @@ pub fn return_type(fun: &WindowFunction, arg_types: &[DataType]) -> Result<DataT
     data_types(arg_types, &signature(fun))?;
 
     match fun {
-        WindowFunction::AggregateFunction(fun) => match fun {
-            AggregateFunction::Count => Ok(DataType::UInt64),
-            AggregateFunction::Max | AggregateFunction::Min => Ok(arg_types[0].clone()),
-            AggregateFunction::Sum => sum_return_type(&arg_types[0]),
-            AggregateFunction::Avg => avg_return_type(&arg_types[0]),
-        },
-        WindowFunction::BuiltInWindowFunction(_) => Ok(arg_types[0].clone()),
+        WindowFunction::AggregateFunction(fun) => aggregates::return_type(fun, arg_types),
+        WindowFunction::BuiltInWindowFunction(fun) => Ok(match fun {
+            BuiltInWindowFunction::RowNumber
+            | BuiltInWindowFunction::Rank
+            | BuiltInWindowFunction::DenseRank => DataType::UInt64,
+            BuiltInWindowFunction::Lag
+            | BuiltInWindowFunction::Lead
+            | BuiltInWindowFunction::FirstValue
+            | BuiltInWindowFunction::LastValue => arg_types[0].clone(),
+        }),
     }
 }
 
-static NUMERICS: &[DataType] = &[
-    DataType::Int8,
-    DataType::Int16,
-    DataType::Int32,
-    DataType::Int64,
-    DataType::UInt8,
-    DataType::UInt16,
-    DataType::UInt32,
-    DataType::UInt64,
-    DataType::Float32,
-    DataType::Float64,
-];
-
-static STRINGS: &[DataType] = &[DataType::Utf8, DataType::LargeUtf8];
-
-static TIMESTAMPS: &[DataType] = &[
-    DataType::Timestamp(TimeUnit::Second, None),
-    DataType::Timestamp(TimeUnit::Millisecond, None),
-    DataType::Timestamp(TimeUnit::Microsecond, None),
-    DataType::Timestamp(TimeUnit::Nanosecond, None),
-];
-
 /// the signatures supported by the function `fun`.
 fn signature(fun: &WindowFunction) -> Signature {
     // note: the physical expression must accept the type returned by this function or the execution panics.
     match fun {
-        WindowFunction::AggregateFunction(fun) => match fun {
-            AggregateFunction::Count => Signature::Any(1),
-            AggregateFunction::Min | AggregateFunction::Max => {
-                let valid = STRINGS
-                    .iter()
-                    .chain(NUMERICS.iter())
-                    .chain(TIMESTAMPS.iter())
-                    .cloned()
-                    .collect::<Vec<_>>();
-                Signature::Uniform(1, valid)
-            }
-            AggregateFunction::Avg | AggregateFunction::Sum => {
-                Signature::Uniform(1, NUMERICS.to_vec())
-            }
+        WindowFunction::AggregateFunction(fun) => aggregates::signature(fun),
+        WindowFunction::BuiltInWindowFunction(fun) => match fun {
+            BuiltInWindowFunction::RowNumber
+            | BuiltInWindowFunction::Rank
+            | BuiltInWindowFunction::DenseRank => Signature::Any(0),
+            BuiltInWindowFunction::Lag
+            | BuiltInWindowFunction::Lead
+            | BuiltInWindowFunction::FirstValue
+            | BuiltInWindowFunction::LastValue => Signature::Any(1),
         },
-        WindowFunction::BuiltInWindowFunction(_) => {
-            Signature::Uniform(1, NUMERICS.to_vec())
-        }
     }
 }