perf: Optimize trunc scalar performance

datafusion/functions/benches/trunc.rs

@@ -32,12 +32,13 @@ 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 f32_array = Arc::new(create_primitive_array::<Float32Type>(size, 0.2));
         let f32_args = vec![ColumnarValue::Array(f32_array)];
         let arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
         let return_field = Field::new("f", DataType::Float32, true).into();
-        let config_options = Arc::new(ConfigOptions::default());
 
         c.bench_function(&format!("trunc f32 array: {size}"), |b| {
             b.iter(|| {
@@ -74,6 +75,51 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
         });
     }
+
+    // Scalar benchmarks - to measure optimized performance
+    let scalar_f64_args = vec![ColumnarValue::Scalar(
+        datafusion_common::ScalarValue::Float64(Some(std::f64::consts::PI)),
+    )];
+    let scalar_arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
+    let scalar_return_field = Field::new("f", DataType::Float64, false).into();
+
+    c.bench_function("trunc f64 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                trunc
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f64_args.clone(),
+                        arg_fields: scalar_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&scalar_return_field),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
+
+    let scalar_f32_args = vec![ColumnarValue::Scalar(
+        datafusion_common::ScalarValue::Float32(Some(std::f32::consts::PI)),
+    )];
+    let scalar_f32_arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
+    let scalar_f32_return_field = Field::new("f", DataType::Float32, false).into();
+
+    c.bench_function("trunc f32 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                trunc
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f32_args.clone(),
+                        arg_fields: scalar_f32_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&scalar_f32_return_field),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

datafusion/functions/src/math/trunc.rs

@@ -24,7 +24,7 @@ use arrow::array::{ArrayRef, AsArray, PrimitiveArray};
 use arrow::datatypes::DataType::{Float32, Float64};
 use arrow::datatypes::{DataType, Float32Type, Float64Type, Int64Type};
 use datafusion_common::ScalarValue::Int64;
-use datafusion_common::{Result, exec_err};
+use datafusion_common::{Result, ScalarValue, exec_err};
 use datafusion_expr::TypeSignature::Exact;
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
@@ -110,6 +110,60 @@ impl ScalarUDFImpl for TruncFunc {
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let value = &args.args[0];
+
+        // Scalar fast path for float types with scalar or default precision
+        if let ColumnarValue::Scalar(scalar) = value {
+            // Get precision: default 0 or from second scalar arg
+            let precision = if args.args.len() >= 2 {
+                match &args.args[1] {
+                    ColumnarValue::Scalar(Int64(Some(p))) => *p,
+                    ColumnarValue::Scalar(Int64(None)) => {
+                        // Return null with the same type as the input
+                        return match scalar {
+                            ScalarValue::Float32(_) => {
+                                Ok(ColumnarValue::Scalar(ScalarValue::Float32(None)))
+                            }
+                            _ => Ok(ColumnarValue::Scalar(ScalarValue::Float64(None))),
+                        };
+                    }
+                    _ => {
+                        // Precision is an array - fall through to array path
+                        return make_scalar_function(trunc, vec![])(&args.args);
+                    }
+                }
+            } else {
+                0 // default precision
+            };
+
+            match scalar {
+                ScalarValue::Float64(v) => {
+                    let result = v.map(|x| {
+                        if precision == 0 {
+                            x.trunc()
+                        } else {
+                            compute_truncate64(x, precision)
+                        }
+                    });
+                    return Ok(ColumnarValue::Scalar(ScalarValue::Float64(result)));
+                }
+                ScalarValue::Float32(v) => {
+                    let result = v.map(|x| {
+                        if precision == 0 {
+                            x.trunc()
+                        } else {
+                            compute_truncate32(x, precision)
+                        }
+                    });
+                    return Ok(ColumnarValue::Scalar(ScalarValue::Float32(result)));
+                }
+                ScalarValue::Null => {
+                    return Ok(ColumnarValue::Scalar(ScalarValue::Float64(None)));
+                }
+                _ => {}
+            }
+        }
+
         make_scalar_function(trunc, vec![])(&args.args)
     }