Add Wrapping version ops for primitiveArray (#496)

src/compute/arithmetics/basic/add.rs

@@ -1,9 +1,10 @@
 //! Definition of basic add operations with primitive arrays
 use std::ops::Add;
 
-use num_traits::{ops::overflowing::OverflowingAdd, CheckedAdd, SaturatingAdd, Zero};
+use num_traits::{ops::overflowing::OverflowingAdd, CheckedAdd, SaturatingAdd, WrappingAdd, Zero};
 
 use crate::compute::arithmetics::basic::check_same_type;
+use crate::compute::arithmetics::ArrayWrappingAdd;
 use crate::{
     array::{Array, PrimitiveArray},
     bitmap::Bitmap,
@@ -42,6 +43,34 @@ where
     binary(lhs, rhs, lhs.data_type().clone(), |a, b| a + b)
 }
 
+/// Wrapping addition of two [`PrimitiveArray`]s.
+/// It wraps around at the boundary of the type if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_add;
+/// use arrow2::array::PrimitiveArray;
+///
+/// let a = PrimitiveArray::from([Some(-100i8), Some(100i8), Some(100i8)]);
+/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]);
+/// let result = wrapping_add(&a, &b).unwrap();
+/// let expected = PrimitiveArray::from([Some(-100i8), Some(-56i8), Some(100i8)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_add<T>(
+    lhs: &PrimitiveArray<T>,
+    rhs: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: NativeType + WrappingAdd<Output = T>,
+{
+    check_same_type(lhs, rhs)?;
+
+    let op = move |a: T, b: T| a.wrapping_add(&b);
+
+    binary(lhs, rhs, lhs.data_type().clone(), op)
+}
+
 /// Checked addition of two primitive arrays. If the result from the sum
 /// overflows, the validity for that index is changed to None
 ///
@@ -138,6 +167,17 @@ where
     }
 }
 
+impl<T> ArrayWrappingAdd<PrimitiveArray<T>> for PrimitiveArray<T>
+where
+    T: NativeType + WrappingAdd<Output = T> + NotI128,
+{
+    type Output = Self;
+
+    fn wrapping_add(&self, rhs: &PrimitiveArray<T>) -> Result<Self::Output> {
+        wrapping_add(self, rhs)
+    }
+}
+
 // Implementation of ArrayCheckedAdd trait for PrimitiveArrays
 impl<T> ArrayCheckedAdd<PrimitiveArray<T>> for PrimitiveArray<T>
 where
@@ -195,6 +235,26 @@ where
     unary(lhs, |a| a + rhs, lhs.data_type().clone())
 }
 
+/// Wrapping addition of a scalar T to a [`PrimitiveArray`] of type T.
+/// It do nothing if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_add_scalar;
+/// use arrow2::array::Int8Array;
+///
+/// let a = Int8Array::from(&[None, Some(100)]);
+/// let result = wrapping_add_scalar(&a, &100i8);
+/// let expected = Int8Array::from(&[None, Some(-56)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
+where
+    T: NativeType + WrappingAdd<Output = T>,
+{
+    unary(lhs, |a| a.wrapping_add(rhs), lhs.data_type().clone())
+}
+
 /// Checked addition of a scalar T to a primitive array of type T. If the
 /// result from the sum overflows then the validity index for that value is
 /// changed to None

src/compute/arithmetics/basic/mul.rs

@@ -1,9 +1,10 @@
 //! Definition of basic mul operations with primitive arrays
 use std::ops::Mul;
 
-use num_traits::{ops::overflowing::OverflowingMul, CheckedMul, SaturatingMul, Zero};
+use num_traits::{ops::overflowing::OverflowingMul, CheckedMul, SaturatingMul, WrappingMul, Zero};
 
 use crate::compute::arithmetics::basic::check_same_type;
