[arrow-rs] branch master updated: Add multiply_scalar (#1159)

[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 d03cd47  Add multiply_scalar (#1159)
d03cd47 is described below

commit d03cd4702f015e07fbf389a4e33ddd7a5146efdf
Author: Liang-Chi Hsieh <vii...@gmail.com>
AuthorDate: Wed Jan 12 04:21:37 2022 -0800

    Add multiply_scalar (#1159)
---
 arrow/src/compute/kernels/arithmetic.rs | 44 +++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/arrow/src/compute/kernels/arithmetic.rs 
b/arrow/src/compute/kernels/arithmetic.rs
index 590bb64..4195ea5 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -1112,6 +1112,31 @@ where
     return math_op(left, right, |a, b| a * b);
 }
 
+/// Multiply every value in an array by a scalar. If any value in the array is 
null then the
+/// result is also null.
+pub fn multiply_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>
+        + Rem<Output = T::Native>
+        + Zero
+        + One,
+{
+    #[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 `left % right` operation on two arrays. If either left or right 
value is null
 /// then the result is also null. If any right hand value is zero then the 
result of this
 /// operation will be `Err(ArrowError::DivideByZero)`.
@@ -1299,6 +1324,25 @@ mod tests {
     }
 
     #[test]
+    fn test_primitive_array_multiply_scalar() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = 3;
+        let c = multiply_scalar(&a, b).unwrap();
+        let expected = Int32Array::from(vec![45, 42, 27, 24, 3]);
+        assert_eq!(c, expected);
+    }
+
+    #[test]
+    fn test_primitive_array_multiply_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 = multiply_scalar(a, 3).unwrap();
+        let expected = Int32Array::from(vec![None, Some(27), Some(24), None]);
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
     fn test_primitive_array_divide() {
         let a = Int32Array::from(vec![15, 15, 8, 1, 9]);
         let b = Int32Array::from(vec![5, 6, 8, 9, 1]);