From e77d4e1e6ca31e3f8d7c51a32dfdf1146a6768a7 Mon Sep 17 00:00:00 2001
From: yjh <465402634@qq.com>
Date: Sat, 9 Oct 2021 11:54:43 +0800
Subject: [PATCH] Added `checked_negate` and `wrapping_negate` for
 `PrimitiveArray` (#506)

---
 src/compute/arithmetics/mod.rs | 43 ++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/compute/arithmetics/mod.rs b/src/compute/arithmetics/mod.rs
index e53fd96a191..e024779ca76 100644
--- a/src/compute/arithmetics/mod.rs
+++ b/src/compute/arithmetics/mod.rs
@@ -61,14 +61,14 @@ pub mod time;
 
 use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
 
-use num_traits::{NumCast, Zero};
+use num_traits::{CheckedNeg, NumCast, WrappingNeg, Zero};
 
 use crate::datatypes::{DataType, IntervalUnit, TimeUnit};
 use crate::error::{ArrowError, Result};
 use crate::types::NativeType;
 use crate::{array::*, bitmap::Bitmap};
 
-use super::arity::unary;
+use super::arity::{unary, unary_checked};
 
 // Macro to evaluate match branch in arithmetic function.
 // The macro is used to downcast both arrays to a primitive_array_type. If there
@@ -302,6 +302,45 @@ where
     unary(array, |a| -a, array.data_type().clone())
 }
 
+/// Checked negates values from array.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::checked_negate;
+/// use arrow2::array::{Array, PrimitiveArray};
+///
+/// let a = PrimitiveArray::from([None, Some(6), Some(i8::MIN), Some(7)]);
+/// let result = checked_negate(&a);
+/// let expected = PrimitiveArray::from([None, Some(-6), None, Some(-7)]);
+/// assert_eq!(result, expected);
+/// assert!(!result.is_valid(2))
+/// ```
+pub fn checked_negate<T>(array: &PrimitiveArray<T>) -> PrimitiveArray<T>
+where
+    T: NativeType + CheckedNeg,
+{
+    unary_checked(array, |a| a.checked_neg(), array.data_type().clone())
+}
+
+/// Wrapping negates values from array.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::wrapping_negate;
+/// use arrow2::array::{Array, PrimitiveArray};
+///
+/// let a = PrimitiveArray::from([None, Some(6), Some(i8::MIN), Some(7)]);
+/// let result = wrapping_negate(&a);
+/// let expected = PrimitiveArray::from([None, Some(-6), Some(i8::MIN), Some(-7)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_negate<T>(array: &PrimitiveArray<T>) -> PrimitiveArray<T>
+where
+    T: NativeType + WrappingNeg,
+{
+    unary(array, |a| a.wrapping_neg(), array.data_type().clone())
+}
+
 /// Defines basic addition operation for primitive arrays
 pub trait ArrayAdd<Rhs> {
     type Output;