This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 599a63e ARROW-11428: [Rust] Add power_scalar kernel
599a63e is described below
commit 599a63ed25cfc2ff59415ec4d04e9ee90c022e7e
Author: Ritchie Vink <[email protected]>
AuthorDate: Tue Feb 9 10:12:40 2021 +0200
ARROW-11428: [Rust] Add power_scalar kernel
Adds a SISD and SIMD kernel to raise a `Float32/64` array to a power of a
`scalar` of the same type. We could also make a thin `sqrt` wrapper.
I also added a `unary_op` fn to `ArrowNumeric` type as this seemed the most
generic way to implement this. Next PR I could add support for a binary version
of this (e.g. array to the power of array).
_edit_:
The `ArrowFloatNumericType` trait was added because the
[Simd::powf](https://rust-lang.github.io/packed_simd/packed_simd_2/struct.Simd.html#method.powf-6)
is only available for float arrays (e.g. `[f32, N]`, `[f64, N]`). However, the
*packed_simd* crate doesn't expose this functionality via a trait, but directly
on the type, hence the extra trait.
Closes #9361 from ritchie46/power_kernel
Authored-by: Ritchie Vink <[email protected]>
Signed-off-by: Neville Dipale <[email protected]>
---
rust/arrow/src/compute/kernels/arithmetic.rs | 99 +++++++++++++++++++++-------
rust/arrow/src/datatypes.rs | 36 ++++++++++
2 files changed, 112 insertions(+), 23 deletions(-)
diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs
b/rust/arrow/src/compute/kernels/arithmetic.rs
index 06d7c90..0677566 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<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,
@@ -65,26 +84,26 @@ where
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);
+ }
}
diff --git a/rust/arrow/src/datatypes.rs b/rust/arrow/src/datatypes.rs
index 2231da8..096a930 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<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 {}
