diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..58c4d02d9f394 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -98,6 +98,11 @@ env_logger = { workspace = true } rand = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync"] } +[[bench]] +harness = false +name = "round_dense" +required-features = ["math_expressions"] + [[bench]] harness = false name = "ascii" diff --git a/datafusion/functions/benches/round_dense.rs b/datafusion/functions/benches/round_dense.rs new file mode 100644 index 0000000000000..2c37849bde489 --- /dev/null +++ b/datafusion/functions/benches/round_dense.rs @@ -0,0 +1,94 @@ +// 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. + +//! Microbenchmark for `round(float_array, scalar_decimal_places)` over a +//! Float column with no NULLs — the dense elementwise-rounding path. + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field, Float32Type, Float64Type}; +use arrow::util::bench_util::create_primitive_array; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::math::round::RoundFunc; +use std::hint::black_box; +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let round_fn = RoundFunc::new(); + let config_options = Arc::new(ConfigOptions::default()); + + for size in [1024usize, 4096, 8192] { + // Float64, no nulls. + let f64_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f64_args = vec![ + ColumnarValue::Array(Arc::clone(&f64_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f64", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float64, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float64, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + + // Float32, no nulls. + let f32_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f32_args = vec![ + ColumnarValue::Array(Arc::clone(&f32_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f32", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float32, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float32, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs index 62f1c3540b9ce..3bb9739bf715b 100644 --- a/datafusion/functions/src/math/round.rs +++ b/datafusion/functions/src/math/round.rs @@ -23,9 +23,9 @@ use arrow::datatypes::DataType::{ Int64, UInt8, UInt16, UInt32, UInt64, }; use arrow::datatypes::{ - ArrowNativeTypeOp, DataType, Decimal32Type, Decimal64Type, Decimal128Type, - Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, - Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + ArrowNativeTypeOp, ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use arrow::datatypes::{Field, FieldRef}; use arrow::error::ArrowError; @@ -495,22 +495,8 @@ fn round_columnar( )?, } } - (Float64, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ - } - (Float32, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ - } + (Float64, _) => round_float_column::(&value_array, decimal_places)?, + (Float32, _) => round_float_column::(&value_array, decimal_places)?, (Decimal32(input_precision, scale), Decimal32(precision, new_scale)) => { // reduce scale to reclaim integer precision let result = calculate_binary_decimal_math_cast::< @@ -859,15 +845,84 @@ fn round_integer_array( } } -fn round_float(value: T, decimal_places: i32) -> Result +/// Rounds a float array to `decimal_places`, taking the hoisted-factor fast +/// path when it applies and falling back to the shared binary-math kernel +/// otherwise. +fn round_float_column( + value_array: &ArrayRef, + decimal_places: &ColumnarValue, +) -> Result where - T: num_traits::Float, + PT: ArrowPrimitiveType, + PT::Native: num_traits::Float, { - let factor = T::from(10_f64.powi(decimal_places)).ok_or_else(|| { + if let Some(arr) = round_float_fast::(value_array, decimal_places)? { + return Ok(arr); + } + let result = calculate_binary_math::( + value_array.as_ref(), + decimal_places, + round_float::, + )?; + Ok(result as _) +} + +/// Fast path for rounding a float array to a scalar number of `decimal_places`. +/// +/// The shared `calculate_binary_math` kernel routes through `try_unary` and +/// re-evaluates `round_float` (including `10f64.powi(decimal_places)` and a +/// `Result` check) for every element. When `decimal_places` is a non-null +/// scalar and the value column has no nulls, we can hoist the scaling factor +/// out of the loop and use the infallible `unary` kernel instead, which the +/// compiler can autovectorize. Requiring no null slots keeps the output +/// bit-identical to the fallible kernel: `unary` writes a computed value into +/// every slot while `try_unary` leaves null slots zeroed, so the two only agree +/// when there are no null slots. Returns `Ok(None)` when the fast path does not +/// apply, so the caller falls back to the shared kernel. +fn round_float_fast( + value_array: &ArrayRef, + decimal_places: &ColumnarValue, +) -> Result> +where + PT: ArrowPrimitiveType, + PT::Native: num_traits::Float, +{ + // Bring `Float` into scope so `.round()` resolves on the `PT::Native` + // projection below. + use num_traits::Float; + + let ColumnarValue::Scalar(ScalarValue::Int32(Some(decimal_places))) = decimal_places + else { + return Ok(None); + }; + + let prim = value_array.as_primitive::(); + if prim.null_count() != 0 { + return Ok(None); + } + + // Compute the scaling factor once, via the same helper `round_float` uses, + // so the arithmetic (and therefore every rounded bit) is identical. + let factor = round_factor::(*decimal_places)?; + + let result = prim.unary::<_, PT>(|value| (value * factor).round() / factor); + Ok(Some(Arc::new(result) as ArrayRef)) +} + +/// Computes the power-of-ten scaling factor used to round to `decimal_places`. +fn round_factor(decimal_places: i32) -> Result { + T::from(10_f64.powi(decimal_places)).ok_or_else(|| { ArrowError::ComputeError(format!( "Invalid value for decimal places: {decimal_places}" )) - })?; + }) +} + +fn round_float(value: T, decimal_places: i32) -> Result +where + T: num_traits::Float, +{ + let factor = round_factor::(decimal_places)?; Ok((value * factor).round() / factor) }