diff --git a/.gitignore b/.gitignore index c6d763755bf..39564fe2b50 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ Cargo.lock fixtures settings.json dev/ +.idea/ diff --git a/Cargo.toml b/Cargo.toml index 0775ae6a8d0..a96330f0603 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,9 @@ ahash = { version = "0.7", optional = true } parquet2 = { version = "0.3", optional = true, default_features = false, features = ["stream"] } +# for division/remainder optimization at runtime +strength_reduce = "0.2" + [dev-dependencies] rand = "0.8" criterion = "0.3" @@ -167,3 +170,7 @@ harness = false [[bench]] name = "write_ipc" harness = false + +[[bench]] +name = "arithmetic_kernels" +harness = false diff --git a/benches/arithmetic_kernels.rs b/benches/arithmetic_kernels.rs new file mode 100644 index 00000000000..f9c26cf8ab6 --- /dev/null +++ b/benches/arithmetic_kernels.rs @@ -0,0 +1,53 @@ +// 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; + +use arrow2::array::*; +use arrow2::util::bench_util::*; +use arrow2::{ + compute::arithmetics::basic::div::div_scalar, datatypes::DataType, types::NativeType, +}; +use num::NumCast; +use std::ops::Div; + +fn bench_div_scalar(lhs: &PrimitiveArray, rhs: &T) +where + T: NativeType + Div + NumCast, +{ + criterion::black_box(div_scalar(lhs, rhs)); +} + +fn add_benchmark(c: &mut Criterion) { + let size = 65536; + let arr = create_primitive_array_with_seed::(size, DataType::UInt64, 0.0, 43); + let arr_i64 = create_primitive_array_with_seed::(size, DataType::Int64, 0.0, 43); + + c.bench_function("divide_scalar 4", |b| { + // 4 is a very fast optimizable divisor + b.iter(|| bench_div_scalar(&arr, &4)) + }); + c.bench_function("divide_scalar prime", |b| { + // large prime number that is probably harder to simplify + b.iter(|| bench_div_scalar(&arr, &524287)) + }); +} + +criterion_group!(benches, add_benchmark); +criterion_main!(benches); diff --git a/src/compute/arithmetics/basic/div.rs b/src/compute/arithmetics/basic/div.rs index 8556c995a9a..84b55ed5f5a 100644 --- a/src/compute/arithmetics/basic/div.rs +++ b/src/compute/arithmetics/basic/div.rs @@ -1,8 +1,9 @@ //! Definition of basic div operations with primitive arrays use std::ops::Div; -use num::{CheckedDiv, Zero}; +use num::{CheckedDiv, NumCast, Zero}; +use crate::datatypes::DataType; use crate::{ array::{Array, PrimitiveArray}, compute::{ @@ -12,6 +13,7 @@ use crate::{ error::{ArrowError, Result}, types::NativeType, }; +use strength_reduce::StrengthReducedU64; /// Divides two primitive arrays with the same type. /// Panics if the divisor is zero of one pair of values overflows. @@ -109,10 +111,28 @@ where /// ``` pub fn div_scalar(lhs: &PrimitiveArray, rhs: &T) -> PrimitiveArray where - T: NativeType + Div, + T: NativeType + Div + NumCast, { let rhs = *rhs; - unary(lhs, |a| a / rhs, lhs.data_type().clone()) + match T::DATA_TYPE { + DataType::UInt64 => { + // Safety: + // We need to go from generic to concrete type and back to generic. + // We know the types are the same, but the compiler needs some convincing + unsafe { + let lhs = &*(lhs as *const PrimitiveArray as *const PrimitiveArray); + let rhs = rhs.to_u64().unwrap(); + + let reduced_div = StrengthReducedU64::new(rhs); + std::mem::transmute::, PrimitiveArray>(unary( + lhs, + |a| a / reduced_div, + lhs.data_type().clone(), + )) + } + } + _ => unary(lhs, |a| a / rhs, lhs.data_type().clone()), + } } /// Checked division of a primitive array of type T by a scalar T. If the @@ -141,7 +161,7 @@ where // Implementation of ArrayDiv trait for PrimitiveArrays with a scalar impl ArrayDiv for PrimitiveArray where - T: NativeType + Div + NotI128, + T: NativeType + Div + NotI128 + NumCast, { type Output = Self; diff --git a/src/compute/arithmetics/mod.rs b/src/compute/arithmetics/mod.rs index 496f02d72f3..ad7eb846c80 100644 --- a/src/compute/arithmetics/mod.rs +++ b/src/compute/arithmetics/mod.rs @@ -61,7 +61,7 @@ pub mod time; use std::ops::{Add, Div, Mul, Neg, Rem, Sub}; -use num::Zero; +use num::{NumCast, Zero}; use crate::datatypes::{DataType, TimeUnit}; use crate::error::{ArrowError, Result}; @@ -265,7 +265,8 @@ where + Add + Sub + Mul - + Rem, + + Rem + + NumCast, { match op { Operator::Add => Ok(basic::add::add_scalar(lhs, rhs)),