diff --git a/datafusion/functions/src/macros.rs b/datafusion/functions/src/macros.rs index f196870e97228..8a6607c46b45e 100644 --- a/datafusion/functions/src/macros.rs +++ b/datafusion/functions/src/macros.rs @@ -207,20 +207,22 @@ macro_rules! downcast_arg { /// $NAME: the name of the function /// $UNARY_FUNC: the unary function to apply to the argument /// $OUTPUT_ORDERING: the output ordering calculation method of the function +/// $STRICT: whether the function returns NULL when any argument is NULL /// $GET_DOC: the function to get the documentation of the UDF macro_rules! make_math_unary_udf { - ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $GET_DOC:expr) => { + ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr) => { make_math_unary_udf!( $UDF, $NAME, $UNARY_FUNC, $OUTPUT_ORDERING, $EVALUATE_BOUNDS, + $STRICT, $GET_DOC, None:: Result<()>> ); }; - ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $GET_DOC:expr, $VALIDATOR:expr) => { + ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr, $VALIDATOR:expr) => { $crate::make_udf_function!($NAME::$UDF, $NAME); mod $NAME { @@ -273,6 +275,10 @@ macro_rules! make_math_unary_udf { } } + fn is_strict(&self) -> bool { + $STRICT + } + fn output_ordering( &self, input: &[ExprProperties], @@ -354,9 +360,10 @@ macro_rules! make_math_unary_udf { /// $NAME: the name of the function /// $BINARY_FUNC: the binary function to apply to the argument /// $OUTPUT_ORDERING: the output ordering calculation method of the function +/// $STRICT: whether the function returns NULL when any argument is NULL /// $GET_DOC: the function to get the documentation of the UDF macro_rules! make_math_binary_udf { - ($UDF:ident, $NAME:ident, $BINARY_FUNC:ident, $OUTPUT_ORDERING:expr, $GET_DOC:expr) => { + ($UDF:ident, $NAME:ident, $BINARY_FUNC:ident, $OUTPUT_ORDERING:expr, $STRICT:expr, $GET_DOC:expr) => { $crate::make_udf_function!($NAME::$UDF, $NAME); mod $NAME { @@ -414,6 +421,10 @@ macro_rules! make_math_binary_udf { } } + fn is_strict(&self) -> bool { + $STRICT + } + fn output_ordering( &self, input: &[ExprProperties], diff --git a/datafusion/functions/src/math/ceil.rs b/datafusion/functions/src/math/ceil.rs index 395cb4eae03f5..7b2c0c35e4cad 100644 --- a/datafusion/functions/src/math/ceil.rs +++ b/datafusion/functions/src/math/ceil.rs @@ -89,6 +89,10 @@ impl ScalarUDFImpl for CeilFunc { } } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let arg = &args.args[0]; diff --git a/datafusion/functions/src/math/cot.rs b/datafusion/functions/src/math/cot.rs index 24f0a412e3a8a..ca207778f7eb7 100644 --- a/datafusion/functions/src/math/cot.rs +++ b/datafusion/functions/src/math/cot.rs @@ -86,6 +86,10 @@ impl ScalarUDFImpl for CotFunc { } } + fn is_strict(&self) -> bool { + true + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } diff --git a/datafusion/functions/src/math/factorial.rs b/datafusion/functions/src/math/factorial.rs index 3b4f973f19d62..f4e9b60dd3799 100644 --- a/datafusion/functions/src/math/factorial.rs +++ b/datafusion/functions/src/math/factorial.rs @@ -76,6 +76,10 @@ impl ScalarUDFImpl for FactorialFunc { Ok(Int64) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/floor.rs b/datafusion/functions/src/math/floor.rs index e02aa141c5b71..4ab6e0eb5effd 100644 --- a/datafusion/functions/src/math/floor.rs +++ b/datafusion/functions/src/math/floor.rs @@ -129,6 +129,10 @@ impl ScalarUDFImpl for FloorFunc { } } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let arg = &args.args[0]; diff --git a/datafusion/functions/src/math/gcd.rs b/datafusion/functions/src/math/gcd.rs index aeddc3f27c409..6a4e69620e060 100644 --- a/datafusion/functions/src/math/gcd.rs +++ b/datafusion/functions/src/math/gcd.rs @@ -82,6 +82,10 @@ impl ScalarUDFImpl for GcdFunc { Ok(arg_types[0].clone()) } + fn is_strict(&self) -> bool { + true + } + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let [arg1, arg2] = take_function_args(self.name(), arg_types)?; diff --git a/datafusion/functions/src/math/iszero.rs b/datafusion/functions/src/math/iszero.rs index de6fc669692ee..62cfdd4c839ec 100644 --- a/datafusion/functions/src/math/iszero.rs +++ b/datafusion/functions/src/math/iszero.rs @@ -85,6 +85,10 @@ impl ScalarUDFImpl for IsZeroFunc { Ok(Boolean) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/lcm.rs b/datafusion/functions/src/math/lcm.rs index 245dba0ba3938..248e4b93ffd8b 100644 --- a/datafusion/functions/src/math/lcm.rs +++ b/datafusion/functions/src/math/lcm.rs @@ -78,6 +78,10 @@ impl ScalarUDFImpl for LcmFunc { Ok(arg_types[0].clone()) } + fn is_strict(&self) -> bool { + true + } + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let [arg1, arg2] = take_function_args(self.name(), arg_types)?; diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs index 2ca2ed1b572be..11d76d8086be9 100644 --- a/datafusion/functions/src/math/log.rs +++ b/datafusion/functions/src/math/log.rs @@ -203,6 +203,10 @@ impl ScalarUDFImpl for LogFunc { } } + fn is_strict(&self) -> bool { + true + } + fn output_ordering(&self, input: &[ExprProperties]) -> Result { let (base_sort_properties, num_sort_properties) = if input.len() == 1 { // log(x) defaults to log(10, x) diff --git a/datafusion/functions/src/math/mod.rs b/datafusion/functions/src/math/mod.rs index a5d45380ecf0a..4b79866895d84 100644 --- a/datafusion/functions/src/math/mod.rs +++ b/datafusion/functions/src/math/mod.rs @@ -60,6 +60,7 @@ make_math_unary_udf!( acos, super::acos_order, super::bounds::acos_bounds, + true, super::get_acos_doc ); make_math_unary_udf!( @@ -68,6 +69,7 @@ make_math_unary_udf!( acosh, super::acosh_order, super::bounds::acosh_bounds, + true, super::get_acosh_doc ); make_math_unary_udf!( @@ -76,6 +78,7 @@ make_math_unary_udf!( asin, super::asin_order, super::bounds::asin_bounds, + true, super::get_asin_doc ); make_math_unary_udf!( @@ -84,6 +87,7 @@ make_math_unary_udf!( asinh, super::asinh_order, super::bounds::unbounded_bounds, + true, super::get_asinh_doc ); make_math_unary_udf!( @@ -92,6 +96,7 @@ make_math_unary_udf!( atan, super::atan_order, super::bounds::atan_bounds, + true, super::get_atan_doc ); make_math_unary_udf!( @@ -100,6 +105,7 @@ make_math_unary_udf!( atanh, super::atanh_order, super::bounds::unbounded_bounds, + true, super::get_atanh_doc ); make_math_binary_udf!( @@ -107,6 +113,7 @@ make_math_binary_udf!( atan2, atan2, super::atan2_order, + true, super::get_atan2_doc ); make_math_unary_udf!( @@ -115,6 +122,7 @@ make_math_unary_udf!( cbrt, super::cbrt_order, super::bounds::unbounded_bounds, + true, super::get_cbrt_doc ); make_udf_function!(ceil::CeilFunc, ceil); @@ -124,6 +132,7 @@ make_math_unary_udf!( cos, super::cos_order, super::bounds::cos_bounds, + true, super::get_cos_doc ); make_math_unary_udf!( @@ -132,6 +141,7 @@ make_math_unary_udf!( cosh, super::cosh_order, super::bounds::cosh_bounds, + true, super::get_cosh_doc ); make_udf_function!(cot::CotFunc, cot); @@ -141,6 +151,7 @@ make_math_unary_udf!( to_degrees, super::degrees_order, super::bounds::unbounded_bounds, + true, super::get_degrees_doc ); make_math_unary_udf!( @@ -149,6 +160,7 @@ make_math_unary_udf!( exp, super::exp_order, super::bounds::exp_bounds, + true, super::get_exp_doc ); make_udf_function!(factorial::FactorialFunc, factorial); @@ -164,6 +176,7 @@ make_math_unary_udf!( ln, super::ln_order, super::bounds::unbounded_bounds, + true, super::get_ln_doc ); make_math_unary_udf!( @@ -172,6 +185,7 @@ make_math_unary_udf!( log2, super::log2_order, super::bounds::unbounded_bounds, + true, super::get_log2_doc ); make_math_unary_udf!( @@ -180,6 +194,7 @@ make_math_unary_udf!( log10, super::log10_order, super::bounds::unbounded_bounds, + true, super::get_log10_doc ); make_udf_function!(nanvl::NanvlFunc, nanvl); @@ -191,6 +206,7 @@ make_math_unary_udf!( to_radians, super::radians_order, super::bounds::radians_bounds, + true, super::get_radians_doc ); make_udf_function!(random::RandomFunc, random); @@ -202,6 +218,7 @@ make_math_unary_udf!( sin, super::sin_order, super::bounds::sin_bounds, + true, super::get_sin_doc ); make_math_unary_udf!( @@ -210,6 +227,7 @@ make_math_unary_udf!( sinh, super::sinh_order, super::bounds::unbounded_bounds, + true, super::get_sinh_doc ); make_math_unary_udf!( @@ -218,6 +236,7 @@ make_math_unary_udf!( sqrt, super::sqrt_order, super::bounds::sqrt_bounds, + true, super::get_sqrt_doc, Some(super::validate_sqrt_input) ); @@ -227,6 +246,7 @@ make_math_unary_udf!( tan, super::tan_order, super::bounds::unbounded_bounds, + true, super::get_tan_doc ); make_math_unary_udf!( @@ -235,10 +255,146 @@ make_math_unary_udf!( tanh, super::tanh_order, super::bounds::tanh_bounds, + true, super::get_tanh_doc ); make_udf_function!(trunc::TruncFunc, trunc); +#[cfg(test)] +mod strict_tests { + use super::*; + use arrow::datatypes::Field; + use datafusion_common::ScalarValue; + use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, + }; + use std::sync::Arc; + + #[test] + fn strict_math_functions_propagate_nulls() { + let cases = vec![ + (abs(), vec![ScalarValue::from(1.0)]), + (acos(), vec![ScalarValue::from(0.5)]), + (acosh(), vec![ScalarValue::from(1.5)]), + (asin(), vec![ScalarValue::from(0.5)]), + (asinh(), vec![ScalarValue::from(0.5)]), + (atan(), vec![ScalarValue::from(0.5)]), + ( + atan2(), + vec![ScalarValue::from(0.5), ScalarValue::from(1.0)], + ), + (atanh(), vec![ScalarValue::from(0.5)]), + (cbrt(), vec![ScalarValue::from(8.0)]), + (ceil(), vec![ScalarValue::from(1.5)]), + (cos(), vec![ScalarValue::from(0.5)]), + (cosh(), vec![ScalarValue::from(0.5)]), + (cot(), vec![ScalarValue::from(0.5)]), + (degrees(), vec![ScalarValue::from(0.5)]), + (exp(), vec![ScalarValue::from(0.5)]), + (factorial(), vec![ScalarValue::from(5_i64)]), + (floor(), vec![ScalarValue::from(1.5)]), + ( + gcd(), + vec![ScalarValue::from(48_i64), ScalarValue::from(18_i64)], + ), + (isnan(), vec![ScalarValue::from(1.0)]), + (iszero(), vec![ScalarValue::from(1.0)]), + ( + lcm(), + vec![ScalarValue::from(4_i64), ScalarValue::from(5_i64)], + ), + (ln(), vec![ScalarValue::from(2.0)]), + (log(), vec![ScalarValue::from(10.0)]), + ( + log(), + vec![ScalarValue::from(10.0), ScalarValue::from(100.0)], + ), + (log2(), vec![ScalarValue::from(2.0)]), + (log10(), vec![ScalarValue::from(10.0)]), + ( + power(), + vec![ScalarValue::from(2.0), ScalarValue::from(3.0)], + ), + (radians(), vec![ScalarValue::from(90.0)]), + (round(), vec![ScalarValue::from(1.5)]), + ( + round(), + vec![ScalarValue::from(1.5), ScalarValue::from(1_i32)], + ), + (signum(), vec![ScalarValue::from(-1.0)]), + (sin(), vec![ScalarValue::from(0.5)]), + (sinh(), vec![ScalarValue::from(0.5)]), + (sqrt(), vec![ScalarValue::from(4.0)]), + (tan(), vec![ScalarValue::from(0.5)]), + (tanh(), vec![ScalarValue::from(0.5)]), + (trunc(), vec![ScalarValue::from(1.5)]), + ( + trunc(), + vec![ScalarValue::from(1.5), ScalarValue::from(1_i64)], + ), + ]; + + for (func, valid_args) in cases { + assert!(func.is_strict(), "{} should be marked strict", func.name()); + + for null_mask in 0..(1 << valid_args.len()) { + let mut args = valid_args.clone(); + for (arg_idx, arg) in args.iter_mut().enumerate() { + if null_mask & (1 << arg_idx) != 0 { + *arg = ScalarValue::try_new_null(&arg.data_type()).unwrap(); + } + } + + let result = + invoke_with_scalar_args(&func, args).unwrap_or_else(|error| { + panic!( + "{} failed for NULL mask {null_mask:b}: {error}", + func.name() + ) + }); + let expected_null = null_mask != 0; + let result = result.into_array(1).unwrap(); + assert_eq!( + result.null_count() == result.len(), + expected_null, + "{} returned {result:?} for NULL mask {null_mask:0width$b}", + func.name(), + width = valid_args.len(), + ); + } + } + } + + fn invoke_with_scalar_args( + func: &ScalarUDF, + args: Vec, + ) -> Result { + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Arc::new(Field::new( + format!("arg_{idx}"), + arg.data_type(), + arg.is_null(), + )) + }) + .collect::>(); + let scalar_arguments = args.iter().map(Some).collect::>(); + let return_field = func.return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + })?; + func.invoke_with_args(ScalarFunctionArgs { + args: args.into_iter().map(ColumnarValue::Scalar).collect(), + arg_fields, + number_rows: 1, + return_field, + config_options: Arc::new(Default::default()), + }) + } +} + pub mod expr_fn { export_functions!( (abs, "returns the absolute value of a given number", num), diff --git a/datafusion/functions/src/math/nans.rs b/datafusion/functions/src/math/nans.rs index c5ea2fa079a45..c313db30378bf 100644 --- a/datafusion/functions/src/math/nans.rs +++ b/datafusion/functions/src/math/nans.rs @@ -83,6 +83,10 @@ impl ScalarUDFImpl for IsNanFunc { Ok(DataType::Boolean) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 252a3ea0b31d7..54ba0d3581e3a 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -99,6 +99,10 @@ impl ScalarUDFImpl for PowerFunc { Ok(DataType::Float64) } + fn is_strict(&self) -> bool { + true + } + fn aliases(&self) -> &[String] { &self.aliases } diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs index 62f1c3540b9ce..8fb81366bf2ac 100644 --- a/datafusion/functions/src/math/round.rs +++ b/datafusion/functions/src/math/round.rs @@ -227,6 +227,10 @@ impl ScalarUDFImpl for RoundFunc { "round" } + fn is_strict(&self) -> bool { + true + } + fn signature(&self) -> &Signature { &self.signature } diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index 8c8eeacf12394..05b78fcffe2a7 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -86,6 +86,10 @@ impl ScalarUDFImpl for SignumFunc { } } + fn is_strict(&self) -> bool { + true + } + fn output_ordering(&self, input: &[ExprProperties]) -> Result { // Non-decreasing for all real numbers x. Ok(input[0].sort_properties) diff --git a/datafusion/functions/src/math/trunc.rs b/datafusion/functions/src/math/trunc.rs index 7b11e19bdb648..00bbf8901eb2a 100644 --- a/datafusion/functions/src/math/trunc.rs +++ b/datafusion/functions/src/math/trunc.rs @@ -122,6 +122,10 @@ impl ScalarUDFImpl for TruncFunc { "trunc" } + fn is_strict(&self) -> bool { + true + } + fn signature(&self) -> &Signature { &self.signature } diff --git a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt index 52ae3e37efca0..afd491b0b64c8 100644 --- a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt +++ b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt @@ -717,6 +717,147 @@ logical_plan 05)--------TableScan: t1 projection=[a] 06)--------TableScan: t2 projection=[x, y] +### +### Strict math function matrix +### + +# Unary function on the nullable side of a LEFT JOIN -> INNER JOIN. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where ceil(t2.y) > 150; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: ceil(CAST(t2.y AS Float64)) > Float64(150) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where ceil(t2.y) > 150; +---- +2 + +# Unary function on the nullable side of a RIGHT JOIN -> INNER JOIN. +query TT +explain +select t2.x +from t1 right join t2 on t1.a = t2.x +where floor(t1.b) > 15; +---- +logical_plan +01)Projection: t2.x +02)--Inner Join: t1.a = t2.x +03)----Projection: t1.a +04)------Filter: CAST(t1.b AS Float64) >= Float64(16) +05)--------TableScan: t1 projection=[a, b] +06)----TableScan: t2 projection=[x] + +query I rowsort +select t2.x +from t1 right join t2 on t1.a = t2.x +where floor(t1.b) > 15; +---- +2 + +# Binary function with the nullable column as its first argument. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where atan2(t2.y, 1) > 1; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: atan2(CAST(t2.y AS Float64), Float64(1)) > Float64(1) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where atan2(t2.y, 1) > 1; +---- +1 +2 + +# Binary function with the nullable column as its second argument. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where power(2, t2.y) > 100; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: power(Float64(2), CAST(t2.y AS Float64)) > Float64(100) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where power(2, t2.y) > 100; +---- +1 +2 + +# A strict function on only the right side of a FULL JOIN -> RIGHT JOIN. +query TT +explain +select t1.a, t2.y +from t1 full join t2 on t1.a = t2.x +where round(t2.y, -2) >= 100; +---- +logical_plan +01)Projection: t1.a, t2.y +02)--Right Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Filter: round(t2.y, Int32(-2)) >= Int32(100) +05)------TableScan: t2 projection=[x, y] + +query II rowsort +select t1.a, t2.y +from t1 full join t2 on t1.a = t2.x +where round(t2.y, -2) >= 100; +---- +1 100 +2 200 +NULL 300 + +# A strict function on only the left side of a FULL JOIN -> LEFT JOIN. +query TT +explain +select t1.a, t1.b +from t1 full join t2 on t1.a = t2.x +where trunc(t1.b, -1) >= 10; +---- +logical_plan +01)Projection: t1.a, t1.b +02)--Left Join: t1.a = t2.x +03)----Filter: trunc(CAST(t1.b AS Float64), Int64(-1)) >= Float64(10) +04)------TableScan: t1 projection=[a, b] +05)----TableScan: t2 projection=[x] + +query II rowsort +select t1.a, t1.b +from t1 full join t2 on t1.a = t2.x +where trunc(t1.b, -1) >= 10; +---- +1 10 +2 20 +3 30 +NULL 40 + ### ### Cleanup ###