diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index 9e731f3f653..b1e07e7dbb7 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -105,8 +105,8 @@ where } /// This is similar to `math_op` as it performs given operation between two input primitive arrays. -/// But the given can return `None` if overflow is detected. For the case, this function returns an -/// `Err`. +/// But the given operation can return `None` if overflow is detected. For the case, this function +/// returns an `Err`. fn math_checked_op( left: &PrimitiveArray, right: &PrimitiveArray, @@ -149,6 +149,54 @@ where Ok(PrimitiveArray::::from_iter(values)) } +/// This is similar to `math_checked_op` but just for divide op. +fn math_checked_divide( + left: &PrimitiveArray, + right: &PrimitiveArray, + op: F, +) -> Result> +where + LT: ArrowNumericType, + RT: ArrowNumericType, + RT::Native: One + Zero, + F: Fn(LT::Native, RT::Native) -> Option, +{ + if left.len() != right.len() { + return Err(ArrowError::ComputeError( + "Cannot perform math operation on arrays of different length".to_string(), + )); + } + + let left_iter = ArrayIter::new(left); + let right_iter = ArrayIter::new(right); + + let values: Result::Native>>> = left_iter + .into_iter() + .zip(right_iter.into_iter()) + .map(|(l, r)| { + if let (Some(l), Some(r)) = (l, r) { + let result = op(l, r); + if let Some(r) = result { + Ok(Some(r)) + } else { + if r.is_zero() { + Err(ArrowError::ComputeError("DivideByZero".to_string())) + } else { + // Overflow + Err(ArrowError::ComputeError("Overflow happened".to_string())) + } + } + } else { + Ok(None) + } + }) + .collect(); + + let values = values?; + + Ok(PrimitiveArray::::from_iter(values)) +} + /// Helper function for operations where a valid `0` on the right array should /// result in an [ArrowError::DivideByZero], namely the division and modulo operations /// @@ -1130,7 +1178,7 @@ where #[cfg(feature = "simd")] return simd_checked_divide_op(&left, &right, simd_checked_divide::, |a, b| a / b); #[cfg(not(feature = "simd"))] - return math_checked_op(left, right, |a, b| a.checked_div_if_applied(b)); + return math_checked_divide(left, right, |a, b| a.checked_div_if_applied(b)); } /// Perform `left / right` operation on two arrays. If either left or right value is null @@ -2027,7 +2075,7 @@ mod tests { } #[test] - #[should_panic(expected = "Overflow happened")] + #[should_panic(expected = "DivideByZero")] fn test_primitive_array_divide_by_zero() { let a = Int32Array::from(vec![15]); let b = Int32Array::from(vec![0]);