diff --git a/datafusion/functions/src/math/gcd.rs b/datafusion/functions/src/math/gcd.rs index 90153dd752f9..78deb2a7965b 100644 --- a/datafusion/functions/src/math/gcd.rs +++ b/datafusion/functions/src/math/gcd.rs @@ -15,17 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{new_null_array, ArrayRef, AsArray, Int64Array}; +use arrow::array::{new_null_array, ArrayRef, AsArray, Int64Array, PrimitiveArray}; +use arrow::compute::try_binary; +use arrow::datatypes::{DataType, Int64Type}; use arrow::error::ArrowError; use std::any::Any; use std::mem::swap; use std::sync::Arc; -use arrow::datatypes::{DataType, Int64Type}; - use datafusion_common::{ - arrow_datafusion_err, exec_err, internal_datafusion_err, DataFusionError, Result, - ScalarValue, + exec_err, internal_datafusion_err, internal_err, Result, ScalarValue, }; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -113,17 +112,20 @@ impl ScalarUDFImpl for GcdFunc { } fn compute_gcd_for_arrays(a: &ArrayRef, b: &ArrayRef) -> Result { - let result: Result = a - .as_primitive::() - .iter() - .zip(b.as_primitive::().iter()) - .map(|(a, b)| match (a, b) { - (Some(a), Some(b)) => Ok(Some(compute_gcd(a, b)?)), - _ => Ok(None), + let a = a.as_primitive::(); + let b = b.as_primitive::(); + if a.len() != b.len() { + return internal_err!( + "Length of arguments for function gcd do not match: {} vs {}", + a.len(), + b.len() + ); + } + try_binary(a, b, compute_gcd) + .map(|arr: PrimitiveArray| { + ColumnarValue::Array(Arc::new(arr) as ArrayRef) }) - .collect(); - - result.map(|arr| ColumnarValue::Array(Arc::new(arr) as ArrayRef)) + .map_err(Into::into) // convert ArrowError to DataFusionError } fn compute_gcd_with_scalar(arr: &ArrayRef, scalar: Option) -> Result { @@ -171,14 +173,12 @@ pub(super) fn unsigned_gcd(mut a: u64, mut b: u64) -> u64 { } /// Computes greatest common divisor using Binary GCD algorithm. -pub fn compute_gcd(x: i64, y: i64) -> Result { +pub fn compute_gcd(x: i64, y: i64) -> Result { let a = x.unsigned_abs(); let b = y.unsigned_abs(); let r = unsigned_gcd(a, b); // gcd(i64::MIN, i64::MIN) = i64::MIN.unsigned_abs() cannot fit into i64 r.try_into().map_err(|_| { - arrow_datafusion_err!(ArrowError::ComputeError(format!( - "Signed integer overflow in GCD({x}, {y})" - ))) + ArrowError::ComputeError(format!("Signed integer overflow in GCD({x}, {y})")) }) }