Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support overflow-checking variant of negate kernel #2893

Merged
merged 1 commit into from Oct 18, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
32 changes: 28 additions & 4 deletions arrow/src/compute/kernels/arithmetic.rs
Expand Up @@ -22,8 +22,6 @@
//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.

use std::ops::Neg;

use crate::array::*;
#[cfg(feature = "simd")]
use crate::buffer::MutableBuffer;
Expand Down Expand Up @@ -1116,12 +1114,27 @@ where
}

/// Perform `-` operation on an array. If value is null then the result is also null.
///
/// This doesn't detect overflow. Once overflowing, the result will wrap around.
/// For an overflow-checking variant, use `negate_checked` instead.
pub fn negate<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: Neg<Output = T::Native>,
T::Native: ArrowNativeTypeOp,
{
Ok(unary(array, |x| x.neg_wrapping()))
}

/// Perform `-` operation on an array. If value is null then the result is also null.
///
/// This detects overflow and returns an `Err` for that. For an non-overflow-checking variant,
/// use `negate` instead.
pub fn negate_checked<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: ArrowNativeTypeOp,
{
Ok(unary(array, |x| -x))
try_unary(array, |value| value.neg_checked())
}

/// Raise array with floating point values to the power of a scalar.
Expand Down Expand Up @@ -2567,6 +2580,17 @@ mod tests {
assert_eq!(expected, actual);
}

#[test]
fn test_primitive_array_negate_checked_overflow() {
let a = Int32Array::from(vec![i32::MIN]);
let actual = negate(&a).unwrap();
let expected = Int32Array::from(vec![i32::MIN]);
assert_eq!(expected, actual);

let err = negate_checked(&a);
err.expect_err("negate_checked should detect overflow");
}

#[test]
fn test_arithmetic_kernel_should_not_rely_on_padding() {
let a: UInt8Array = (0..128_u8).into_iter().map(Some).collect();
Expand Down
22 changes: 22 additions & 0 deletions arrow/src/datatypes/native.rs
Expand Up @@ -64,6 +64,10 @@ pub trait ArrowNativeTypeOp: ArrowNativeType {

fn mod_wrapping(self, rhs: Self) -> Self;

fn neg_checked(self) -> Result<Self>;

fn neg_wrapping(self) -> Self;

fn is_zero(self) -> bool;

fn is_eq(self, rhs: Self) -> bool;
Expand Down Expand Up @@ -158,6 +162,16 @@ macro_rules! native_type_op {
self.wrapping_rem(rhs)
}

fn neg_checked(self) -> Result<Self> {
self.checked_neg().ok_or_else(|| {
ArrowError::ComputeError(format!("Overflow happened on: {:?}", self))
})
}

fn neg_wrapping(self) -> Self {
self.wrapping_neg()
}

fn is_zero(self) -> bool {
self == 0
}
Expand Down Expand Up @@ -253,6 +267,14 @@ macro_rules! native_type_float_op {
self % rhs
}

fn neg_checked(self) -> Result<Self> {
Ok(-self)
}

fn neg_wrapping(self) -> Self {
-self
}

fn is_zero(self) -> bool {
self == $zero
}
Expand Down