datafusion_functions/math/
cot.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow::array::AsArray;
21use arrow::datatypes::DataType::{Float32, Float64};
22use arrow::datatypes::{DataType, Float32Type, Float64Type};
23
24use datafusion_common::utils::take_function_args;
25use datafusion_common::{Result, ScalarValue, internal_err};
26use datafusion_expr::{ColumnarValue, Documentation, ScalarFunctionArgs};
27use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
28use datafusion_macros::user_doc;
29
30#[user_doc(
31    doc_section(label = "Math Functions"),
32    description = "Returns the cotangent of a number.",
33    syntax_example = r#"cot(numeric_expression)"#,
34    sql_example = r#"```sql
35> SELECT cot(1);
36+---------+
37| cot(1)  |
38+---------+
39| 0.64209 |
40+---------+
41```"#,
42    standard_argument(name = "numeric_expression", prefix = "Numeric")
43)]
44#[derive(Debug, PartialEq, Eq, Hash)]
45pub struct CotFunc {
46    signature: Signature,
47}
48
49impl Default for CotFunc {
50    fn default() -> Self {
51        CotFunc::new()
52    }
53}
54
55impl CotFunc {
56    pub fn new() -> Self {
57        use DataType::*;
58        Self {
59            // math expressions expect 1 argument of type f64 or f32
60            // priority is given to f64 because e.g. `sqrt(1i32)` is in IR (real numbers) and thus we
61            // return the best approximation for it (in f64).
62            // We accept f32 because in this case it is clear that the best approximation
63            // will be as good as the number of digits in the number
64            signature: Signature::uniform(
65                1,
66                vec![Float64, Float32],
67                Volatility::Immutable,
68            ),
69        }
70    }
71}
72
73impl ScalarUDFImpl for CotFunc {
74    fn name(&self) -> &str {
75        "cot"
76    }
77
78    fn signature(&self) -> &Signature {
79        &self.signature
80    }
81
82    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
83        match arg_types[0] {
84            Float32 => Ok(Float32),
85            _ => Ok(Float64),
86        }
87    }
88
89    fn is_strict(&self) -> bool {
90        true
91    }
92
93    fn documentation(&self) -> Option<&Documentation> {
94        self.doc()
95    }
96
97    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
98        let return_field = args.return_field;
99        let [arg] = take_function_args(self.name(), args.args)?;
100
101        match arg {
102            ColumnarValue::Scalar(scalar) => {
103                if scalar.is_null() {
104                    return ColumnarValue::Scalar(ScalarValue::Null)
105                        .cast_to(return_field.data_type(), None);
106                }
107
108                match scalar {
109                    ScalarValue::Float64(Some(v)) => Ok(ColumnarValue::Scalar(
110                        ScalarValue::Float64(Some(compute_cot64(v))),
111                    )),
112                    ScalarValue::Float32(Some(v)) => Ok(ColumnarValue::Scalar(
113                        ScalarValue::Float32(Some(compute_cot32(v))),
114                    )),
115                    _ => {
116                        internal_err!(
117                            "Unexpected scalar type for cot: {:?}",
118                            scalar.data_type()
119                        )
120                    }
121                }
122            }
123            ColumnarValue::Array(array) => match array.data_type() {
124                Float64 => Ok(ColumnarValue::Array(Arc::new(
125                    array
126                        .as_primitive::<Float64Type>()
127                        .unary::<_, Float64Type>(compute_cot64),
128                ))),
129                Float32 => Ok(ColumnarValue::Array(Arc::new(
130                    array
131                        .as_primitive::<Float32Type>()
132                        .unary::<_, Float32Type>(compute_cot32),
133                ))),
134                other => {
135                    internal_err!("Unexpected data type {other:?} for function cot")
136                }
137            },
138        }
139    }
140}
141
142fn compute_cot32(x: f32) -> f32 {
143    let a = f32::tan(x);
144    1.0 / a
145}
146
147fn compute_cot64(x: f64) -> f64 {
148    let a = f64::tan(x);
149    1.0 / a
150}
151
152#[cfg(test)]
153mod test {
154    use std::sync::Arc;
155
156    use arrow::array::{ArrayRef, Float32Array, Float64Array};
157    use arrow::datatypes::{DataType, Field};
158    use datafusion_common::ScalarValue;
159    use datafusion_common::cast::{as_float32_array, as_float64_array};
160    use datafusion_common::config::ConfigOptions;
161    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
162
163    use crate::math::cot::CotFunc;
164
165    #[test]
166    fn test_cot_f32() {
167        let array = Arc::new(Float32Array::from(vec![12.1, 30.0, 90.0, -30.0]));
168        let arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
169        let args = ScalarFunctionArgs {
170            args: vec![ColumnarValue::Array(Arc::clone(&array) as ArrayRef)],
171            arg_fields,
172            number_rows: array.len(),
173            return_field: Field::new("f", DataType::Float32, true).into(),
174            config_options: Arc::new(ConfigOptions::default()),
175        };
176        let result = CotFunc::new()
177            .invoke_with_args(args)
178            .expect("failed to initialize function cot");
179
180        match result {
181            ColumnarValue::Array(arr) => {
182                let floats = as_float32_array(&arr)
183                    .expect("failed to convert result to a Float32Array");
184
185                let expected = Float32Array::from(vec![
186                    -1.986_460_4,
187                    -0.156_119_96,
188                    -0.501_202_8,
189                    0.156_119_96,
190                ]);
191
192                let eps = 1e-6;
193                assert_eq!(floats.len(), 4);
194                assert!((floats.value(0) - expected.value(0)).abs() < eps);
195                assert!((floats.value(1) - expected.value(1)).abs() < eps);
196                assert!((floats.value(2) - expected.value(2)).abs() < eps);
197                assert!((floats.value(3) - expected.value(3)).abs() < eps);
198            }
199            ColumnarValue::Scalar(_) => {
200                panic!("Expected an array value")
201            }
202        }
203    }
204
205    #[test]
206    fn test_cot_f64() {
207        let array = Arc::new(Float64Array::from(vec![12.1, 30.0, 90.0, -30.0]));
208        let arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
209        let args = ScalarFunctionArgs {
210            args: vec![ColumnarValue::Array(Arc::clone(&array) as ArrayRef)],
211            arg_fields,
212            number_rows: array.len(),
213            return_field: Field::new("f", DataType::Float64, true).into(),
214            config_options: Arc::new(ConfigOptions::default()),
215        };
216        let result = CotFunc::new()
217            .invoke_with_args(args)
218            .expect("failed to initialize function cot");
219
220        match result {
221            ColumnarValue::Array(arr) => {
222                let floats = as_float64_array(&arr)
223                    .expect("failed to convert result to a Float64Array");
224
225                let expected = Float64Array::from(vec![
226                    -1.986_458_685_881_4,
227                    -0.156_119_952_161_6,
228                    -0.501_202_783_380_1,
229                    0.156_119_952_161_6,
230                ]);
231
232                let eps = 1e-12;
233                assert_eq!(floats.len(), 4);
234                assert!((floats.value(0) - expected.value(0)).abs() < eps);
235                assert!((floats.value(1) - expected.value(1)).abs() < eps);
236                assert!((floats.value(2) - expected.value(2)).abs() < eps);
237                assert!((floats.value(3) - expected.value(3)).abs() < eps);
238            }
239            ColumnarValue::Scalar(_) => {
240                panic!("Expected an array value")
241            }
242        }
243    }
244
245    #[test]
246    fn test_cot_scalar_f64() {
247        let arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
248        let args = ScalarFunctionArgs {
249            args: vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0)))],
250            arg_fields,
251            number_rows: 1,
252            return_field: Field::new("f", DataType::Float64, false).into(),
253            config_options: Arc::new(ConfigOptions::default()),
254        };
255        let result = CotFunc::new()
256            .invoke_with_args(args)
257            .expect("cot scalar should succeed");
258
259        match result {
260            ColumnarValue::Scalar(ScalarValue::Float64(Some(v))) => {
261                // cot(1.0) = 1/tan(1.0) ≈ 0.6420926159343306
262                let expected = 1.0_f64 / 1.0_f64.tan();
263                assert!((v - expected).abs() < 1e-12);
264            }
265            _ => panic!("Expected Float64 scalar"),
266        }
267    }
268
269    #[test]
270    fn test_cot_scalar_f32() {
271        let arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
272        let args = ScalarFunctionArgs {
273            args: vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(1.0)))],
274            arg_fields,
275            number_rows: 1,
276            return_field: Field::new("f", DataType::Float32, false).into(),
277            config_options: Arc::new(ConfigOptions::default()),
278        };
279        let result = CotFunc::new()
280            .invoke_with_args(args)
281            .expect("cot scalar should succeed");
282
283        match result {
284            ColumnarValue::Scalar(ScalarValue::Float32(Some(v))) => {
285                let expected = 1.0_f32 / 1.0_f32.tan();
286                assert!((v - expected).abs() < 1e-6);
287            }
288            _ => panic!("Expected Float32 scalar"),
289        }
290    }
291
292    #[test]
293    fn test_cot_scalar_null() {
294        let arg_fields = vec![Field::new("a", DataType::Float64, true).into()];
295        let args = ScalarFunctionArgs {
296            args: vec![ColumnarValue::Scalar(ScalarValue::Float64(None))],
297            arg_fields,
298            number_rows: 1,
299            return_field: Field::new("f", DataType::Float64, true).into(),
300            config_options: Arc::new(ConfigOptions::default()),
301        };
302        let result = CotFunc::new()
303            .invoke_with_args(args)
304            .expect("cot null should succeed");
305
306        match result {
307            ColumnarValue::Scalar(scalar) => {
308                assert!(scalar.is_null());
309            }
310            _ => panic!("Expected scalar result"),
311        }
312    }
313
314    #[test]
315    fn test_cot_scalar_zero() {
316        let arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
317        let args = ScalarFunctionArgs {
318            args: vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0)))],
319            arg_fields,
320            number_rows: 1,
321            return_field: Field::new("f", DataType::Float64, false).into(),
322            config_options: Arc::new(ConfigOptions::default()),
323        };
324        let result = CotFunc::new()
325            .invoke_with_args(args)
326            .expect("cot zero should succeed");
327
328        match result {
329            ColumnarValue::Scalar(ScalarValue::Float64(Some(v))) => {
330                // cot(0) = 1/tan(0) = infinity
331                assert!(v.is_infinite());
332            }
333            _ => panic!("Expected Float64 scalar"),
334        }
335    }
336
337    #[test]
338    fn test_cot_scalar_pi() {
339        let arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
340        let args = ScalarFunctionArgs {
341            args: vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(
342                std::f64::consts::PI,
343            )))],
344            arg_fields,
345            number_rows: 1,
346            return_field: Field::new("f", DataType::Float64, false).into(),
347            config_options: Arc::new(ConfigOptions::default()),
348        };
349        let result = CotFunc::new()
350            .invoke_with_args(args)
351            .expect("cot pi should succeed");
352
353        match result {
354            ColumnarValue::Scalar(ScalarValue::Float64(Some(v))) => {
355                // cot(PI) = 1/tan(PI) - very large negative number due to floating point
356                let expected = 1.0_f64 / std::f64::consts::PI.tan();
357                assert!((v - expected).abs() < 1e-6);
358            }
359            _ => panic!("Expected Float64 scalar"),
360        }
361    }
362}
datafusion_functions/math/cot.rs

datafusion_functions/math/
cot.rs
