From c9a0c00557555181472b4b2844102ccd6f2f20fe Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 10 Jan 2022 15:51:13 -0800 Subject: [PATCH 1/2] 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..2b11e4e3e16 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -990,6 +990,69 @@ where Ok(PrimitiveArray::::from(data)) } +/// SIMD vectorized version of adding a scalar to an array. +#[cfg(feature = "simd")] +fn simd_add_scalar( + array: &PrimitiveArray, + scalar: T::Native, +) -> Result> +where + T: ArrowNumericType, + T::Native: Add + + Sub + + Mul + + Div + + Rem + + Zero + + One, +{ + let lanes = T::lanes(); + let buffer_size = array.len() * std::mem::size_of::(); + 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::::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( @@ -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( + array: &PrimitiveArray, + scalar: T::Native, +) -> Result> +where + T: datatypes::ArrowNumericType, + T::Native: Add + + Sub + + Mul + + Div + + Rem + + 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( @@ -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]); From 9aec3762294b3c7eb63234eea22dbb88f3f02219 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 11 Jan 2022 09:09:55 -0800 Subject: [PATCH 2/2] move simd_float_unary_math_op to simd_unary_math_op --- arrow/src/compute/kernels/arithmetic.rs | 78 +++---------------------- 1 file changed, 7 insertions(+), 71 deletions(-) diff --git a/arrow/src/compute/kernels/arithmetic.rs b/arrow/src/compute/kernels/arithmetic.rs index 2b11e4e3e16..46bd42faf4f 100644 --- a/arrow/src/compute/kernels/arithmetic.rs +++ b/arrow/src/compute/kernels/arithmetic.rs @@ -97,13 +97,13 @@ where } #[cfg(feature = "simd")] -fn simd_float_unary_math_op( +fn simd_unary_math_op( array: &PrimitiveArray, simd_op: SIMD_OP, scalar_op: SCALAR_OP, ) -> Result> where - T: datatypes::ArrowFloatNumericType, + T: ArrowNumericType, SIMD_OP: Fn(T::Simd) -> T::Simd, SCALAR_OP: Fn(T::Native) -> T::Native, { @@ -990,69 +990,6 @@ where Ok(PrimitiveArray::::from(data)) } -/// SIMD vectorized version of adding a scalar to an array. -#[cfg(feature = "simd")] -fn simd_add_scalar( - array: &PrimitiveArray, - scalar: T::Native, -) -> Result> -where - T: ArrowNumericType, - T::Native: Add - + Sub - + Mul - + Div - + Rem - + Zero - + One, -{ - let lanes = T::lanes(); - let buffer_size = array.len() * std::mem::size_of::(); - 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::::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( @@ -1090,7 +1027,10 @@ where + One, { #[cfg(feature = "simd")] - return simd_add_scalar(&array, scalar); + { + 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)); } @@ -1139,11 +1079,7 @@ where #[cfg(feature = "simd")] { let raise_vector = T::init(raise); - return simd_float_unary_math_op( - array, - |x| T::pow(x, raise_vector), - |x| x.pow(raise), - ); + return simd_unary_math_op(array, |x| T::pow(x, raise_vector), |x| x.pow(raise)); } #[cfg(not(feature = "simd"))] return Ok(unary(array, |x| x.pow(raise)));