Skip to content

Commit

Permalink
Fix: Issue 2721 : binary function should not panic but return error w…
Browse files Browse the repository at this point in the history
…hen array lengths are unequal (#2750)
  • Loading branch information
aksharau committed Sep 19, 2022
1 parent 5e83ef9 commit 3bf6eb9
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 31 deletions.
14 changes: 4 additions & 10 deletions arrow/src/compute/kernels/arithmetic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
);
}
Expand Down
36 changes: 22 additions & 14 deletions arrow/src/compute/kernels/arity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -386,7 +394,7 @@ where
}
});

values.collect()
Ok(values.collect())
}
}

Expand Down
9 changes: 2 additions & 7 deletions arrow/src/compute/kernels/bitwise.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 3bf6eb9

Please sign in to comment.