diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 7677e8b2af95..5b790fb56ddf 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!( 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);