perf: optimize spark_ceil (3x faster)

native/spark-expr/Cargo.toml

@@ -163,6 +163,14 @@ harness = false
 name = "to_json"
 harness = false
 
+[[bench]]
+name = "floor"
+harness = false
+
+[[bench]]
+name = "ceil"
+harness = false
+
 [[bench]]
 name = "cast_float_to_string"
 harness = false

native/spark-expr/benches/ceil.rs

@@ -0,0 +1,65 @@
+// 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::Decimal128Array;
+use arrow::datatypes::DataType;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::spark_ceil;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const NUM_ROWS: i128 = 8192;
+
+fn decimal_args(precision: u8, scale: i8, spread: i128) -> Vec<ColumnarValue> {
+    let values: Vec<i128> = (0..NUM_ROWS)
+        .map(|i| (i - NUM_ROWS / 2) * spread + i % 7)
+        .collect();
+    let array = Decimal128Array::from(values)
+        .with_precision_and_scale(precision, scale)
+        .unwrap();
+    vec![ColumnarValue::Array(Arc::new(array))]
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    // decimal(18, 4): unscaled values fit comfortably in 64 bits
+    let narrow = decimal_args(18, 4, 1_000_003);
+    // decimal(38, 6): unscaled values exceed the 64-bit range
+    let wide = decimal_args(38, 6, 1_000_000_000_000_000_003);
+
+    let mut group = c.benchmark_group("ceil");
+    group.bench_function("ceil_decimal_18_4", |b| {
+        b.iter(|| {
+            black_box(spark_ceil(
+                black_box(&narrow),
+                black_box(&DataType::Decimal128(18, 0)),
+            ))
+        })
+    });
+    group.bench_function("ceil_decimal_38_6", |b| {
+        b.iter(|| {
+            black_box(spark_ceil(
+                black_box(&wide),
+                black_box(&DataType::Decimal128(38, 0)),
+            ))
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

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/ceil.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_ceil(
                 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_ceil_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_ceil_pow10::<EXP>),
+                    make_decimal_array(array, precision, scale, decimal_ceil_f(*input_scale))
+                )
             }
             other => Err(DataFusionError::Internal(format!(
                 "Unsupported data type {other:?} for function ceil",
@@ -62,10 +67,10 @@ pub fn spark_ceil(
                 a.map(|x| x.ceil() as i64),
             ))),
             ScalarValue::Int64(a) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(a.map(|x| x)))),
-            ScalarValue::Decimal128(a, _, scale) if *scale > 0 => {
-                let f = decimal_ceil_f(scale);
+            ScalarValue::Decimal128(a, _, input_scale) if *input_scale > 0 => {
+                let f = decimal_ceil_f(*input_scale);
                 let (precision, scale) = get_precision_scale(data_type);
-                make_decimal_scalar(a, precision, scale, &f)
+                make_decimal_scalar(a, precision, scale, f)
             }
             _ => Err(DataFusionError::Internal(format!(
                 "Unsupported data type {:?} for function ceil",
@@ -76,11 +81,32 @@ pub fn spark_ceil(
 }
 
 #[inline]
-fn decimal_ceil_f(scale: &i8) -> impl Fn(i128) -> i128 {
-    let div = 10_i128.pow_wrapping(*scale as u32);
+fn decimal_ceil_f(scale: i8) -> impl Fn(i128) -> i128 {
+    let div = 10_i128.pow_wrapping(scale as u32);
     move |x: i128| div_ceil(x, div)
 }
 
+/// Ceiling-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_ceil_pow10<const EXP: u32>(x: i128) -> i128 {
+    match i64::try_from(x) {
+        Ok(x) => div_ceil(x, const { 10_i64.pow(EXP) }) as i128,
+        Err(_) => decimal_ceil_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_ceil_pow10`] instantiation.
+#[cold]
+#[inline(never)]
+fn decimal_ceil_wide(x: i128, div: i128) -> i128 {
+    div_ceil(x, div)
+}
+
 #[cfg(test)]
 mod test {
     use crate::spark_ceil;
@@ -187,6 +213,36 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn test_ceil_decimal128_wide_array() -> Result<()> {
+        // Unscaled values that exceed i64::MAX (~9.22e18) exercise the wide fallback in
+        // decimal_ceil_pow10. Values chosen at scale 6 targeting Decimal128(38, 0):
+        //   20_000_000_000_000_000_000 / 10^6 = 20_000_000_000_000 (exact multiple)
+        //   20_000_000_000_000_000_001 / 10^6 -> ceil 20_000_000_000_001 (positive remainder)
+        //  -20_000_000_000_000_000_001 / 10^6 -> ceil -20_000_000_000_000 (negative remainder)
+        let array = Decimal128Array::from(vec![
+            Some(20_000_000_000_000_000_000_i128),
+            Some(20_000_000_000_000_000_001_i128),
+            Some(-20_000_000_000_000_000_001_i128),
+            None,
+        ])
+        .with_precision_and_scale(38, 6)?;
+        let args = vec![ColumnarValue::Array(Arc::new(array))];
+        let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Decimal128(38, 0))? else {
+            unreachable!()
+        };
+        let expected = Decimal128Array::from(vec![
+            Some(20_000_000_000_000_i128),
+            Some(20_000_000_000_001_i128),
+            Some(-20_000_000_000_000_i128),
+            None,
+        ])
+        .with_precision_and_scale(38, 0)?;
+        let actual = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
+        assert_eq!(actual, &expected);
+        Ok(())
+    }
+
     #[test]
     fn test_ceil_f32_scalar() -> Result<()> {
         let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.2345)))];

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,10 +67,10 @@ 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)
+                make_decimal_scalar(a, precision, scale, f)
             }
             _ => Err(DataFusionError::Internal(format!(
                 "Unsupported data type {:?} for function floor",
@@ -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;