perf: optimize trunc for scalar precision case (10x faster) (#23593)

## Which issue does this PR close?

N/A

## Rationale for this change

Optimize existing expression.

## What changes are included in this PR?

Added an array fast path for trunc(value, scalar precision) that hoists
10^precision out of the per-element loop and uses a unary kernel instead
of broadcasting the precision and recomputing powi per element.

## Are these changes tested?

Existing tests.

Benchmark (criterion):

- trunc f32 precision array_ 1024: 87.531% faster (base 1284ns -> cand
160ns)
- trunc f64 precision array_ 1024: 81.29% faster (base 1313ns -> cand
245ns)
- trunc f32 precision array_ 4096: 91.021% faster (base 4542ns -> cand
407ns)
- trunc f64 precision array_ 4096: 82.983% faster (base 4560ns -> cand
776ns)
- trunc f32 precision array_ 8192: 92.235% faster (base 9475ns -> cand
735ns)
- trunc f64 precision array_ 8192: 81.919% faster (base 10319ns -> cand
1865ns)

Full criterion output:

```text
trunc f64 precision array: 1024
                        time:   [245.17 ns 245.40 ns 245.70 ns]
                        change: [−81.321% −81.290% −81.259%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
  2 (2.00%) high mild
  1 (1.00%) high severe

trunc f32 precision array: 1024
                        time:   [159.81 ns 160.13 ns 160.65 ns]
                        change: [−87.558% −87.531% −87.495%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 15 outliers among 100 measurements (15.00%)
  6 (6.00%) high mild
  9 (9.00%) high severe

trunc f64 precision array: 4096
                        time:   [764.34 ns 768.58 ns 773.05 ns]
                        change: [−83.091% −82.983% −82.845%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 8 outliers among 100 measurements (8.00%)
  2 (2.00%) high mild
  6 (6.00%) high severe

trunc f32 precision array: 4096
                        time:   [404.97 ns 405.35 ns 405.93 ns]
                        change: [−91.049% −91.021% −90.991%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 13 outliers among 100 measurements (13.00%)
  2 (2.00%) high mild
  11 (11.00%) high severe

trunc f64 precision array: 8192
                        time:   [1.8676 µs 1.8741 µs 1.8813 µs]
                        change: [−81.960% −81.919% −81.874%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 100 measurements (1.00%)
  1 (1.00%) high mild

trunc f32 precision array: 8192
                        time:   [729.32 ns 731.77 ns 734.41 ns]
                        change: [−92.264% −92.235% −92.203%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  6 (6.00%) high mild
```


## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: andygrove

datafusion/functions/Cargo.toml

@@ -301,6 +301,11 @@ harness = false
 name = "trunc"
 required-features = ["math_expressions"]
 
+[[bench]]
+harness = false
+name = "trunc_precision"
+required-features = ["math_expressions"]
+
 [[bench]]
 harness = false
 name = "initcap"

datafusion/functions/benches/trunc_precision.rs

@@ -0,0 +1,91 @@
+// 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.
+
+//! Benchmarks the `trunc(value, precision)` array path where `precision` is a
+//! constant (scalar) argument.
+
+use arrow::datatypes::{DataType, Field, Float32Type, Float64Type};
+use arrow::util::bench_util::create_primitive_array;
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
+use datafusion_functions::math::trunc;
+use std::hint::black_box;
+use std::sync::Arc;
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let trunc = trunc();
+    let config_options = Arc::new(ConfigOptions::default());
+
+    for size in [1024, 4096, 8192] {
+        let f64_array = Arc::new(create_primitive_array::<Float64Type>(size, 0.2));
+        let f64_args = vec![
+            ColumnarValue::Array(f64_array),
+            ColumnarValue::Scalar(ScalarValue::Int64(Some(3))),
+        ];
+        let arg_fields = vec![
+            Field::new("a", DataType::Float64, true).into(),
+            Field::new("p", DataType::Int64, false).into(),
+        ];
+        let return_field = Field::new("f", DataType::Float64, true).into();
+        c.bench_function(&format!("trunc f64 precision array: {size}"), |b| {
+            b.iter(|| {
+                black_box(
+                    trunc
+                        .invoke_with_args(ScalarFunctionArgs {
+                            args: f64_args.clone(),
+                            arg_fields: arg_fields.clone(),
+                            number_rows: size,
+                            return_field: Arc::clone(&return_field),
+                            config_options: Arc::clone(&config_options),
+                        })
+                        .unwrap(),
+                )
+            })
+        });
+
+        let f32_array = Arc::new(create_primitive_array::<Float32Type>(size, 0.2));
+        let f32_args = vec![
+            ColumnarValue::Array(f32_array),
+            ColumnarValue::Scalar(ScalarValue::Int64(Some(3))),
+        ];
+        let arg_fields = vec![
+            Field::new("a", DataType::Float32, true).into(),
+            Field::new("p", DataType::Int64, false).into(),
+        ];
+        let return_field = Field::new("f", DataType::Float32, true).into();
+        c.bench_function(&format!("trunc f32 precision array: {size}"), |b| {
+            b.iter(|| {
+                black_box(
+                    trunc
+                        .invoke_with_args(ScalarFunctionArgs {
+                            args: f32_args.clone(),
+                            arg_fields: arg_fields.clone(),
+                            number_rows: size,
+                            return_field: Arc::clone(&return_field),
+                            config_options: Arc::clone(&config_options),
+                        })
+                        .unwrap(),
+                )
+            })
+        });
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/functions/src/math/trunc.rs

@@ -25,8 +25,8 @@ use arrow::datatypes::DataType::{
     Decimal32, Decimal64, Decimal128, Decimal256, Float32, Float64,
 };
 use arrow::datatypes::{
-    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
-    Float32Type, Float64Type, Int64Type,
+    ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, Decimal128Type,
+    Decimal256Type, DecimalType, Float32Type, Float64Type, Int64Type,
 };
 use datafusion_common::ScalarValue::Int64;
 use datafusion_common::types::{
@@ -40,7 +40,7 @@ use datafusion_expr::{
 };
 use datafusion_expr_common::signature::{Coercion, TypeSignature, TypeSignatureClass};
 use datafusion_macros::user_doc;
-use num_traits::{One, Zero, pow};
+use num_traits::{Float, NumCast, One, Zero, pow};
 
 #[user_doc(
     doc_section(label = "Math Functions"),
@@ -158,6 +158,12 @@ impl ScalarUDFImpl for TruncFunc {
             }
         };
 
+        // Whether an explicit precision argument was supplied. The array fast
+        // paths below must only apply to the two-argument form: single-argument
+        // `trunc(x)` uses a different zero handling (mapping `-0.0` to `0.0`)
+        // that must be preserved.
+        let has_precision_arg = args.args.len() == 2;
+
         // Scalar fast path using tuple matching for (value, precision)
         match (&args.args[0], precision) {
             // Null cases
@@ -231,6 +237,25 @@ impl ScalarUDFImpl for TruncFunc {
                 *lscale,
             ))),
 
+            // Array value with a constant (scalar) precision: hoist the power
+            // of ten out of the per-element loop instead of broadcasting the
+            // scalar into a full precision array and recomputing `10^p` for
+            // every element (see `truncate_float_array`).
+            (ColumnarValue::Array(arr), Some(p))
+                if has_precision_arg && arr.data_type() == &Float64 =>
+            {
+                Ok(ColumnarValue::Array(truncate_float_array::<Float64Type>(
+                    arr, p,
+                )))
+            }
+            (ColumnarValue::Array(arr), Some(p))
+                if has_precision_arg && arr.data_type() == &Float32 =>
+            {
+                Ok(ColumnarValue::Array(truncate_float_array::<Float32Type>(
+                    arr, p,
+                )))
+            }
+
             // Array path for everything else
             _ => make_scalar_function(trunc, vec![])(&args.args),
         }
