diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 2ffe93a0e567..eb120d1dadb1 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -182,6 +182,11 @@ harness = false name = "character_length" required-features = ["unicode_expressions"] +[[bench]] +harness = false +name = "cot" +required-features = ["math_expressions"] + [[bench]] harness = false name = "strpos" diff --git a/datafusion/functions/benches/cot.rs b/datafusion/functions/benches/cot.rs new file mode 100644 index 000000000000..e655d82dec91 --- /dev/null +++ b/datafusion/functions/benches/cot.rs @@ -0,0 +1,47 @@ +// 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. + +extern crate criterion; + +use arrow::{ + datatypes::{Float32Type, Float64Type}, + util::bench_util::create_primitive_array, +}; +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use datafusion_expr::ColumnarValue; +use datafusion_functions::math::cot; + +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let cot_fn = cot(); + for size in [1024, 4096, 8192] { + let f32_array = Arc::new(create_primitive_array::(size, 0.2)); + let f32_args = vec![ColumnarValue::Array(f32_array)]; + c.bench_function(&format!("cot f32 array: {}", size), |b| { + b.iter(|| black_box(cot_fn.invoke(&f32_args).unwrap())) + }); + let f64_array = Arc::new(create_primitive_array::(size, 0.2)); + let f64_args = vec![ColumnarValue::Array(f64_array)]; + c.bench_function(&format!("cot f64 array: {}", size), |b| { + b.iter(|| black_box(cot_fn.invoke(&f64_args).unwrap())) + }); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/math/cot.rs b/datafusion/functions/src/math/cot.rs index 66219960d9a2..f039767536fa 100644 --- a/datafusion/functions/src/math/cot.rs +++ b/datafusion/functions/src/math/cot.rs @@ -18,11 +18,11 @@ use std::any::Any; use std::sync::Arc; -use arrow::array::{ArrayRef, Float32Array, Float64Array}; -use arrow::datatypes::DataType; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType::{Float32, Float64}; +use arrow::datatypes::{DataType, Float32Type, Float64Type}; -use datafusion_common::{exec_err, DataFusionError, Result}; +use datafusion_common::{exec_err, Result}; use datafusion_expr::ColumnarValue; use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; @@ -85,18 +85,16 @@ impl ScalarUDFImpl for CotFunc { ///cot SQL function fn cot(args: &[ArrayRef]) -> Result { match args[0].data_type() { - Float64 => Ok(Arc::new(make_function_scalar_inputs!( - &args[0], - "x", - Float64Array, - { compute_cot64 } - )) as ArrayRef), - Float32 => Ok(Arc::new(make_function_scalar_inputs!( - &args[0], - "x", - Float32Array, - { compute_cot32 } - )) as ArrayRef), + Float64 => Ok(Arc::new( + args[0] + .as_primitive::() + .unary::<_, Float64Type>(|x: f64| compute_cot64(x)), + ) as ArrayRef), + Float32 => Ok(Arc::new( + args[0] + .as_primitive::() + .unary::<_, Float32Type>(|x: f32| compute_cot32(x)), + ) as ArrayRef), other => exec_err!("Unsupported data type {other:?} for function cot"), } }