You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/24 00:21:42 UTC
(datafusion-comet) branch main updated: perf: optimize `spark_ceil` (3x faster) (#4926)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 406e18271c perf: optimize `spark_ceil` (3x faster) (#4926)
406e18271c is described below
commit 406e18271c59c73ffd21c2f9ed5b3ff23a26ab2d
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Jul 23 18:21:37 2026 -0600
perf: optimize `spark_ceil` (3x faster) (#4926)
* perf: optimize spark_floor in datafusion-comet-spark-expr
* refactor: unify make_decimal_scalar/array signatures on impl Fn
Adopts impl Fn on both helpers so they are symmetric and no callsite needs to pass &f, addressing coordination feedback with #4926 which uses impl Fn too. Floor's scalar path drops its stray &f.
* chore: Codegen fallback explain for plan (#4891)
* perf: optimize spark_cast_int_to_int in datafusion-comet-spark-expr (#4920)
* perf: optimize spark_ceil in datafusion-comet-spark-expr
* perf: share dispatch_pow10 with floor; add wide-path decimal test
Ceil now uses the same dispatch_pow10 + decimal_ceil_pow10<const EXP> pattern that floor uses in the same batch, replacing the per-row runtime i64::try_from branch. The wide fallback is an out-of-line cold function mirroring floor. Adds test_ceil_decimal128_wide_array covering Decimal128(38, 6) unscaled values above i64::MAX for positive-remainder, negative-remainder, and exact-multiple cases so the fallback arm and the negative-remainder branch have direct coverage.
---------
Co-authored-by: Oleks V <co...@users.noreply.github.com>
---
native/spark-expr/Cargo.toml | 8 +++
native/spark-expr/benches/ceil.rs | 65 +++++++++++++++++++
native/spark-expr/benches/floor.rs | 102 ++++++++++++++++++++++++++++++
native/spark-expr/src/math_funcs/ceil.rs | 74 +++++++++++++++++++---
native/spark-expr/src/math_funcs/floor.rs | 44 ++++++++++---
native/spark-expr/src/math_funcs/utils.rs | 94 ++++++++++++++++++++++++++-
6 files changed, 367 insertions(+), 20 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 08a5584501..b5297d3da0 100644
--- a/native/spark-expr/Cargo.toml
+++ b/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
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/benches/floor.rs b/native/spark-expr/benches/floor.rs
new file mode 100644
index 0000000000..5005c40f00
--- /dev/null
+++ b/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);
diff --git a/native/spark-expr/src/math_funcs/ceil.rs b/native/spark-expr/src/math_funcs/ceil.rs
index 324a118fcb..292ef117e9 100644
--- a/native/spark-expr/src/math_funcs/ceil.rs
+++ b/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)))];
diff --git a/native/spark-expr/src/math_funcs/floor.rs b/native/spark-expr/src/math_funcs/floor.rs
index c7c1c9e904..973bbab708 100644
--- a/native/spark-expr/src/math_funcs/floor.rs
+++ b/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;
diff --git a/native/spark-expr/src/math_funcs/utils.rs b/native/spark-expr/src/math_funcs/utils.rs
index 6109c8d161..23e7b336ef 100644
--- a/native/spark-expr/src/math_funcs/utils.rs
+++ b/native/spark-expr/src/math_funcs/utils.rs
@@ -41,23 +41,113 @@ 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>,
precision: u8,
scale: i8,
- f: &dyn Fn(i128) -> i128,
+ f: impl Fn(i128) -> i128,
) -> Result<ColumnarValue, DataFusionError> {
let result = ScalarValue::Decimal128(a.map(f), precision, scale);
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(
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);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org