ARROW-11428: [Rust] Add power_scalar kernel

rust/arrow/src/compute/kernels/arithmetic.rs

@@ -30,61 +30,80 @@ 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<T, F>(
+/// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values.
+#[cfg(simd)]
+fn simd_signed_unary_math_op<T, SIMD_OP, SCALAR_OP>(
     array: &PrimitiveArray<T>,
-    op: F,
+    simd_op: SIMD_OP,
+    scalar_op: SCALAR_OP,
 ) -> Result<PrimitiveArray<T>>
 where
     T: datatypes::ArrowSignedNumericType,
-    T::Native: Neg<Output = T::Native>,
-    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::<T::Native>();
+    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,
         array.len(),
         None,
         array.data_ref().null_buffer().cloned(),
         0,
-        vec![buffer],
+        vec![result.into()],
         vec![],
     );
     Ok(PrimitiveArray::<T>::from(Arc::new(data)))
 }
 
-/// SIMD vectorized version of `signed_unary_math_op` above.
 #[cfg(simd)]
-fn simd_signed_unary_math_op<T, SIMD_OP, SCALAR_OP>(
+fn simd_float_unary_math_op<T, SIMD_OP, SCALAR_OP>(
     array: &PrimitiveArray<T>,
     simd_op: SIMD_OP,
     scalar_op: SCALAR_OP,
 ) -> Result<PrimitiveArray<T>>
 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::<T::Native>();
+
     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<T>(
+    array: &PrimitiveArray<T>,
+    raise: T::Native,
+) -> Result<PrimitiveArray<T>>
+where
+    T: datatypes::ArrowFloatNumericType,
+    T::Native: Pow<T::Native, Output = T::Native>,
+{
+    #[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);
+    }
 }

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<F: Fn(Self::Simd) -> 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<F: Fn(Self::Simd) -> 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 {}