@@ -375,14 +400,35 @@ fn trunc(args: &[ArrayRef]) -> Result<ArrayRef> {
     }
 }
 
-fn compute_truncate32(x: f32, y: i64) -> f32 {
-    let factor = 10.0_f32.powi(y as i32);
+/// Truncates `x` using a pre-computed `factor` of `10^precision`. Taking the
+/// factor as an argument lets callers hoist `10^precision` out of a per-element
+/// loop when the precision is constant.
+fn truncate_with_factor<F: Float>(x: F, factor: F) -> F {
     (x * factor).trunc() / factor
 }
 
+/// Truncates every element of a float array to `precision` decimal places,
+/// computing the `10^precision` factor once and reusing it for every element.
+fn truncate_float_array<T>(arr: &ArrayRef, precision: i64) -> ArrayRef
+where
+    T: ArrowPrimitiveType,
+    T::Native: Float,
+{
+    let factor = <T::Native as NumCast>::from(10.0_f64)
+        .unwrap()
+        .powi(precision as i32);
+    Arc::new(
+        arr.as_primitive::<T>()
+            .unary::<_, T>(|x| truncate_with_factor(x, factor)),
+    )
+}
+
+fn compute_truncate32(x: f32, y: i64) -> f32 {
+    truncate_with_factor(x, 10.0_f32.powi(y as i32))
+}
+
 fn compute_truncate64(x: f64, y: i64) -> f64 {
-    let factor = 10.0_f64.powi(y as i32);
-    (x * factor).trunc() / factor
+    truncate_with_factor(x, 10.0_f64.powi(y as i32))
 }
 
 /// Truncates a decimal value to `truncate_precision` fractional digits.