datafusion_spark/function/math/
atan2.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::{ArrayRef, AsArray, Float64Array};
21use arrow::compute::kernels::arity::binary;
22use arrow::datatypes::{DataType, Float64Type};
23use datafusion_common::Result;
24use datafusion_common::utils::take_function_args;
25use datafusion_expr::{
26    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
27};
28use datafusion_functions::utils::make_scalar_function;
29
30/// Spark-compatible `atan2` function.
31///
32/// <https://spark.apache.org/docs/latest/api/sql/index.html#atan2>
33///
34/// `atan2(exprY, exprX)` returns the angle in radians between the positive
35/// x-axis and the point given by the coordinates (exprX, exprY).
36#[derive(Debug, PartialEq, Eq, Hash)]
37pub struct SparkAtan2 {
38    signature: Signature,
39}
40
41impl Default for SparkAtan2 {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl SparkAtan2 {
48    pub fn new() -> Self {
49        Self {
50            // Spark only defines atan2 over doubles
51            signature: Signature::exact(
52                vec![DataType::Float64, DataType::Float64],
53                Volatility::Immutable,
54            ),
55        }
56    }
57}
58
59impl ScalarUDFImpl for SparkAtan2 {
60    fn name(&self) -> &str {
61        "atan2"
62    }
63
64    fn signature(&self) -> &Signature {
65        &self.signature
66    }
67
68    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
69        Ok(DataType::Float64)
70    }
71
72    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
73        make_scalar_function(spark_atan2, vec![])(&args.args)
74    }
75}
76
77fn spark_atan2(args: &[ArrayRef]) -> Result<ArrayRef> {
78    // Spark arg order is atan2(exprY, exprX); Rust computes y.atan2(x).
79    let [y, x] = take_function_args("atan2", args)?;
80    let y = y.as_primitive::<Float64Type>();
81    let x = x.as_primitive::<Float64Type>();
82    let result: Float64Array = binary(y, x, |y, x| y.atan2(x))?;
83    Ok(Arc::new(result))
84}
datafusion_spark/function/math/atan2.rs

datafusion_spark/function/math/
atan2.rs
