[arrow-rs] branch master updated: Add subtract_scalar kernel (#1152)

[alamb]

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

alamb 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 a800a88  Add subtract_scalar kernel (#1152)
a800a88 is described below

commit a800a88153cbe62904c62b5dc402cdf534dc7bc1
Author: Liang-Chi Hsieh <vii...@gmail.com>
AuthorDate: Wed Jan 12 04:22:42 2022 -0800

    Add subtract_scalar kernel (#1152)
    
    * Add subtract_scalar
    
    * Rebase
---
 arrow/src/compute/kernels/arithmetic.rs | 42 +++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/arrow/src/compute/kernels/arithmetic.rs 
b/arrow/src/compute/kernels/arithmetic.rs
index 4195ea5..81adb0b 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -1061,6 +1061,29 @@ where
     return math_op(left, right, |a, b| a - b);
 }
 
+/// Subtract every value in an array by a scalar. If any value in the array is 
null then the
+/// result is also null.
+pub fn subtract_scalar<T>(
+    array: &PrimitiveArray<T>,
+    scalar: T::Native,
+) -> Result<PrimitiveArray<T>>
+where
+    T: datatypes::ArrowNumericType,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + Zero,
+{
+    #[cfg(feature = "simd")]
+    {
+        let scalar_vector = T::init(scalar);
+        return simd_unary_math_op(array, |x| x - scalar_vector, |x| x - 
scalar);
+    }
+    #[cfg(not(feature = "simd"))]
+    return Ok(unary(array, |value| value - scalar));
+}
+
 /// Perform `-` operation on an array. If value is null then the result is 
also null.
 pub fn negate<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
 where
@@ -1312,6 +1335,25 @@ mod tests {
     }
 
     #[test]
+    fn test_primitive_array_subtract_scalar() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = 3;
+        let c = subtract_scalar(&a, b).unwrap();
+        let expected = Int32Array::from(vec![12, 11, 6, 5, -2]);
+        assert_eq!(c, expected);
+    }
+
+    #[test]
+    fn test_primitive_array_subtract_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);
+        let actual = subtract_scalar(a, 3).unwrap();
+        let expected = Int32Array::from(vec![None, Some(6), Some(5), None]);
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
     fn test_primitive_array_multiply() {
         let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
         let b = Int32Array::from(vec![6, 7, 8, 9, 8]);