Fix: Issue 2721 : binary function should not panic but return error when array lengths are unequal (#2750)

Author: aksharau

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)
         );
     }

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, B, F, O>(
     a: &PrimitiveArray<A>,
     b: &PrimitiveArray<B>,
     op: F,
-) -> PrimitiveArray<O>
+) -> Result<PrimitiveArray<O>>
 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, B, F, O>(
     a: &PrimitiveArray<A>,
     b: &PrimitiveArray<B>,
     op: F,
-) -> PrimitiveArray<O>
+) -> Result<PrimitiveArray<O>>
 where
     A: ArrowPrimitiveType,
     B: ArrowPrimitiveType,
     O: ArrowPrimitiveType,
     F: Fn(A::Native, B::Native) -> Option<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(),
+        ));
+    }
 
     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())
     }
 }
 

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