Support overflow-checking variant of negate kernel (#2893)

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

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
             }