[Function](exec) upport atan2 math function (#30672)

Co-authored-by: Rohit Satardekar <[email protected]>
apache · Feb 4, 2024 · d1923cd · d1923cd
1 parent 39c5720
commit d1923cd
Show file tree

Hide file tree

Showing 10 changed files with 229 additions and 1 deletion.
diff --git a/be/src/vec/functions/math.cpp b/be/src/vec/functions/math.cpp
@@ -74,6 +74,21 @@ struct AtanName {
 };
 using FunctionAtan = FunctionMathUnary<UnaryFunctionPlain<AtanName, std::atan>>;
 
+template <typename A, typename B>
+struct Atan2Impl {
+    using ResultType = double;
+    static const constexpr bool allow_decimal = false;
+
+    template <typename type>
+    static inline double apply(A a, B b) {
+        return std::atan2((double)a, (double)b);
+    }
+};
+struct Atan2Name {
+    static constexpr auto name = "atan2";
+};
+using FunctionAtan2 = FunctionBinaryArithmetic<Atan2Impl, Atan2Name, false>;
+
 struct CosName {
     static constexpr auto name = "cos";
 };
@@ -397,6 +412,7 @@ void register_function_math(SimpleFunctionFactory& factory) {
     factory.register_function<FunctionAcos>();
     factory.register_function<FunctionAsin>();
     factory.register_function<FunctionAtan>();
+    factory.register_function<FunctionAtan2>();
     factory.register_function<FunctionCos>();
     factory.register_function<FunctionCosh>();
     factory.register_alias("ceil", "dceil");

diff --git a/docs/en/docs/sql-manual/sql-functions/numeric-functions/atan2.md b/docs/en/docs/sql-manual/sql-functions/numeric-functions/atan2.md
@@ -0,0 +1,52 @@
+---
+{
+    "title": "ATAN2",
+    "language": "en"
+}
+---
+
+<!-- 
+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.
+-->
+
+## atan2
+
+### description
+#### Syntax
+
+`DOUBLE atan2(DOUBLE y, DOUBLE x)`
+Returns the arc tangent of 'y' / 'x'.
+
+### example
+
+```
+mysql> select atan2(0.1, 0.2);
++---------------------+
+| atan2(0.1, 0.2)     |
++---------------------+
+| 0.46364760900080609 |
++---------------------+
+
+mysql> select atan2(1.0, 1.0);
++---------------------+
+| atan2(1.0, 1.0)     |
++---------------------+
+| 0.78539816339744828 |
++---------------------+
+```
+
+### keywords
+	ATAN2
diff --git a/docs/sidebars.json b/docs/sidebars.json
@@ -718,6 +718,7 @@
                                 "sql-manual/sql-functions/numeric-functions/asin",
                                 "sql-manual/sql-functions/numeric-functions/acos",
                                 "sql-manual/sql-functions/numeric-functions/atan",
+                                "sql-manual/sql-functions/numeric-functions/atan2",
                                 "sql-manual/sql-functions/numeric-functions/e",
                                 "sql-manual/sql-functions/numeric-functions/pi",
                                 "sql-manual/sql-functions/numeric-functions/exp",

diff --git a/docs/zh-CN/docs/sql-manual/sql-functions/numeric-functions/atan2.md b/docs/zh-CN/docs/sql-manual/sql-functions/numeric-functions/atan2.md
@@ -0,0 +1,52 @@
+---
+{
+    "title": "ATAN2",
+    "language": "zh-CN"
+}
+---
+
+<!-- 
+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.
+-->
+
+## atan2
+
+### description
+#### Syntax
+
+`DOUBLE atan2(DOUBLE y, DOUBLE x)`
+返回 'y' / 'x' 的反正切.
+
+### example
+
+```
+mysql> select atan2(0.1, 0.2);
++---------------------+
+| atan2(0.1, 0.2)     |
++---------------------+
+| 0.46364760900080609 |
++---------------------+
+
+mysql> select atan2(1.0, 1.0);
++---------------------+
+| atan2(1.0, 1.0)     |
++---------------------+
+| 0.78539816339744828 |
++---------------------+
+```
+
+### keywords
+	ATAN2
diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java
@@ -72,6 +72,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Ascii;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Asin;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan2;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Bin;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.BitCount;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.BitLength;
@@ -497,6 +498,7 @@ public class BuiltinScalarFunctions implements FunctionHelper {
             scalar(Ascii.class, "ascii"),
             scalar(Asin.class, "asin"),
             scalar(Atan.class, "atan"),
+            scalar(Atan2.class, "atan2"),
             scalar(Bin.class, "bin"),
             scalar(BitCount.class, "bit_count"),
             scalar(BitLength.class, "bit_length"),

diff --git a/...core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Atan2.java b/...core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Atan2.java
@@ -0,0 +1,68 @@
+// 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.
+
+package org.apache.doris.nereids.trees.expressions.functions.scalar;
+
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
+import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
+import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.types.DoubleType;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+/**
+ * ScalarFunction 'atan2'. This class is generated by GenerateFunction.
+ */
+public class Atan2 extends ScalarFunction
+        implements UnaryExpression, ExplicitlyCastableSignature, PropagateNullable {
+
+    public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
+            FunctionSignature.ret(DoubleType.INSTANCE).args(DoubleType.INSTANCE, DoubleType.INSTANCE)
+    );
+
+    /**
+     * constructor with 2 argument.
+     */
+    public Atan2(Expression arg0, Expression arg1) {
+        super("atan2", arg0, arg1);
+    }
+
+    /**
+     * withChildren.
+     */
+    @Override
+    public Atan2 withChildren(List<Expression> children) {
+        Preconditions.checkArgument(children.size() == 2);
+        return new Atan2(children.get(0), children.get(1));
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+
+    @Override
+    public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitAtan2(this, context);
+    }
+}
diff --git a/...c/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/...c/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java
@@ -72,6 +72,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Ascii;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Asin;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan2;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Bin;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.BitCount;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.BitLength;
@@ -633,6 +634,10 @@ default R visitAtan(Atan atan, C context) {
         return visitScalarFunction(atan, context);
     }
 
+    default R visitAtan2(Atan2 atan2, C context) {
+        return visitScalarFunction(atan2, context);
+    }
+
     default R visitBin(Bin bin, C context) {
         return visitScalarFunction(bin, context);
     }

diff --git a/gensrc/script/doris_builtins_functions.py b/gensrc/script/doris_builtins_functions.py
@@ -1263,6 +1263,7 @@
         [['abs'], 'DECIMAL128', ['DECIMAL128'], ''],
         [['acos'], 'DOUBLE', ['DOUBLE'], ''],
         [['atan'], 'DOUBLE', ['DOUBLE'], ''],
+        [['atan2'], 'DOUBLE', ['DOUBLE', 'DOUBLE'], ''],
         [['asin'], 'DOUBLE', ['DOUBLE'], ''],
 
         [['bin'], 'VARCHAR', ['BIGINT'], ''],

diff --git a/regression-test/data/nereids_function_p0/scalar_function/A.out b/regression-test/data/nereids_function_p0/scalar_function/A.out
@@ -434,3 +434,32 @@ nan
 0.8329812666744317
 0.8760580505981934
 
+-- !sql_atan2_Double --
+\N
+1.4711276743037345
+1.3734007669450159
+1.2793395323170296
+1.1902899496825317
+1.1071487177940904
+1.0303768265243125
+0.960070362405688
+0.89605538457134393
+0.83798122500839
+0.78539816339744828
+0.73781506012046483
+0.69473827619670314
+
+-- !sql_atan2_Double_notnull --
+1.4711276743037345
+1.3734007669450159
+1.2793395323170296
+1.1902899496825317
+1.1071487177940904
+1.0303768265243125
+0.960070362405688
+0.89605538457134393
+0.83798122500839
+0.78539816339744828
+0.73781506012046483
+0.69473827619670314
+
diff --git a/regression-test/suites/nereids_function_p0/scalar_function/A.groovy b/regression-test/suites/nereids_function_p0/scalar_function/A.groovy
@@ -73,4 +73,6 @@ suite("nereids_scalar_fn_A") {
 	qt_sql_asin_Double_notnull "select asin(kdbl) from fn_test_not_nullable order by kdbl"
 	qt_sql_atan_Double "select atan(kdbl) from fn_test order by kdbl"
 	qt_sql_atan_Double_notnull "select atan(kdbl) from fn_test_not_nullable order by kdbl"
-}
+	qt_sql_atan2_Double "select atan2(kdbl, kdbl*kdbl) from fn_test order by kdbl"
+	qt_sql_atan2_Double_notnull "select atan2(kdbl, kdbl*kdbl) from fn_test_not_nullable order by kdbl"
+}