diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs b/rust/arrow/src/compute/kernels/arithmetic.rs index 06d7c90f6a5..067756662cf 100644 --- a/rust/arrow/src/compute/kernels/arithmetic.rs +++ b/rust/arrow/src/compute/kernels/arithmetic.rs @@ -30,34 +30,53 @@ use num::{One, Zero}; use crate::buffer::Buffer; #[cfg(simd)] use crate::buffer::MutableBuffer; -use crate::compute::util::combine_option_bitmap; +use crate::compute::{kernels::arity::unary, util::combine_option_bitmap}; use crate::datatypes; use crate::datatypes::ArrowNumericType; use crate::error::{ArrowError, Result}; use crate::{array::*, util::bit_util}; +use num::traits::Pow; #[cfg(simd)] use std::borrow::BorrowMut; #[cfg(simd)] use std::slice::{ChunksExact, ChunksExactMut}; -/// Helper function to perform math lambda function on values from single array of signed numeric -/// type. If value is null then the output value is also null, so `-null` is `null`. -pub fn signed_unary_math_op( +/// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values. +#[cfg(simd)] +fn simd_signed_unary_math_op( array: &PrimitiveArray, - op: F, + simd_op: SIMD_OP, + scalar_op: SCALAR_OP, ) -> Result> where T: datatypes::ArrowSignedNumericType, - T::Native: Neg, - F: Fn(T::Native) -> T::Native, + SIMD_OP: Fn(T::SignedSimd) -> T::SignedSimd, + SCALAR_OP: Fn(T::Native) -> T::Native, { - let values = array.values().iter().map(|v| op(*v)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size. - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; + 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); + + let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes); + let mut array_chunks = array.values().chunks_exact(lanes); + + result_chunks + .borrow_mut() + .zip(array_chunks.borrow_mut()) + .for_each(|(result_slice, input_slice)| { + let simd_input = T::load_signed(input_slice); + let simd_result = T::signed_unary_op(simd_input, &simd_op); + T::write_signed(simd_result, result_slice); + }); + + let result_remainder = result_chunks.into_remainder(); + let array_remainder = array_chunks.remainder(); + + result_remainder.into_iter().zip(array_remainder).for_each( + |(scalar_result, scalar_input)| { + *scalar_result = scalar_op(*scalar_input); + }, + ); let data = ArrayData::new( T::DATA_TYPE, @@ -65,26 +84,26 @@ where None, array.data_ref().null_buffer().cloned(), 0, - vec![buffer], + vec![result.into()], vec![], ); Ok(PrimitiveArray::::from(Arc::new(data))) } -/// SIMD vectorized version of `signed_unary_math_op` above. #[cfg(simd)] -fn simd_signed_unary_math_op( +fn simd_float_unary_math_op( array: &PrimitiveArray, simd_op: SIMD_OP, scalar_op: SCALAR_OP, ) -> Result> where - T: datatypes::ArrowSignedNumericType, - SIMD_OP: Fn(T::SignedSimd) -> T::SignedSimd, + T: datatypes::ArrowFloatNumericType, + SIMD_OP: Fn(T::Simd) -> T::Simd, SCALAR_OP: Fn(T::Native) -> T::Native, { 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); let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes); @@ -94,9 +113,9 @@ where .borrow_mut() .zip(array_chunks.borrow_mut()) .for_each(|(result_slice, input_slice)| { - let simd_input = T::load_signed(input_slice); - let simd_result = T::signed_unary_op(simd_input, &simd_op); - T::write_signed(simd_result, result_slice); + let simd_input = T::load(input_slice); + let simd_result = T::unary_op(simd_input, &simd_op); + T::write(simd_result, result_slice); }); let result_remainder = result_chunks.into_remainder(); @@ -536,7 +555,29 @@ where #[cfg(simd)] return simd_signed_unary_math_op(array, |x| -x, |x| -x); #[cfg(not(simd))] - return signed_unary_math_op(array, |x| -x); + return Ok(unary(array, |x| -x)); +} + +/// Raise array with floating point values to the power of a scalar. +pub fn powf_scalar( + array: &PrimitiveArray, + raise: T::Native, +) -> Result> +where + T: datatypes::ArrowFloatNumericType, + T::Native: Pow, +{ + #[cfg(simd)] + { + let raise_vector = T::init(raise); + return simd_float_unary_math_op( + array, + |x| T::pow(x, raise_vector), + |x| x.pow(raise), + ); + } + #[cfg(not(simd))] + return Ok(unary(array, |x| x.pow(raise))); } /// Perform `left * right` operation on two arrays. If either left or right value is null @@ -808,4 +849,16 @@ mod tests { .collect(); assert_eq!(expected, actual); } + + #[test] + fn test_primitive_array_raise_power_scalar() { + let a = Float64Array::from(vec![1.0, 2.0, 3.0]); + let actual = powf_scalar(&a, 2.0).unwrap(); + let expected = Float64Array::from(vec![1.0, 4.0, 9.0]); + assert_eq!(expected, actual); + let a = Float64Array::from(vec![Some(1.0), None, Some(3.0)]); + let actual = powf_scalar(&a, 2.0).unwrap(); + let expected = Float64Array::from(vec![Some(1.0), None, Some(9.0)]); + assert_eq!(expected, actual); + } } diff --git a/rust/arrow/src/datatypes.rs b/rust/arrow/src/datatypes.rs index 2231da81c37..096a9305891 100644 --- a/rust/arrow/src/datatypes.rs +++ b/rust/arrow/src/datatypes.rs @@ -605,6 +605,8 @@ where /// Writes a SIMD result back to a slice fn write(simd_result: Self::Simd, slice: &mut [Self::Native]); + + fn unary_op Self::Simd>(a: Self::Simd, op: F) -> Self::Simd; } #[cfg(not(simd))] @@ -806,6 +808,14 @@ macro_rules! make_numeric_type { fn write(simd_result: Self::Simd, slice: &mut [Self::Native]) { unsafe { simd_result.write_to_slice_unaligned_unchecked(slice) }; } + + #[inline] + fn unary_op Self::Simd>( + a: Self::Simd, + op: F, + ) -> Self::Simd { + op(a) + } } #[cfg(not(simd))] @@ -909,6 +919,32 @@ make_signed_numeric_type!(Int64Type, i64x8); make_signed_numeric_type!(Float32Type, f32x16); make_signed_numeric_type!(Float64Type, f64x8); +#[cfg(simd)] +pub trait ArrowFloatNumericType: ArrowNumericType { + fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd; +} + +#[cfg(not(simd))] +pub trait ArrowFloatNumericType: ArrowNumericType {} + +macro_rules! make_float_numeric_type { + ($impl_ty:ty, $simd_ty:ident) => { + #[cfg(simd)] + impl ArrowFloatNumericType for $impl_ty { + #[inline] + fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd { + base.powf(raise) + } + } + + #[cfg(not(simd))] + impl ArrowFloatNumericType for $impl_ty {} + }; +} + +make_float_numeric_type!(Float32Type, f32x16); +make_float_numeric_type!(Float64Type, f64x8); + /// A subtype of primitive type that represents temporal values. pub trait ArrowTemporalType: ArrowPrimitiveType {}