Add modulus ops into ArrowNativeTypeOp

arrow/src/compute/kernels/arithmetic.rs

@@ -22,7 +22,7 @@
 //! `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::{Div, Neg, Rem};
+use std::ops::{Div, Neg};
 
 use num::{One, Zero};
 
@@ -168,7 +168,7 @@ fn simd_checked_modulus<T: ArrowNumericType>(
     right: T::Simd,
 ) -> Result<T::Simd>
 where
-    T::Native: One + Zero,
+    T::Native: ArrowNativeTypeOp + One,
 {
     let zero = T::init(T::Native::zero());
     let one = T::init(T::Native::one());
@@ -291,7 +291,7 @@ fn simd_checked_divide_op<T, SI, SC>(
 ) -> Result<PrimitiveArray<T>>
 where
     T: ArrowNumericType,
-    T::Native: One + Zero,
+    T::Native: ArrowNativeTypeOp,
     SI: Fn(Option<u64>, T::Simd, T::Simd) -> Result<T::Simd>,
     SC: Fn(T::Native, T::Native) -> T::Native,
 {
@@ -1075,20 +1075,14 @@ pub fn modulus<T>(
 ) -> Result<PrimitiveArray<T>>
 where
     T: ArrowNumericType,
-    T::Native: Rem<Output = T::Native> + Zero + One,
+    T::Native: ArrowNativeTypeOp + One,
 {
     #[cfg(feature = "simd")]
     return simd_checked_divide_op(&left, &right, simd_checked_modulus::<T>, |a, b| {
-        a % b
+        a.mod_wrapping(b)
     });
     #[cfg(not(feature = "simd"))]
-    return try_binary(left, right, |a, b| {
-        if b.is_zero() {
-            Err(ArrowError::DivideByZero)
-        } else {
-            Ok(a % b)
-        }
-    });
+    return try_binary(left, right, |a, b| a.mod_checked_divide_by_zero(b));
 }
 
 /// Perform `left / right` operation on two arrays. If either left or right value is null
@@ -1191,13 +1185,13 @@ pub fn modulus_scalar<T>(
 ) -> Result<PrimitiveArray<T>>
 where
     T: ArrowNumericType,
-    T::Native: Rem<Output = T::Native> + Zero,
+    T::Native: ArrowNativeTypeOp,
 {
     if modulo.is_zero() {
         return Err(ArrowError::DivideByZero);
     }
 
-    Ok(unary(array, |a| a % modulo))
+    Ok(unary(array, |a| a.mod_wrapping(modulo)))
 }
 
 /// Divide every value in an array by a scalar. If any value in the array is null then the
@@ -1769,7 +1763,7 @@ mod tests {
     }
 
     #[test]
-    fn test_primitive_array_modulus() {
+    fn test_int_array_modulus() {
         let a = Int32Array::from(vec![15, 15, 8, 1, 9]);
         let b = Int32Array::from(vec![5, 6, 8, 9, 1]);
         let c = modulus(&a, &b).unwrap();
@@ -1780,6 +1774,24 @@ mod tests {
         assert_eq!(0, c.value(4));
     }
 
+    #[test]
+    #[should_panic(
+        expected = "called `Result::unwrap()` on an `Err` value: DivideByZero"
+    )]
+    fn test_int_array_modulus_divide_by_zero() {
+        let a = Int32Array::from(vec![1]);
+        let b = Int32Array::from(vec![0]);
+        modulus(&a, &b).unwrap();
+    }
+
+    #[test]
+    fn test_int_array_modulus_overflow_wrapping() {
+        let a = Int32Array::from(vec![i32::MIN]);
+        let b = Int32Array::from(vec![-1]);
+        let result = modulus(&a, &b).unwrap();
+        assert_eq!(0, result.value(0))
+    }
+
     #[test]
     fn test_primitive_array_divide_scalar() {
         let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
@@ -1842,7 +1854,7 @@ mod tests {
     }
 
     #[test]
-    fn test_primitive_array_modulus_scalar() {
+    fn test_int_array_modulus_scalar() {
         let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
         let b = 3;
         let c = modulus_scalar(&a, b).unwrap();
@@ -1851,7 +1863,7 @@ mod tests {
     }
 
     #[test]
-    fn test_primitive_array_modulus_scalar_sliced() {
+    fn test_int_array_modulus_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
         let a = as_primitive_array(&a);
@@ -1860,6 +1872,22 @@ mod tests {
         assert_eq!(actual, expected);
     }
 
+    #[test]
+    #[should_panic(
+        expected = "called `Result::unwrap()` on an `Err` value: DivideByZero"
+    )]
+    fn test_int_array_modulus_scalar_divide_by_zero() {
+        let a = Int32Array::from(vec![1]);
+        modulus_scalar(&a, 0).unwrap();
+    }
+
+    #[test]
+    fn test_int_array_modulus_scalar_overflow_wrapping() {
+        let a = Int32Array::from(vec![i32::MIN]);
+        let result = modulus_scalar(&a, -1).unwrap();
+        assert_eq!(0, result.value(0))
+    }
+
     #[test]
     fn test_primitive_array_divide_sliced() {
         let a = Int32Array::from(vec![0, 0, 0, 15, 15, 8, 1, 9, 0]);

arrow/src/datatypes/native.rs

@@ -47,7 +47,7 @@ pub(crate) mod native_op {
     use super::ArrowNativeType;
     use crate::error::{ArrowError, Result};
     use num::Zero;
-    use std::ops::{Add, Div, Mul, Sub};
+    use std::ops::{Add, Div, Mul, Rem, Sub};
 
     /// Trait for ArrowNativeType to provide overflow-checking and non-overflow-checking
     /// variants for arithmetic operations. For floating point types, this provides some
@@ -65,6 +65,7 @@ pub(crate) mod native_op {
         + Sub<Output = Self>
         + Mul<Output = Self>
         + Div<Output = Self>
+        + Rem<Output = Self>
         + Zero
     {
         fn add_checked(self, rhs: Self) -> Result<Self> {
@@ -102,6 +103,28 @@ pub(crate) mod native_op {
         fn div_wrapping(self, rhs: Self) -> Self {
             self / rhs
         }
+
+        /// Check `DivideByZero` error and `Overflow` error.
+        fn mod_fully_checked(self, rhs: Self) -> Result<Self> {
+            if rhs.is_zero() {
+                Err(ArrowError::DivideByZero)
+            } else {
+                Ok(self.mod_wrapping(rhs))
+            }
+        }
+
+        /// Only check `DivideByZero` error
+        fn mod_checked_divide_by_zero(self, rhs: Self) -> Result<Self> {
+            if rhs.is_zero() {
+                Err(ArrowError::DivideByZero)
+            } else {
+                Ok(self.mod_wrapping(rhs))
+            }
+        }
+
+        fn mod_wrapping(self, rhs: Self) -> Self {
+            self % rhs
+        }
     }
 }
 
@@ -163,6 +186,23 @@ macro_rules! native_type_op {
             fn div_wrapping(self, rhs: Self) -> Self {
                 self.wrapping_div(rhs)
             }
+
+            fn mod_fully_checked(self, rhs: Self) -> Result<Self> {
+                if rhs.is_zero() {
+                    Err(ArrowError::DivideByZero)
+                } else {
+                    self.checked_rem(rhs).ok_or_else(|| {
+                        ArrowError::ComputeError(format!(
+                            "Overflow happened on: {:?} % {:?}",
+                            self, rhs
+                        ))
+                    })
+                }
+            }
+
+            fn mod_wrapping(self, rhs: Self) -> Self {
+                self.wrapping_rem(rhs)
+            }
         }
     };
 }