PoC strength_reduce division

.gitignore

@@ -6,3 +6,4 @@ Cargo.lock
 fixtures
 settings.json
 dev/
+.idea/

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

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<T>(lhs: &PrimitiveArray<T>, rhs: &T)
+where
+    T: NativeType + Div<Output = T> + NumCast,
+{
+    criterion::black_box(div_scalar(lhs, rhs));
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let size = 65536;
+    let arr = create_primitive_array_with_seed::<u64>(size, DataType::UInt64, 0.0, 43);
+    let arr_i64 = create_primitive_array_with_seed::<i64>(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);

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<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
 where
-    T: NativeType + Div<Output = T>,
+    T: NativeType + Div<Output = T> + 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<T> as *const PrimitiveArray<u64>);
+                let rhs = rhs.to_u64().unwrap();
+
+                let reduced_div = StrengthReducedU64::new(rhs);
+                std::mem::transmute::<PrimitiveArray<u64>, PrimitiveArray<T>>(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<T> ArrayDiv<T> for PrimitiveArray<T>
 where
-    T: NativeType + Div<Output = T> + NotI128,
+    T: NativeType + Div<Output = T> + NotI128 + NumCast,
 {
     type Output = Self;
 

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<Output = T>
         + Sub<Output = T>
         + Mul<Output = T>
-        + Rem<Output = T>,
+        + Rem<Output = T>
+        + NumCast,
 {
     match op {
         Operator::Add => Ok(basic::add::add_scalar(lhs, rhs)),