Define overrides for atan function to allow conditional sql call rewriting (1579/2015)

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

core/src/main/java/org/opensearch/sql/calcite/validate/PplConvertletTable.java

@@ -9,6 +9,7 @@
 import java.util.Map;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.SqlBasicCall;
 import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlFunction;
 import org.apache.calcite.sql.SqlOperator;
@@ -37,6 +38,14 @@ private PplConvertletTable() {
     registerOperator(SqlStdOperatorTable.LESS_THAN, ipConvertlet(PPLBuiltinOperators.LESS_IP));
     registerOperator(
         SqlStdOperatorTable.LESS_THAN_OR_EQUAL, ipConvertlet(PPLBuiltinOperators.LTE_IP));
+    // There is no implementation for PPLBuiltinOperators.ATAN. It needs to be replaced to
+    // SqlStdOperatorTable.ATAN when converted to RelNode
+    registerOperator(
+        PPLBuiltinOperators.ATAN,
+        (cx, call) -> {
+          ((SqlBasicCall) call).setOperator(SqlStdOperatorTable.ATAN);
+          return StandardConvertletTable.INSTANCE.convertCall(cx, call);
+        });
   }
 
   @Override

core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java

@@ -23,11 +23,19 @@
 import org.apache.calcite.linq4j.tree.Expression;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlBasicCall;
+import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeTransforms;
 import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
+import org.apache.calcite.sql.validate.SqlValidator;
 import org.apache.calcite.util.BuiltInMethod;
 import org.opensearch.sql.calcite.udf.udaf.FirstAggFunction;
 import org.opensearch.sql.calcite.udf.udaf.LastAggFunction;
@@ -477,6 +485,25 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
   public static final SqlFunction ENHANCED_COALESCE =
       new EnhancedCoalesceFunction().toUDF("COALESCE");
 
+  public static final SqlFunction ATAN =
+      new SqlFunction(
+          "ATAN",
+          SqlKind.OTHER_FUNCTION,
+          ReturnTypes.DOUBLE_NULLABLE,
+          null,
+          OperandTypes.NUMERIC_OPTIONAL_NUMERIC,
+          SqlFunctionCategory.NUMERIC) {
+        @Override
+        public SqlNode rewriteCall(SqlValidator validator, SqlCall call) {
+          SqlOperator op =
+              call.getOperandList().size() == 2
+                  ? SqlStdOperatorTable.ATAN2
+                  : SqlStdOperatorTable.ATAN;
+          ((SqlBasicCall) call).setOperator(op);
+          return call;
+        }
+      };
+
   /**
    * Returns the PPL specific operator table, creating it if necessary.
    *

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -784,7 +784,7 @@ void populate() {
       registerOperator(ABS, SqlStdOperatorTable.ABS);
       registerOperator(ACOS, SqlStdOperatorTable.ACOS);
       registerOperator(ASIN, SqlStdOperatorTable.ASIN);
-      registerOperator(ATAN, SqlStdOperatorTable.ATAN);
+      registerOperator(ATAN, PPLBuiltinOperators.ATAN);
       registerOperator(ATAN2, SqlStdOperatorTable.ATAN2);
       // TODO, workaround to support sequence CompositeOperandTypeChecker.
       registerOperator(

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteArrayFunctionIT.java

@@ -65,9 +65,7 @@ public void testArrayWithMix() {
                         TEST_INDEX_BANK)));
 
     verifyErrorMessageContains(
-        e,
-        "Cannot resolve function: ARRAY, arguments: [INTEGER,BOOLEAN], caused by: fail to create"
-            + " array with fixed type");
+        e, "Cannot infer return type for array; operand types: [INTEGER, BOOLEAN]");
   }
 
   @Test