perf: optimize spark_floor (up to 4x faster)

native/spark-expr/Cargo.toml

@@ -139,3 +139,7 @@ harness = false
 [[bench]]
 name = "to_json"
 harness = false
+
+[[bench]]
+name = "floor"
+harness = false

native/spark-expr/benches/floor.rs

@@ -0,0 +1,102 @@
+// 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.
+
+use arrow::array::{ArrayRef, Decimal128Array, Float64Array};
+use arrow::datatypes::DataType;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::spark_floor;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const ROWS: usize = 8192;
+
+/// Unscaled decimal values that fit in 64 bits, which is the common case, with every 10th row
+/// null.
+fn decimal_array(precision: u8, scale: i8) -> ArrayRef {
+    let values: Vec<Option<i128>> = (0..ROWS)
+        .map(|i| {
+            if i % 10 == 0 {
+                None
+            } else {
+                let v = (i as i128) * 987_654_321 % 100_000_000_000_000_000;
+                Some(if i % 2 == 0 { v } else { -v })
+            }
+        })
+        .collect();
+    Arc::new(
+        Decimal128Array::from(values)
+            .with_precision_and_scale(precision, scale)
+            .unwrap(),
+    )
+}
+
+/// Unscaled decimal values too wide for 64 bits, exercising the 128-bit fallback.
+fn wide_decimal_array(precision: u8, scale: i8) -> ArrayRef {
+    let values: Vec<Option<i128>> = (0..ROWS)
+        .map(|i| {
+            if i % 10 == 0 {
+                None
+            } else {
+                let v = (i as i128 + 1) * 10_000_000_000_000_000_000_000_000_000;
+                Some(if i % 2 == 0 { v } else { -v })
+            }
+        })
+        .collect();
+    Arc::new(
+        Decimal128Array::from(values)
+            .with_precision_and_scale(precision, scale)
+            .unwrap(),
+    )
+}
+
+fn float_array() -> ArrayRef {
+    Arc::new(Float64Array::from(
+        (0..ROWS)
+            .map(|i| i as f64 * 1.5 - 4096.0)
+            .collect::<Vec<_>>(),
+    ))
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let dec_2 = decimal_array(38, 2);
+    c.bench_function("spark_floor: decimal128(38,2)", |b| {
+        let args = vec![ColumnarValue::Array(Arc::clone(&dec_2))];
+        b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(37, 0)).unwrap()))
+    });
+
+    let dec_18 = decimal_array(38, 18);
+    c.bench_function("spark_floor: decimal128(38,18)", |b| {
+        let args = vec![ColumnarValue::Array(Arc::clone(&dec_18))];
+        b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(21, 0)).unwrap()))
+    });
+
+    let wide = wide_decimal_array(38, 6);
+    c.bench_function("spark_floor: decimal128(38,6) wide values", |b| {
+        let args = vec![ColumnarValue::Array(Arc::clone(&wide))];
+        b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Decimal128(33, 0)).unwrap()))
+    });
+
+    let floats = float_array();
+    c.bench_function("spark_floor: float64", |b| {
+        let args = vec![ColumnarValue::Array(Arc::clone(&floats))];
+        b.iter(|| black_box(spark_floor(black_box(&args), &DataType::Float64).unwrap()))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

native/spark-expr/src/math_funcs/floor.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 use crate::downcast_compute_op;
-use crate::math_funcs::utils::{get_precision_scale, make_decimal_array, make_decimal_scalar};
+use crate::math_funcs::utils::{
+    dispatch_pow10, get_precision_scale, make_decimal_array, make_decimal_scalar,
+};
 use arrow::array::{Array, ArrowNativeTypeOp};
 use arrow::array::{Float32Array, Float64Array, Int64Array};
 use arrow::datatypes::DataType;
@@ -45,10 +47,13 @@ pub fn spark_floor(
                 let result = array.as_any().downcast_ref::<Int64Array>().unwrap();
                 Ok(ColumnarValue::Array(Arc::new(result.clone())))
             }
-            DataType::Decimal128(_, scale) if *scale > 0 => {
-                let f = decimal_floor_f(scale);
+            DataType::Decimal128(_, input_scale) if *input_scale > 0 => {
                 let (precision, scale) = get_precision_scale(data_type);
-                make_decimal_array(array, precision, scale, &f)
+                dispatch_pow10!(
+                    *input_scale,
+                    EXP => make_decimal_array(array, precision, scale, decimal_floor_pow10::<EXP>),
+                    make_decimal_array(array, precision, scale, decimal_floor_f(*input_scale))
+                )
             }
             other => Err(DataFusionError::Internal(format!(
                 "Unsupported data type {other:?} for function floor",
@@ -62,8 +67,8 @@ pub fn spark_floor(
                 a.map(|x| x.floor() as i64),
             ))),
             ScalarValue::Int64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))),
-            ScalarValue::Decimal128(a, _, scale) if *scale > 0 => {
-                let f = decimal_floor_f(scale);
+            ScalarValue::Decimal128(a, _, input_scale) if *input_scale > 0 => {
+                let f = decimal_floor_f(*input_scale);
                 let (precision, scale) = get_precision_scale(data_type);
                 make_decimal_scalar(a, precision, scale, &f)
             }
@@ -76,11 +81,32 @@ pub fn spark_floor(
 }
 
 #[inline]
-fn decimal_floor_f(scale: &i8) -> impl Fn(i128) -> i128 {
-    let div = 10_i128.pow_wrapping(*scale as u32);
+fn decimal_floor_f(scale: i8) -> impl Fn(i128) -> i128 {
+    let div = 10_i128.pow_wrapping(scale as u32);
     move |x: i128| div_floor(x, div)
 }
 
+/// Floor-divides an unscaled decimal by `10^EXP`.
+///
+/// `EXP` is a compile-time constant so that the divisor is folded in and the division lowered to a
+/// multiply-and-shift. A 128-bit division is always a libcall, even by a constant, so values that
+/// fit in 64 bits take a 64-bit path; unscaled decimals rarely exceed that range.
+#[inline]
+fn decimal_floor_pow10<const EXP: u32>(x: i128) -> i128 {
+    match i64::try_from(x) {
+        Ok(x) => div_floor(x, const { 10_i64.pow(EXP) }) as i128,
+        Err(_) => decimal_floor_wide(x, const { 10_i128.pow(EXP) }),
+    }
+}
+
+/// Kept out of line so that the libcall and its stack frame stay out of the loop body of every
+/// [`decimal_floor_pow10`] instantiation.
+#[cold]
+#[inline(never)]
+fn decimal_floor_wide(x: i128, div: i128) -> i128 {
+    div_floor(x, div)
+}
+
 #[cfg(test)]
 mod test {
     use crate::spark_floor;

native/spark-expr/src/math_funcs/utils.rs

@@ -41,6 +41,94 @@ macro_rules! downcast_compute_op {
     }};
 }
 
+/// Evaluates `$body` with `$exp` bound as a compile-time constant equal to `$scale`, for every
+/// decimal scale whose divisor `10^scale` fits in an `i64`. Larger scales evaluate `$fallback`,
+/// which must derive the divisor at runtime.
+///
+/// Specializing on the scale lets the divisor be folded into the generated code, turning the
+/// division by `10^scale` into a multiply-and-shift.
+macro_rules! dispatch_pow10 {
+    ($scale:expr, $exp:ident => $body:expr, $fallback:expr) => {
+        match $scale {
+            1 => {
+                const $exp: u32 = 1;
+                $body
+            }
+            2 => {
+                const $exp: u32 = 2;
+                $body
+            }
+            3 => {
+                const $exp: u32 = 3;
+                $body
+            }
+            4 => {
+                const $exp: u32 = 4;
+                $body
+            }
+            5 => {
+                const $exp: u32 = 5;
+                $body
+            }
+            6 => {
+                const $exp: u32 = 6;
+                $body
+            }
+            7 => {
+                const $exp: u32 = 7;
+                $body
+            }
+            8 => {
+                const $exp: u32 = 8;
+                $body
+            }
+            9 => {
+                const $exp: u32 = 9;
+                $body
+            }
+            10 => {
+                const $exp: u32 = 10;
+                $body
+            }
+            11 => {
+                const $exp: u32 = 11;
+                $body
+            }
+            12 => {
+                const $exp: u32 = 12;
+                $body
+            }
+            13 => {
+                const $exp: u32 = 13;
+                $body
+            }
+            14 => {
+                const $exp: u32 = 14;
+                $body
+            }
+            15 => {
+                const $exp: u32 = 15;
+                $body
+            }
+            16 => {
+                const $exp: u32 = 16;
+                $body
+            }
+            17 => {
+                const $exp: u32 = 17;
+                $body
+            }
+            18 => {
+                const $exp: u32 = 18;
+                $body
+            }
+            _ => $fallback,
+        }
+    };
+}
+
+pub(crate) use dispatch_pow10;
+
 #[inline]
 pub(crate) fn make_decimal_scalar(
     a: &Option<i128>,
@@ -52,12 +140,14 @@ pub(crate) fn make_decimal_scalar(
     Ok(ColumnarValue::Scalar(result))
 }
 
+/// Generic over the operation so that it is monomorphized into the element loop rather than
+/// called through a vtable for every value.
 #[inline]
-pub(crate) fn make_decimal_array(
+pub(crate) fn make_decimal_array<F: Fn(i128) -> i128>(
     array: &ArrayRef,
     precision: u8,
     scale: i8,
-    f: &dyn Fn(i128) -> i128,
+    f: F,
 ) -> Result<ColumnarValue, DataFusionError> {
     let array = array.as_primitive::<Decimal128Type>();
     let result: Decimal128Array = arrow::compute::kernels::arity::unary(array, f);