From 3bf6eb98ceb3962e1d9419da6dc93e609f7893e6 Mon Sep 17 00:00:00 2001 From: aksharau Date: Mon, 19 Sep 2022 11:18:23 +0530 Subject: [PATCH] Fix: Issue 2721 : binary function should not panic but return error when array lengths are unequal (#2750) --- arrow/src/compute/kernels/arithmetic.rs | 14 +++------- arrow/src/compute/kernels/arity.rs | 36 +++++++++++++++---------- arrow/src/compute/kernels/bitwise.rs | 9 ++----- 3 files changed, 28 insertions(+), 31 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index 7b91a261c7e1..b1a62ccfd6af 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -69,13 +69,7 @@ where RT: ArrowNumericType, F: Fn(LT::Native, RT::Native) -> LT::Native, { - if left.len() != right.len() { - return Err(ArrowError::ComputeError( - "Cannot perform math operation on arrays of different length".to_string(), - )); - } - - Ok(binary(left, right, op)) + binary(left, right, op) } /// Helper function for operations where a valid `0` on the right array should @@ -1128,13 +1122,13 @@ where T: ArrowNumericType, T::Native: ArrowNativeTypeOp + Zero + One, { - Ok(binary_opt(left, right, |a, b| { + binary_opt(left, right, |a, b| { if b.is_zero() { None } else { Some(a.div_wrapping(b)) } - })) + }) } /// Perform `left / right` operation on two arrays. If either left or right value is null @@ -1670,7 +1664,7 @@ mod tests { let b = Int32Array::from(vec![6, 7, 8]); let e = add(&a, &b).expect_err("should have failed due to different lengths"); assert_eq!( - "ComputeError(\"Cannot perform math operation on arrays of different length\")", + "ComputeError(\"Cannot perform binary operation on arrays of different length\")", format!("{:?}", e) ); } diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs index 12cf9721f976..2347502f96e7 100644 --- a/arrow/src/compute/kernels/arity.rs +++ b/arrow/src/compute/kernels/arity.rs @@ -235,25 +235,29 @@ where /// especially when the operation can be vectorised, however, requires `op` to be infallible /// for all possible values of its inputs /// -/// # Panic +/// # Error /// -/// Panics if the arrays have different lengths +/// This function gives error if the arrays have different lengths pub fn binary( a: &PrimitiveArray, b: &PrimitiveArray, op: F, -) -> PrimitiveArray +) -> Result> where A: ArrowPrimitiveType, B: ArrowPrimitiveType, O: ArrowPrimitiveType, F: Fn(A::Native, B::Native) -> O::Native, { - assert_eq!(a.len(), b.len()); + if a.len() != b.len() { + return Err(ArrowError::ComputeError( + "Cannot perform binary operation on arrays of different length".to_string(), + )); + } let len = a.len(); if a.is_empty() { - return PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)); + return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE))); } let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap(); @@ -270,7 +274,7 @@ where // `values` is an iterator with a known size from a PrimitiveArray let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; - unsafe { build_primitive_array(len, buffer, null_count, null_buffer) } + Ok(unsafe { build_primitive_array(len, buffer, null_count, null_buffer) }) } /// Applies the provided fallible binary operation across `a` and `b`, returning any error, @@ -344,32 +348,36 @@ where /// /// The function is only evaluated for non-null indices /// -/// # Panic +/// # Error /// -/// Panics if the arrays have different lengths +/// This function gives error if the arrays have different lengths pub(crate) fn binary_opt( a: &PrimitiveArray, b: &PrimitiveArray, op: F, -) -> PrimitiveArray +) -> Result> where A: ArrowPrimitiveType, B: ArrowPrimitiveType, O: ArrowPrimitiveType, F: Fn(A::Native, B::Native) -> Option, { - assert_eq!(a.len(), b.len()); + if a.len() != b.len() { + return Err(ArrowError::ComputeError( + "Cannot perform binary operation on arrays of different length".to_string(), + )); + } if a.is_empty() { - return PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)); + return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE))); } if a.null_count() == 0 && b.null_count() == 0 { - a.values() + Ok(a.values() .iter() .zip(b.values().iter()) .map(|(a, b)| op(*a, *b)) - .collect() + .collect()) } else { let iter_a = ArrayIter::new(a); let iter_b = ArrayIter::new(b); @@ -386,7 +394,7 @@ where } }); - values.collect() + Ok(values.collect()) } } diff --git a/arrow/src/compute/kernels/bitwise.rs b/arrow/src/compute/kernels/bitwise.rs index 2f3c9e490f4c..0b877b326482 100644 --- a/arrow/src/compute/kernels/bitwise.rs +++ b/arrow/src/compute/kernels/bitwise.rs @@ -18,7 +18,7 @@ use crate::array::PrimitiveArray; use crate::compute::{binary, unary}; use crate::datatypes::ArrowNumericType; -use crate::error::{ArrowError, Result}; +use crate::error::Result; use std::ops::{BitAnd, BitOr, BitXor, Not}; // The helper function for bitwise operation with two array @@ -31,12 +31,7 @@ where T: ArrowNumericType, F: Fn(T::Native, T::Native) -> T::Native, { - if left.len() != right.len() { - return Err(ArrowError::ComputeError( - "Cannot perform bitwise operation on arrays of different length".to_string(), - )); - } - Ok(binary(left, right, op)) + binary(left, right, op) } /// Perform `left & right` operation on two arrays. If either left or right value is null