Skip to content

Commit

Permalink
cleanup return_type()
Browse files Browse the repository at this point in the history
  • Loading branch information
2010YOUY01 committed Nov 8, 2023
1 parent 0506a5c commit c871714
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 42 deletions.
38 changes: 9 additions & 29 deletions datafusion/expr/src/built_in_function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,12 @@ use crate::signature::TIMEZONE_WILDCARD;
use crate::type_coercion::binary::get_wider_type;
use crate::type_coercion::functions::data_types;
use crate::{
conditional_expressions, struct_expressions, utils, FuncMonotonicity, Signature,
conditional_expressions, struct_expressions, FuncMonotonicity, Signature,
TypeSignature, Volatility,
};

use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit};
use datafusion_common::{
internal_err, plan_datafusion_err, plan_err, DataFusionError, Result,
};
use datafusion_common::{internal_err, plan_err, DataFusionError, Result};

use strum::IntoEnumIterator;
use strum_macros::EnumIter;
Expand Down Expand Up @@ -483,38 +481,20 @@ impl BuiltinScalarFunction {
}

/// Returns the output [`DataType`] of this function
///
/// This method should be invoked only after `input_expr_types` have been validated
/// against the function's `TypeSignature` using `type_coercion::functions::data_types()`.
///
/// This method will:
/// 1. Perform additional checks on `input_expr_types` that are beyond the scope of `TypeSignature` validation.
/// 2. Deduce the output `DataType` based on the provided `input_expr_types`.
pub fn return_type(self, input_expr_types: &[DataType]) -> Result<DataType> {
use DataType::*;
use TimeUnit::*;

// Note that this function *must* return the same type that the respective physical expression returns
// or the execution panics.

if input_expr_types.is_empty()
&& !self.signature().type_signature.supports_zero_argument()
{
return plan_err!(
"{}",
utils::generate_signature_error_msg(
&format!("{self}"),
self.signature(),
input_expr_types
)
);
}

// verify that this is a valid set of data types for this function
data_types(input_expr_types, &self.signature()).map_err(|_| {
plan_datafusion_err!(
"{}",
utils::generate_signature_error_msg(
&format!("{self}"),
self.signature(),
input_expr_types,
)
)
})?;

// the return type of the built in function.
// Some built-in functions' return type depends on the incoming type.
match self {
Expand Down
19 changes: 16 additions & 3 deletions datafusion/expr/src/expr_schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ use crate::expr::{
};
use crate::field_util::GetFieldAccessSchema;
use crate::type_coercion::binary::get_result_type;
use crate::{LogicalPlan, Projection, Subquery};
use crate::type_coercion::functions::data_types;
use crate::{utils, LogicalPlan, Projection, Subquery};
use arrow::compute::can_cast_types;
use arrow::datatypes::{DataType, Field};
use datafusion_common::{
Expand Down Expand Up @@ -89,12 +90,24 @@ impl ExprSchemable for Expr {
Ok((fun.return_type)(&data_types)?.as_ref().clone())
}
Expr::ScalarFunction(ScalarFunction { fun, args }) => {
let data_types = args
let arg_data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;

fun.return_type(&data_types)
// verify that input data types is consistent with function's `TypeSignature`
data_types(&arg_data_types, &fun.signature()).map_err(|_| {
plan_datafusion_err!(
"{}",
utils::generate_signature_error_msg(
&format!("{fun}"),
fun.signature(),
&arg_data_types,
)
)
})?;

fun.return_type(&arg_data_types)
}
Expr::WindowFunction(WindowFunction { fun, args, .. }) => {
let data_types = args
Expand Down
11 changes: 10 additions & 1 deletion datafusion/expr/src/type_coercion/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,17 @@ pub fn data_types(
signature: &Signature,
) -> Result<Vec<DataType>> {
if current_types.is_empty() {
return Ok(vec![]);
if signature.type_signature.supports_zero_argument() {
return Ok(vec![]);
} else {
return plan_err!(
"Coercion from {:?} to the signature {:?} failed.",
current_types,
&signature.type_signature
);
}
}

let valid_types = get_valid_types(&signature.type_signature, current_types)?;

if valid_types
Expand Down
15 changes: 6 additions & 9 deletions datafusion/physical-expr/src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ use arrow::{
use datafusion_common::{internal_err, DataFusionError, Result, ScalarValue};
pub use datafusion_expr::FuncMonotonicity;
use datafusion_expr::{
BuiltinScalarFunction, ColumnarValue, ScalarFunctionImplementation,
type_coercion::functions::data_types, BuiltinScalarFunction, ColumnarValue,
ScalarFunctionImplementation,
};
use std::ops::Neg;
use std::sync::Arc;
Expand All @@ -65,6 +66,9 @@ pub fn create_physical_expr(
.map(|e| e.data_type(input_schema))
.collect::<Result<Vec<_>>>()?;

// verify that input data types is consistent with function's `TypeSignature`
data_types(&input_expr_types, &fun.signature())?;

let data_type = fun.return_type(&input_expr_types)?;

let fun_expr: ScalarFunctionImplementation = match fun {
Expand Down Expand Up @@ -2952,14 +2956,7 @@ mod tests {
"Builtin scalar function {fun} does not support empty arguments"
);
}
Err(DataFusionError::Plan(err)) => {
if !err
.contains("No function matches the given name and argument types")
{
return plan_err!(
"Builtin scalar function {fun} didn't got the right error message with empty arguments");
}
}
Err(DataFusionError::Plan(_)) => return Ok(()),
Err(..) => {
return internal_err!(
"Builtin scalar function {fun} didn't got the right error with empty arguments");
Expand Down

0 comments on commit c871714

Please sign in to comment.