perf: Optimize scalar performance for cot

datafusion/functions/benches/cot.rs

@@ -27,11 +27,15 @@ use datafusion_functions::math::cot;
 use std::hint::black_box;
 
 use arrow::datatypes::{DataType, Field};
+use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
 use std::sync::Arc;
 
 fn criterion_benchmark(c: &mut Criterion) {
     let cot_fn = cot();
+    let config_options = Arc::new(ConfigOptions::default());
+
+    // Array benchmarks - run for different sizes
     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)];
@@ -42,7 +46,6 @@ fn criterion_benchmark(c: &mut Criterion) {
                 Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
             })
             .collect::<Vec<_>>();
-        let config_options = Arc::new(ConfigOptions::default());
 
         c.bench_function(&format!("cot f32 array: {size}"), |b| {
             b.iter(|| {
@@ -59,6 +62,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 )
             })
         });
+
         let f64_array = Arc::new(create_primitive_array::<Float64Type>(size, 0.2));
         let f64_args = vec![ColumnarValue::Array(f64_array)];
         let arg_fields = f64_args
@@ -86,6 +90,47 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
         });
     }
+
+    // Scalar benchmarks - run only once since size doesn't affect scalar performance
+    let scalar_f32_args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(1.0)))];
+    let scalar_f32_arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
+    let return_field_f32 = Field::new("f", DataType::Float32, false).into();
+
+    c.bench_function("cot f32 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                cot_fn
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f32_args.clone(),
+                        arg_fields: scalar_f32_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&return_field_f32),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
+
+    let scalar_f64_args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0)))];
+    let scalar_f64_arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
+    let return_field_f64 = Field::new("f", DataType::Float64, false).into();
+
+    c.bench_function("cot f64 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                cot_fn
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f64_args.clone(),
+                        arg_fields: scalar_f64_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&return_field_f64),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);