+use crate::compute::arithmetics::ArrayWrappingMul;
 use crate::{
     array::{Array, PrimitiveArray},
     bitmap::Bitmap,
@@ -42,6 +43,34 @@ where
     binary(lhs, rhs, lhs.data_type().clone(), |a, b| a * b)
 }
 
+/// Wrapping multiplication of two [`PrimitiveArray`]s.
+///  It wraps around at the boundary of the type if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_mul;
+/// use arrow2::array::PrimitiveArray;
+///
+/// let a = PrimitiveArray::from([Some(100i8), Some(0x10i8), Some(100i8)]);
+/// let b = PrimitiveArray::from([Some(0i8), Some(0x10i8), Some(0i8)]);
+/// let result = wrapping_mul(&a, &b).unwrap();
+/// let expected = PrimitiveArray::from([Some(0), Some(0), Some(0)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_mul<T>(
+    lhs: &PrimitiveArray<T>,
+    rhs: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: NativeType + WrappingMul<Output = T>,
+{
+    check_same_type(lhs, rhs)?;
+
+    let op = move |a: T, b: T| a.wrapping_mul(&b);
+
+    binary(lhs, rhs, lhs.data_type().clone(), op)
+}
+
 /// Checked multiplication of two primitive arrays. If the result from the
 /// multiplications overflows, the validity for that index is changed
 /// returned.
@@ -139,6 +168,17 @@ where
     }
 }
 
+impl<T> ArrayWrappingMul<PrimitiveArray<T>> for PrimitiveArray<T>
+where
+    T: NativeType + WrappingMul<Output = T> + NotI128,
+{
+    type Output = Self;
+
+    fn wrapping_mul(&self, rhs: &PrimitiveArray<T>) -> Result<Self::Output> {
+        wrapping_mul(self, rhs)
+    }
+}
+
 // Implementation of ArrayCheckedMul trait for PrimitiveArrays
 impl<T> ArrayCheckedMul<PrimitiveArray<T>> for PrimitiveArray<T>
 where
@@ -195,6 +235,26 @@ where
     unary(lhs, |a| a * rhs, lhs.data_type().clone())
 }
 
+/// Wrapping multiplication of a scalar T to a [`PrimitiveArray`] of type T.
+/// It do nothing if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_mul_scalar;
+/// use arrow2::array::Int8Array;
+///
+/// let a = Int8Array::from(&[None, Some(0x10)]);
+/// let result = wrapping_mul_scalar(&a, &0x10);
+/// let expected = Int8Array::from(&[None, Some(0)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_mul_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
+where
+    T: NativeType + WrappingMul<Output = T>,
+{
+    unary(lhs, |a| a.wrapping_mul(rhs), lhs.data_type().clone())
+}
+
 /// Checked multiplication of a scalar T to a primitive array of type T. If the
 /// result from the multiplication overflows, then the validity for that index is
 /// changed to None

src/compute/arithmetics/basic/sub.rs

@@ -1,9 +1,10 @@
 //! Definition of basic sub operations with primitive arrays
 use std::ops::Sub;
 
-use num_traits::{ops::overflowing::OverflowingSub, CheckedSub, SaturatingSub, Zero};
+use num_traits::{ops::overflowing::OverflowingSub, CheckedSub, SaturatingSub, WrappingSub, Zero};
 
 use crate::compute::arithmetics::basic::check_same_type;
+use crate::compute::arithmetics::ArrayWrappingSub;
 use crate::{
     array::{Array, PrimitiveArray},
     bitmap::Bitmap,
@@ -42,6 +43,34 @@ where
     binary(lhs, rhs, lhs.data_type().clone(), |a, b| a - b)
 }
 
