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

datafusion_spark/function/math/
trigonometry.rs
