datafusion_spark/function/math/
trigonometry.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 crate::function::error_utils::unsupported_data_type_exec_err;
19use arrow::array::{ArrayRef, AsArray};
20use arrow::datatypes::{DataType, Float64Type};
21use datafusion_common::utils::take_function_args;
22use datafusion_common::{Result, ScalarValue};
23use datafusion_expr::{
24    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
25};
26use std::sync::Arc;
27
28static CSC_FUNCTION_NAME: &str = "csc";
29
30/// <https://spark.apache.org/docs/latest/api/sql/index.html#csc>
31#[derive(Debug, PartialEq, Eq, Hash)]
32pub struct SparkCsc {
33    signature: Signature,
34}
35
36impl Default for SparkCsc {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl SparkCsc {
43    pub fn new() -> Self {
44        Self {
45            signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
46        }
47    }
48}
49
50impl ScalarUDFImpl for SparkCsc {
51    fn name(&self) -> &str {
52        CSC_FUNCTION_NAME
53    }
54
55    fn signature(&self) -> &Signature {
56        &self.signature
57    }
58
59    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
60        Ok(DataType::Float64)
61    }
62
63    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
64        let [arg] = take_function_args(self.name(), &args.args)?;
65        spark_csc(arg)
66    }
67}
68
69fn spark_csc(arg: &ColumnarValue) -> Result<ColumnarValue> {
70    match arg {
71        ColumnarValue::Scalar(ScalarValue::Float64(value)) => Ok(ColumnarValue::Scalar(
72            ScalarValue::Float64(value.map(|x| 1.0 / x.sin())),
73        )),
74        ColumnarValue::Array(array) => match array.data_type() {
75            DataType::Float64 => Ok(ColumnarValue::Array(Arc::new(
76                array
77                    .as_primitive::<Float64Type>()
78                    .unary::<_, Float64Type>(|x| 1.0 / x.sin()),
79            ) as ArrayRef)),
80            other => Err(unsupported_data_type_exec_err(
81                CSC_FUNCTION_NAME,
82                format!("{}", DataType::Float64).as_str(),
83                other,
84            )),
85        },
86        other => Err(unsupported_data_type_exec_err(
87            CSC_FUNCTION_NAME,
88            format!("{}", DataType::Float64).as_str(),
89            &other.data_type(),
90        )),
91    }
92}
93
94static SEC_FUNCTION_NAME: &str = "sec";
95
96/// <https://spark.apache.org/docs/latest/api/sql/index.html#sec>
97#[derive(Debug, PartialEq, Eq, Hash)]
98pub struct SparkSec {
99    signature: Signature,
100}
101
102impl Default for SparkSec {
103    fn default() -> Self {
104        Self::new()
105    }
106}
107
108impl SparkSec {
109    pub fn new() -> Self {
110        Self {
111            signature: Signature::exact(vec![DataType::Float64], Volatility::Immutable),
112        }
113    }
114}
115
116impl ScalarUDFImpl for SparkSec {
117    fn name(&self) -> &str {
118        SEC_FUNCTION_NAME
119    }
120
121    fn signature(&self) -> &Signature {
122        &self.signature
123    }
124
125    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
126        Ok(DataType::Float64)
127    }
128
129    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
130        let [arg] = take_function_args(self.name(), &args.args)?;
131        spark_sec(arg)
132    }
133}
134
135fn spark_sec(arg: &ColumnarValue) -> Result<ColumnarValue> {
136    match arg {
137        ColumnarValue::Scalar(ScalarValue::Float64(value)) => Ok(ColumnarValue::Scalar(
138            ScalarValue::Float64(value.map(|x| 1.0 / x.cos())),
139        )),
140        ColumnarValue::Array(array) => match array.data_type() {
141            DataType::Float64 => Ok(ColumnarValue::Array(Arc::new(
142                array
143                    .as_primitive::<Float64Type>()
144                    .unary::<_, Float64Type>(|x| 1.0 / x.cos()),
145            ) as ArrayRef)),
146            other => Err(unsupported_data_type_exec_err(
147                SEC_FUNCTION_NAME,
148                format!("{}", DataType::Float64).as_str(),
149                other,
150            )),
151        },
152        other => Err(unsupported_data_type_exec_err(
153            SEC_FUNCTION_NAME,
154            format!("{}", DataType::Float64).as_str(),
155            &other.data_type(),
156        )),
157    }
158}
datafusion_spark/function/math/trigonometry.rs

datafusion_spark/function/math/
trigonometry.rs
