This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8aad208299 handle overflow in gcd and return this as an error (#11057)
8aad208299 is described below
commit 8aad2082994bdea14eaf7d0e8c177aa3ce9af67b
Author: Lorrens Pantelis <[email protected]>
AuthorDate: Sun Jun 23 12:22:58 2024 +0200
handle overflow in gcd and return this as an error (#11057)
* handle overflow in gcd and return this as an error
* Add SLT test
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/functions/src/math/gcd.rs | 62 +++++++++++++++++++++--------
datafusion/functions/src/math/lcm.rs | 36 +++++++++++------
datafusion/sqllogictest/test_files/math.slt | 11 +++++
3 files changed, 81 insertions(+), 28 deletions(-)
diff --git a/datafusion/functions/src/math/gcd.rs
b/datafusion/functions/src/math/gcd.rs
index d0199f7a22..6346f96d2c 100644
--- a/datafusion/functions/src/math/gcd.rs
+++ b/datafusion/functions/src/math/gcd.rs
@@ -16,6 +16,7 @@
// under the License.
use arrow::array::{ArrayRef, Int64Array};
+use arrow::error::ArrowError;
use std::any::Any;
use std::mem::swap;
use std::sync::Arc;
@@ -24,7 +25,7 @@ use arrow::datatypes::DataType;
use arrow::datatypes::DataType::Int64;
use crate::utils::make_scalar_function;
-use datafusion_common::{exec_err, DataFusionError, Result};
+use datafusion_common::{arrow_datafusion_err, exec_err, DataFusionError,
Result};
use datafusion_expr::ColumnarValue;
use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
@@ -73,26 +74,31 @@ impl ScalarUDFImpl for GcdFunc {
/// Gcd SQL function
fn gcd(args: &[ArrayRef]) -> Result<ArrayRef> {
match args[0].data_type() {
- Int64 => Ok(Arc::new(make_function_inputs2!(
- &args[0],
- &args[1],
- "x",
- "y",
- Int64Array,
- Int64Array,
- { compute_gcd }
- )) as ArrayRef),
+ Int64 => {
+ let arg1 = downcast_arg!(&args[0], "x", Int64Array);
+ let arg2 = downcast_arg!(&args[1], "y", Int64Array);
+
+ Ok(arg1
+ .iter()
+ .zip(arg2.iter())
+ .map(|(a1, a2)| match (a1, a2) {
+ (Some(a1), Some(a2)) => Ok(Some(compute_gcd(a1, a2)?)),
+ _ => Ok(None),
+ })
+ .collect::<Result<Int64Array>>()
+ .map(Arc::new)? as ArrayRef)
+ }
other => exec_err!("Unsupported data type {other:?} for function gcd"),
}
}
/// Computes greatest common divisor using Binary GCD algorithm.
-pub fn compute_gcd(x: i64, y: i64) -> i64 {
+pub fn compute_gcd(x: i64, y: i64) -> Result<i64> {
if x == 0 {
- return y;
+ return Ok(y);
}
if y == 0 {
- return x;
+ return Ok(x);
}
let mut a = x.unsigned_abs();
@@ -113,7 +119,11 @@ pub fn compute_gcd(x: i64, y: i64) -> i64 {
if b == 0 {
// because the input values are i64, casting this back to i64 is
safe
- return (a << shift) as i64;
+ return (a << shift).try_into().map_err(|_| {
+ arrow_datafusion_err!(ArrowError::ComputeError(format!(
+ "Signed integer overflow in GCD({x}, {y})"
+ )))
+ });
}
}
}
@@ -122,10 +132,13 @@ pub fn compute_gcd(x: i64, y: i64) -> i64 {
mod test {
use std::sync::Arc;
- use arrow::array::{ArrayRef, Int64Array};
+ use arrow::{
+ array::{ArrayRef, Int64Array},
+ error::ArrowError,
+ };
use crate::math::gcd::gcd;
- use datafusion_common::cast::as_int64_array;
+ use datafusion_common::{cast::as_int64_array, DataFusionError};
#[test]
fn test_gcd_i64() {
@@ -143,4 +156,21 @@ mod test {
assert_eq!(ints.value(2), 5);
assert_eq!(ints.value(3), 8);
}
+
+ #[test]
+ fn overflow_on_both_param_i64_min() {
+ let args: Vec<ArrayRef> = vec![
+ Arc::new(Int64Array::from(vec![i64::MIN])), // x
+ Arc::new(Int64Array::from(vec![i64::MIN])), // y
+ ];
+
+ match gcd(&args) {
+ // we expect a overflow
+ Err(DataFusionError::ArrowError(ArrowError::ComputeError(_), _))
=> {}
+ Err(_) => {
+ panic!("failed to initialize function gcd")
+ }
+ Ok(_) => panic!("GCD({0}, {0}) should have overflown", i64::MIN),
+ };
+ }
}
diff --git a/datafusion/functions/src/math/lcm.rs
b/datafusion/functions/src/math/lcm.rs
index 3674f7371d..d58bc5ef0b 100644
--- a/datafusion/functions/src/math/lcm.rs
+++ b/datafusion/functions/src/math/lcm.rs
@@ -22,7 +22,8 @@ use arrow::array::{ArrayRef, Int64Array};
use arrow::datatypes::DataType;
use arrow::datatypes::DataType::Int64;
-use datafusion_common::{exec_err, DataFusionError, Result};
+use arrow::error::ArrowError;
+use datafusion_common::{arrow_datafusion_err, exec_err, DataFusionError,
Result};
use datafusion_expr::ColumnarValue;
use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
@@ -78,21 +79,32 @@ fn lcm(args: &[ArrayRef]) -> Result<ArrayRef> {
let b = y.wrapping_abs();
if a == 0 || b == 0 {
- return 0;
+ return Ok(0);
+ }
+ match compute_gcd(a, b) {
+ Ok(gcd) => Ok(a / gcd * b),
+ // change the error string to use LCM instead of GCD
+ Err(_) =>
Err(arrow_datafusion_err!(ArrowError::ComputeError(format!(
+ "Signed integer overflow in LCM({x}, {y})"
+ )))),
}
- a / compute_gcd(a, b) * b
};
match args[0].data_type() {
- Int64 => Ok(Arc::new(make_function_inputs2!(
- &args[0],
- &args[1],
- "x",
- "y",
- Int64Array,
- Int64Array,
- { compute_lcm }
- )) as ArrayRef),
+ Int64 => {
+ let arg1 = downcast_arg!(&args[0], "x", Int64Array);
+ let arg2 = downcast_arg!(&args[1], "y", Int64Array);
+
+ Ok(arg1
+ .iter()
+ .zip(arg2.iter())
+ .map(|(a1, a2)| match (a1, a2) {
+ (Some(a1), Some(a2)) => Ok(Some(compute_lcm(a1, a2)?)),
+ _ => Ok(None),
+ })
+ .collect::<Result<Int64Array>>()
+ .map(Arc::new)? as ArrayRef)
+ }
other => exec_err!("Unsupported data type {other:?} for function lcm"),
}
}
diff --git a/datafusion/sqllogictest/test_files/math.slt
b/datafusion/sqllogictest/test_files/math.slt
index 78efbb3f56..3a433c76fc 100644
--- a/datafusion/sqllogictest/test_files/math.slt
+++ b/datafusion/sqllogictest/test_files/math.slt
@@ -564,3 +564,14 @@ SELECT c1%0 FROM test_non_nullable_decimal
statement ok
drop table test_non_nullable_decimal
+
+query II
+select gcd(20, 1000), lcm(20, 1000);
+----
+20 1000
+
+query error DataFusion error: Arrow error: Compute error: Signed integer
overflow in GCD\(\-9223372036854775808, \-9223372036854775808\)
+select gcd(-9223372036854775808, -9223372036854775808);
+
+query error DataFusion error: Arrow error: Compute error: Signed integer
overflow in LCM\(\-9223372036854775808, \-9223372036854775808\)
+select lcm(-9223372036854775808, -9223372036854775808);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
