diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..ee8c8423e7 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 = "floor" +harness = false 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> = (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> = (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::>(), + )) +} + +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/floor.rs b/native/spark-expr/src/math_funcs/floor.rs index c7c1c9e904..7a53353cec 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::().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::), + 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(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..3893878381 100644 --- a/native/spark-expr/src/math_funcs/utils.rs +++ b/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, @@ -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 i128>( array: &ArrayRef, precision: u8, scale: i8, - f: &dyn Fn(i128) -> i128, + f: F, ) -> Result { let array = array.as_primitive::(); let result: Decimal128Array = arrow::compute::kernels::arity::unary(array, f);