From d785779ce2f579c8ac9ff848612994579364e2b1 Mon Sep 17 00:00:00 2001 From: LorrensP-2158466 Date: Tue, 25 Jun 2024 19:45:13 +0200 Subject: [PATCH 1/3] checked overflow on integer parameter --- datafusion/functions/src/math/power.rs | 37 +++++++++++++++++--------- 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 7677e8b2af95..d8ccbfac7508 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -17,9 +17,11 @@ //! Math function: `power()`. -use arrow::datatypes::DataType; +use arrow::datatypes::{ArrowNativeTypeOp, DataType}; + use datafusion_common::{ - exec_err, plan_datafusion_err, DataFusionError, Result, ScalarValue, + arrow_datafusion_err, exec_datafusion_err, exec_err, plan_datafusion_err, + DataFusionError, Result, ScalarValue, }; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo}; @@ -94,14 +96,25 @@ impl ScalarUDFImpl for PowerFunc { { f64::powf } )), - DataType::Int64 => Arc::new(make_function_inputs2!( - &args[0], - &args[1], - "base", - "exponent", - Int64Array, - { i64::pow } - )), + DataType::Int64 => { + let bases = downcast_arg!(&args[0], "base", Int64Array); + let exponents = downcast_arg!(&args[1], "exponent", Int64Array); + bases + .iter() + .zip(exponents.iter()) + .map(|(base, exp)| match (base, exp) { + (Some(base), Some(exp)) => Ok(Some(base.pow_checked( + exp.try_into().map_err(|_| { + exec_datafusion_err!( + "Can't use negative exponents: {exp} in integer computation, please use Float." + ) + })?, + ).map_err(|e| arrow_datafusion_err!(e))?)), + _ => Ok(None), + }) + .collect::>() + .map(Arc::new)? as ArrayRef + } other => { return exec_err!( @@ -192,8 +205,8 @@ mod tests { #[test] fn test_power_i64() { let args = [ - ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5]))), // base - ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4]))), // exponent + ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5, 5]))), // base + ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4, -4]))), // exponent ]; let result = PowerFunc::new() From 36ba2168d742aab49d52485483efc0022bbd78eb Mon Sep 17 00:00:00 2001 From: LorrensP-2158466 Date: Tue, 25 Jun 2024 19:45:25 +0200 Subject: [PATCH 2/3] added slt test --- datafusion/sqllogictest/test_files/math.slt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datafusion/sqllogictest/test_files/math.slt b/datafusion/sqllogictest/test_files/math.slt index 3a433c76fc57..58e4f90653e3 100644 --- a/datafusion/sqllogictest/test_files/math.slt +++ b/datafusion/sqllogictest/test_files/math.slt @@ -575,3 +575,6 @@ select gcd(-9223372036854775808, -9223372036854775808); query error DataFusion error: Arrow error: Compute error: Signed integer overflow in LCM\(\-9223372036854775808, \-9223372036854775808\) select lcm(-9223372036854775808, -9223372036854775808); + +query error DataFusion error: Arrow error: Compute error: Overflow happened on: 2107754225 \^ 1221660777 +select power(2107754225, 1221660777); From 6a32515c86deebb96971e625ad513da6b33a9689 Mon Sep 17 00:00:00 2001 From: LorrensP-2158466 Date: Tue, 25 Jun 2024 20:33:44 +0200 Subject: [PATCH 3/3] fix: forgot to remove failling test case --- datafusion/functions/src/math/power.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index d8ccbfac7508..5b790fb56ddf 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -205,8 +205,8 @@ mod tests { #[test] fn test_power_i64() { let args = [ - ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5, 5]))), // base - ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4, -4]))), // exponent + ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5]))), // base + ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4]))), // exponent ]; let result = PowerFunc::new()