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..49385d087af0d 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,59 @@ fn round_integer_array( } } -fn round_float(value: T, decimal_places: i32) -> Result +/// Rounds a float array to `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, the scaling factor can instead be hoisted out of the loop and the +/// infallible `unary` kernel used, which the compiler can autovectorize. +/// `unary` also computes over null slots, but it carries the input null buffer +/// through to the output, so those values stay masked. +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(|| { + // Bring `Float` into scope so `.round()` resolves on the `PT::Native` + // projection below. + use num_traits::Float; + + if let ColumnarValue::Scalar(ScalarValue::Int32(Some(decimal_places))) = + decimal_places + { + let factor = round_factor::(*decimal_places)?; + let result = value_array + .as_primitive::() + .unary::<_, PT>(|value| (value * factor).round() / factor); + return Ok(Arc::new(result) as ArrayRef); + } + + let result = calculate_binary_math::( + value_array.as_ref(), + decimal_places, + round_float::, + )?; + Ok(result as _) +} + +/// 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) } @@ -957,6 +987,7 @@ mod test { use std::sync::Arc; use arrow::array::{ArrayRef, Float32Array, Float64Array, Int64Array}; + use arrow::datatypes::DataType; use datafusion_common::DataFusionError; use datafusion_common::ScalarValue; use datafusion_common::cast::{as_float32_array, as_float64_array}; @@ -1022,6 +1053,35 @@ mod test { assert_eq!(floats, &expected); } + /// A scalar `decimal_places` takes the hoisted-factor `unary` path, which + /// computes over null slots as well. The nulls must survive into the output. + #[test] + fn test_round_f64_scalar_decimal_places_preserves_nulls() { + let value: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(125.2345), + None, + Some(-1.555), + None, + ])); + + let result = super::round_columnar( + &ColumnarValue::Array(value), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + 4, + &DataType::Float64, + ) + .expect("failed to initialize function round"); + let ColumnarValue::Array(result) = result else { + panic!("expected an array result"); + }; + let floats = + as_float64_array(&result).expect("failed to initialize function round"); + + let expected = Float64Array::from(vec![Some(125.23), None, Some(-1.56), None]); + + assert_eq!(floats, &expected); + } + #[test] fn test_round_f32_one_input() { let args: Vec = vec![