From c33e6801f911c4fd0a7249ada7e49a18c9e1b9d6 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Mon, 10 Jan 2022 15:51:13 -0800
Subject: [PATCH] Add add_scalar

---
 arrow/src/compute/kernels/arithmetic.rs | 104 ++++++++++++++++++++++++
 1 file changed, 104 insertions(+)

diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs
index 09d4b9fd6cd..af52eed8ca6 100644
--- a/arrow/src/compute/kernels/arithmetic.rs
+++ b/arrow/src/compute/kernels/arithmetic.rs
@@ -990,6 +990,69 @@ where
     Ok(PrimitiveArray::<T>::from(data))
 }
 
+/// SIMD vectorized version of adding a scalar to an array.
+#[cfg(feature = "simd")]
+fn simd_add_scalar<T>(
+    array: &PrimitiveArray<T>,
+    scalar: T::Native,
+) -> Result<PrimitiveArray<T>>
+where
+    T: 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,
+{
+    let lanes = T::lanes();
+    let buffer_size = array.len() * std::mem::size_of::<T::Native>();
+    let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
+
+    // safety: result is newly created above, always written as a T below
+    let mut result_chunks = unsafe { result.typed_data_mut().chunks_exact_mut(lanes) };
+    let mut array_chunks = array.values().chunks_exact(lanes);
+
+    let simd_right = T::init(scalar);
+
+    result_chunks
+        .borrow_mut()
+        .zip(array_chunks.borrow_mut())
+        .for_each(|(result_slice, array_slice)| {
+            let simd_left = T::load(array_slice);
+
+            let simd_result = T::bin_op(simd_left, simd_right, |a, b| a + b);
+            T::write(simd_result, result_slice);
+        });
+
+    let result_remainder = result_chunks.into_remainder();
+    let array_remainder = array_chunks.remainder();
+
+    result_remainder
+        .iter_mut()
+        .zip(array_remainder.iter())
+        .for_each(|(scalar_result, scalar_array)| {
+            *scalar_result = *scalar_array + scalar;
+        });
+
+    let data = unsafe {
+        ArrayData::new_unchecked(
+            T::DATA_TYPE,
+            array.len(),
+            None,
+            array
+                .data_ref()
+                .null_buffer()
+                .map(|b| b.bit_slice(array.offset(), array.len())),
+            0,
+            vec![result.into()],
+            vec![],
+        )
+    };
+    Ok(PrimitiveArray::<T>::from(data))
+}
+
 /// Perform `left + right` operation on two arrays. If either left or right value is null
 /// then the result is also null.
 pub fn add<T>(
@@ -1010,6 +1073,28 @@ where
     return math_op(left, right, |a, b| a + b);
 }
 
+/// Add every value in an array by a scalar. If any value in the array is null then the
+/// result is also null.
+pub fn add_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")]
+    return simd_add_scalar(&array, 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.
 pub fn subtract<T>(
@@ -1228,6 +1313,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_primitive_array_add_scalar() {
+        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
+        let b = 3;
+        let c = add_scalar(&a, b).unwrap();
+        let expected = Int32Array::from(vec![18, 17, 12, 11, 4]);
+        assert_eq!(c, expected);
+    }
+
+    #[test]
+    fn test_primitive_array_add_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 = add_scalar(a, 3).unwrap();
+        let expected = Int32Array::from(vec![None, Some(12), Some(11), None]);
+        assert_eq!(actual, expected);
+    }
+
     #[test]
     fn test_primitive_array_subtract() {
         let a = Int32Array::from(vec![1, 2, 3, 4, 5]);