[arrow-rs] branch master updated: Support overflow-checking variant of negate kernel (#2893)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new e118ae2d3 Support overflow-checking variant of negate kernel (#2893)
e118ae2d3 is described below

commit e118ae2d3b3ed9945bd0721a10d001a829d41854
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Tue Oct 18 14:58:55 2022 -0700

    Support overflow-checking variant of negate kernel (#2893)
---
 arrow/src/compute/kernels/arithmetic.rs | 32 ++++++++++++++++++++++++++++----
 arrow/src/datatypes/native.rs           | 22 ++++++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
index a73ee7eee..e7fcb50cc 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -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;
@@ -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.
@@ -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();
diff --git a/arrow/src/datatypes/native.rs b/arrow/src/datatypes/native.rs
index 444ba39e0..2643025f1 100644
--- a/arrow/src/datatypes/native.rs
+++ b/arrow/src/datatypes/native.rs
@@ -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;
@@ -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
             }
@@ -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
             }