perf: optimize round in datafusion-functions

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,84 @@ fn round_integer_array(
     }
 }
 
-fn round_float<T>(value: T, decimal_places: i32) -> Result<T, ArrowError>
+/// 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<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(|| {
+    if let Some(arr) = round_float_fast::<PT>(value_array, decimal_places)? {
+        return Ok(arr);
+    }
+    let result = calculate_binary_math::<PT, Int32Type, PT, _>(
+        value_array.as_ref(),
+        decimal_places,
+        round_float::<PT::Native>,
+    )?;
+    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<PT>(
+    value_array: &ArrayRef,
+    decimal_places: &ColumnarValue,
+) -> Result<Option<ArrayRef>>
+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::<PT>();
+    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::<PT::Native>(*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<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)
 }