perf: optimize round expression

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"

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::<Float64Type>(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::<Float32Type>(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);

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::<Float64Type, Int32Type, Float64Type, _>(
-                value_array.as_ref(),
-                decimal_places,
-                round_float::<f64>,
-            )?;
-            result as _
-        }
-        (Float32, _) => {
-            let result = calculate_binary_math::<Float32Type, Int32Type, Float32Type, _>(
-                value_array.as_ref(),
-                decimal_places,
-                round_float::<f32>,
-            )?;
-            result as _
-        }
+        (Float64, _) => round_float_column::<Float64Type>(&value_array, decimal_places)?,
+        (Float32, _) => round_float_column::<Float32Type>(&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<T>(value: T, decimal_places: i32) -> Result<T, ArrowError>
+/// 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<PT>(
+    value_array: &ArrayRef,
+    decimal_places: &ColumnarValue,
+) -> Result<ArrayRef>
 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::<PT::Native>(*decimal_places)?;
+        let result = value_array
+            .as_primitive::<PT>()
+            .unary::<_, PT>(|value| (value * factor).round() / factor);
+        return Ok(Arc::new(result) as ArrayRef);
+    }
+
+    let result = calculate_binary_math::<PT, Int32Type, PT, _>(
+        value_array.as_ref(),
+        decimal_places,
+        round_float::<PT::Native>,
+    )?;
+    Ok(result as _)
+}
+
+/// Computes the power-of-ten scaling factor used to round to `decimal_places`.
+fn round_factor<T: num_traits::Float>(decimal_places: i32) -> Result<T, ArrowError> {
+    T::from(10_f64.powi(decimal_places)).ok_or_else(|| {
         ArrowError::ComputeError(format!(
             "Invalid value for decimal places: {decimal_places}"
         ))
-    })?;
+    })
+}
+
+fn round_float<T>(value: T, decimal_places: i32) -> Result<T, ArrowError>
+where
+    T: num_traits::Float,
+{
+    let factor = round_factor::<T>(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<ArrayRef> = vec![