[Rust] [Experiment] Trigonometry kernels

rust/arrow/Cargo.toml

@@ -106,10 +106,18 @@ harness = false
 name = "filter_kernels"
 harness = false
 
+[[bench]]
+name = "math_kernels"
+harness = false
+
 [[bench]]
 name = "take_kernels"
 harness = false
 
+[[bench]]
+name = "trigonometry_kernels"
+harness = false
+
 [[bench]]
 name = "length_kernel"
 harness = false

rust/arrow/benches/math_kernels.rs

@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+extern crate arrow;
+
+use arrow::datatypes::Float32Type;
+use arrow::util::bench_util::*;
+use arrow::{array::*, compute::*};
+
+fn bench_powf(array: &Float32Array, raise: f32) {
+    criterion::black_box(powf_scalar(array, raise));
+}
+
+fn bench_powi(array: &Float32Array, raise: i32) {
+    criterion::black_box(powi(array, raise));
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let array = create_primitive_array::<Float32Type>(512, 0.0);
+    let array_nulls = create_primitive_array::<Float32Type>(512, 0.5);
+
+    c.bench_function("powf(2.0) 512", |b| b.iter(|| bench_powf(&array, 2.0)));
+    c.bench_function("powf(4.0) 512", |b| b.iter(|| bench_powf(&array, 4.0)));
+    c.bench_function("powf(32.0) 512", |b| b.iter(|| bench_powf(&array, 32.0)));
+    c.bench_function("powf(2.0) nulls 512", |b| {
+        b.iter(|| bench_powf(&array_nulls, 2.0))
+    });
+    c.bench_function("powf(4.0) nulls 512", |b| {
+        b.iter(|| bench_powf(&array_nulls, 4.0))
+    });
+    c.bench_function("powf(32.0) nulls 512", |b| {
+        b.iter(|| bench_powf(&array_nulls, 32.0))
+    });
+
+    c.bench_function("powi(2) 512", |b| b.iter(|| bench_powi(&array, 2)));
+    c.bench_function("powi(4) 512", |b| b.iter(|| bench_powi(&array, 4)));
+    c.bench_function("powi(32) 512", |b| b.iter(|| bench_powi(&array, 32)));
+    c.bench_function("powi(2) nulls 512", |b| {
+        b.iter(|| bench_powi(&array_nulls, 2))
+    });
+    c.bench_function("powi(4) nulls 512", |b| {
+        b.iter(|| bench_powi(&array_nulls, 4))
+    });
+    c.bench_function("powi(32) nulls 512", |b| {
+        b.iter(|| bench_powi(&array_nulls, 32))
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

rust/arrow/benches/trigonometry_kernels.rs

@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+extern crate arrow;
+
+use arrow::{
+    array::*, compute::*, datatypes::Float32Type,
+    util::bench_util::create_primitive_array,
+};
+
+fn bench_haversine(
+    lat_a: &Float32Array,
+    lng_a: &Float32Array,
+    lat_b: &Float32Array,
+    lng_b: &Float32Array,
+) {
+    criterion::black_box(haversine(lat_a, lng_a, lat_b, lng_b, 6371000.0).unwrap());
+}
+fn bench_sin(array: &Float32Array) {
+    criterion::black_box(sin(array));
+}
+
+fn bench_cos(array: &Float32Array) {
+    criterion::black_box(cos(array));
+}
+
+fn bench_tan(array: &Float32Array) {
+    criterion::black_box(tan(array));
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let lat_a = create_primitive_array(512, 0.0);
+    let lng_a = create_primitive_array(512, 0.0);
+    let lat_b = create_primitive_array(512, 0.0);
+    let lng_b = create_primitive_array(512, 0.0);
+
+    c.bench_function("haversine_unary 512", |b| {
+        b.iter(|| bench_haversine(&lat_a, &lng_a, &lat_b, &lng_b))
+    });
+
+    let lat_a = create_primitive_array(512, 0.5);
+    let lng_a = create_primitive_array(512, 0.5);
+    let lat_b = create_primitive_array(512, 0.5);
+    let lng_b = create_primitive_array(512, 0.5);
+
+    c.bench_function("haversine_unary_nulls 512", |b| {
+        b.iter(|| bench_haversine(&lat_a, &lng_a, &lat_b, &lng_b))
+    });
+
+    let array = create_primitive_array::<Float32Type>(512, 0.0);
+    let array_nulls = create_primitive_array::<Float32Type>(512, 0.5);
+
+    c.bench_function("sin 512", |b| b.iter(|| bench_sin(&array)));
+    c.bench_function("cos 512", |b| b.iter(|| bench_cos(&array)));
+    c.bench_function("tan 512", |b| b.iter(|| bench_tan(&array)));
+    c.bench_function("sin nulls 512", |b| b.iter(|| bench_sin(&array_nulls)));
+    c.bench_function("cos nulls 512", |b| b.iter(|| bench_cos(&array_nulls)));
+    c.bench_function("tan nulls 512", |b| b.iter(|| bench_tan(&array_nulls)));
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

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

@@ -22,24 +22,25 @@
 //! `RUSTFLAGS="-C target-feature=+avx2"` for example.  See the documentation
 //! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
 
+#[cfg(simd)]
+use std::borrow::BorrowMut;
 use std::ops::{Add, Div, Mul, Neg, Sub};
+#[cfg(simd)]
+use std::slice::{ChunksExact, ChunksExactMut};
 use std::sync::Arc;
 
 use num::{One, Zero};
 
 use crate::buffer::Buffer;
 #[cfg(simd)]
 use crate::buffer::MutableBuffer;
-use crate::compute::{kernels::arity::unary, util::combine_option_bitmap};
+#[cfg(not(simd))]
+use crate::compute::kernels::arity::unary;
+use crate::compute::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};
 
 /// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values.
 #[cfg(simd)]
@@ -90,55 +91,6 @@ where
     Ok(PrimitiveArray::<T>::from(Arc::new(data)))
 }
 
-#[cfg(simd)]
-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::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);
-    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(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();
-    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![result.into()],
-        vec![],
-    );
-    Ok(PrimitiveArray::<T>::from(Arc::new(data)))
-}
-
 /// Helper function to perform math lambda function on values from two arrays. If either
 /// left or right value is null then the output value is also null, so `1 + null` is
 /// `null`.
@@ -558,28 +510,6 @@ where
     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
 /// then the result is also null.
 pub fn multiply<T>(
@@ -849,16 +779,4 @@ 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);
-    }
 }