+/// Wrapping subtraction of two [`PrimitiveArray`]s.
+///  It wraps around at the boundary of the type if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_sub;
+/// use arrow2::array::PrimitiveArray;
+///
+/// let a = PrimitiveArray::from([Some(-100i8), Some(-100i8), Some(100i8)]);
+/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]);
+/// let result = wrapping_sub(&a, &b).unwrap();
+/// let expected = PrimitiveArray::from([Some(-100i8), Some(56i8), Some(100i8)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_sub<T>(
+    lhs: &PrimitiveArray<T>,
+    rhs: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: NativeType + WrappingSub<Output = T>,
+{
+    check_same_type(lhs, rhs)?;
+
+    let op = move |a: T, b: T| a.wrapping_sub(&b);
+
+    binary(lhs, rhs, lhs.data_type().clone(), op)
+}
+
 /// Checked subtraction of two primitive arrays. If the result from the
 /// subtraction overflow, the validity for that index is changed
 ///
@@ -138,6 +167,17 @@ where
     }
 }
 
+impl<T> ArrayWrappingSub<PrimitiveArray<T>> for PrimitiveArray<T>
+where
+    T: NativeType + WrappingSub<Output = T> + NotI128,
+{
+    type Output = Self;
+
+    fn wrapping_sub(&self, rhs: &PrimitiveArray<T>) -> Result<Self::Output> {
+        wrapping_sub(self, rhs)
+    }
+}
+
 // Implementation of ArrayCheckedSub trait for PrimitiveArrays
 impl<T> ArrayCheckedSub<PrimitiveArray<T>> for PrimitiveArray<T>
 where
@@ -195,6 +235,26 @@ where
     unary(lhs, |a| a - rhs, lhs.data_type().clone())
 }
 
+/// Wrapping subtraction of a scalar T to a [`PrimitiveArray`] of type T.
+/// It do nothing if the result overflows.
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::arithmetics::basic::wrapping_sub_scalar;
+/// use arrow2::array::Int8Array;
+///
+/// let a = Int8Array::from(&[None, Some(-100)]);
+/// let result = wrapping_sub_scalar(&a, &100i8);
+/// let expected = Int8Array::from(&[None, Some(56)]);
+/// assert_eq!(result, expected);
+/// ```
+pub fn wrapping_sub_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
+where
+    T: NativeType + WrappingSub<Output = T>,
+{
+    unary(lhs, |a| a.wrapping_sub(rhs), lhs.data_type().clone())
+}
+
 /// Checked subtraction of a scalar T to a primitive array of type T. If the
 /// result from the subtraction overflows, then the validity for that index
 /// is changed to None

src/compute/arithmetics/mod.rs

@@ -309,6 +309,13 @@ pub trait ArrayAdd<Rhs> {
     fn add(&self, rhs: &Rhs) -> Result<Self::Output>;
 }
 
+/// Defines wrapping addition operation for primitive arrays
+pub trait ArrayWrappingAdd<Rhs> {
+    type Output;
+
+    fn wrapping_add(&self, rhs: &Rhs) -> Result<Self::Output>;
+}
+
 /// Defines checked addition operation for primitive arrays
 pub trait ArrayCheckedAdd<Rhs> {
     type Output;
@@ -337,6 +344,13 @@ pub trait ArraySub<Rhs> {
     fn sub(&self, rhs: &Rhs) -> Result<Self::Output>;
 }
 
+/// Defines wrapping subtraction operation for primitive arrays
+pub trait ArrayWrappingSub<Rhs> {
+    type Output;
+
+    fn wrapping_sub(&self, rhs: &Rhs) -> Result<Self::Output>;
+}
+
 /// Defines checked subtraction operation for primitive arrays
 pub trait ArrayCheckedSub<Rhs> {
     type Output;
@@ -365,6 +379,13 @@ pub trait ArrayMul<Rhs> {
     fn mul(&self, rhs: &Rhs) -> Result<Self::Output>;
 }
 
+/// Defines wrapping multiplication operation for primitive arrays
+pub trait ArrayWrappingMul<Rhs> {
+    type Output;
+
+    fn wrapping_mul(&self, rhs: &Rhs) -> Result<Self::Output>;
+}
+
 /// Defines checked multiplication operation for primitive arrays
 pub trait ArrayCheckedMul<Rhs> {
     type Output;