From 6817f4a6d9d020566c8df5f5f97378cc73d6bfe6 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Tue, 14 Jul 2026 13:20:07 -0600
Subject: [PATCH] perf: optimize spark_ceil in datafusion-comet-spark-expr

---
 native/spark-expr/Cargo.toml              |  4 ++
 native/spark-expr/benches/ceil.rs         | 65 +++++++++++++++++++++++
 native/spark-expr/src/math_funcs/ceil.rs  | 15 +++++-
 native/spark-expr/src/math_funcs/utils.rs |  2 +-
 4 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 native/spark-expr/benches/ceil.rs

diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 6f90181ee0..52235f1c26 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -139,3 +139,7 @@ harness = false
 [[bench]]
 name = "to_json"
 harness = false
+
+[[bench]]
+name = "ceil"
+harness = false
diff --git a/native/spark-expr/benches/ceil.rs b/native/spark-expr/benches/ceil.rs
new file mode 100644
index 0000000000..0436860f11
--- /dev/null
+++ b/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);
diff --git a/native/spark-expr/src/math_funcs/ceil.rs b/native/spark-expr/src/math_funcs/ceil.rs
index 324a118fcb..08f4718cea 100644
--- a/native/spark-expr/src/math_funcs/ceil.rs
+++ b/native/spark-expr/src/math_funcs/ceil.rs
@@ -78,7 +78,20 @@ pub fn spark_ceil(
 #[inline]
 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)
+    // 128-bit division is a compiler intrinsic call, while 64-bit division maps to a single
+    // hardware instruction. Decimals with precision <= 18 always fit in 64 bits, so take that
+    // path whenever both the value and the divisor do, and fall back to 128-bit otherwise.
+    let div_i64 = i64::try_from(div).ok();
+    move |x: i128| match (div_i64, i64::try_from(x)) {
+        (Some(d), Ok(x)) => {
+            let quotient = x / d;
+            // `d` is positive, so a positive remainder means the truncated quotient is below
+            // the true value and has to be rounded up.
+            let adjusted = if x % d > 0 { quotient + 1 } else { quotient };
+            adjusted as i128
+        }
+        _ => div_ceil(x, div),
+    }
 }
 
 #[cfg(test)]
diff --git a/native/spark-expr/src/math_funcs/utils.rs b/native/spark-expr/src/math_funcs/utils.rs
index 6109c8d161..5cec6a2d0a 100644
--- a/native/spark-expr/src/math_funcs/utils.rs
+++ b/native/spark-expr/src/math_funcs/utils.rs
@@ -57,7 +57,7 @@ pub(crate) fn make_decimal_array(
     array: &ArrayRef,
     precision: u8,
     scale: i8,
-    f: &dyn Fn(i128) -> i128,
+    f: impl Fn(i128) -> i128,
 ) -> Result<ColumnarValue, DataFusionError> {
     let array = array.as_primitive::<Decimal128Type>();
     let result: Decimal128Array = arrow::compute::kernels::arity::unary(array, f);