Add interface registerOperator(BuiltinFunctionName, SqlOperator, PPLTypeChecker) to allow registering an operator with a designated type chekcer

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

Author: yuancu

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

@@ -564,73 +564,6 @@ private void compulsoryCast(
     return null;
   }
 
-  /**
-   * Get a string representation of the argument types expressed in ExprType for error messages.
-   *
-   * @param argTypes the list of argument types as {@link RelDataType}
-   * @return a string in the format [type1,type2,...] representing the argument types
-   */
-  private static String getActualSignature(List<RelDataType> argTypes) {
-    return "["
-        + argTypes.stream()
-            .map(OpenSearchTypeFactory::convertRelDataTypeToExprType)
-            .map(Objects::toString)
-            .collect(Collectors.joining(","))
-        + "]";
-  }
-
-  /**
-   * Wraps a {@link SqlOperandTypeChecker} into a {@link PPLTypeChecker} for use in function
-   * signature validation.
-   *
-   * @param typeChecker the original SQL operand type checker
-   * @param functionName the name of the function for error reporting
-   * @param isUserDefinedFunction true if the function is user-defined, false otherwise
-   * @return a {@link PPLTypeChecker} that delegates to the provided {@code typeChecker}
-   */
-  private static PPLTypeChecker wrapSqlOperandTypeChecker(
-      SqlOperandTypeChecker typeChecker, String functionName, boolean isUserDefinedFunction) {
-    PPLTypeChecker pplTypeChecker;
-    // Only the composite operand type checker for UDFs are concerned here.
-    if (isUserDefinedFunction
-        && typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
-      // UDFs implement their own composite type checkers, which always use OR logic for
-      // argument
-      // types. Verifying the composition type would require accessing a protected field in
-      // CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
-      // be skipped, so we avoid checking the composition type here.
-      pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, false);
-    } else if (typeChecker instanceof ImplicitCastOperandTypeChecker implicitCastTypeChecker) {
-      pplTypeChecker = PPLTypeChecker.wrapFamily(implicitCastTypeChecker);
-    } else if (typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
-      // If compositeTypeChecker contains operand checkers other than family type checkers or
-      // other than OR compositions, the function with be registered with a null type checker,
-      // which means the function will not be type checked.
-      try {
-        pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, true);
-      } catch (IllegalArgumentException | UnsupportedOperationException e) {
-        logger.debug(
-            String.format(
-                "Failed to create composite type checker for operator: %s. Will skip its type"
-                    + " checking",
-                functionName),
-            e);
-        pplTypeChecker = null;
-      }
-    } else if (typeChecker instanceof SameOperandTypeChecker comparableTypeChecker) {
-      // Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
-      // SameOperandTypeCheckers like COALESCE, IFNULL, etc.
-      pplTypeChecker = PPLTypeChecker.wrapComparable(comparableTypeChecker);
-    } else if (typeChecker instanceof UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
-      pplTypeChecker = PPLTypeChecker.wrapUDT(udtOperandMetadata.allowedParamTypes());
-    } else {
-      logger.info(
-          "Cannot create type checker for function: {}. Will skip its type checking", functionName);
-      pplTypeChecker = null;
-    }
-    return pplTypeChecker;
-  }
-
   @SuppressWarnings({"UnusedReturnValue", "SameParameterValue"})
   private abstract static class AbstractBuilder {
 
@@ -650,7 +583,7 @@ abstract void register(
      * @param operators the operators to associate with this function name, tried in sequence until
      *     one matches the argument types during resolution
      */
-    public void registerOperator(BuiltinFunctionName functionName, SqlOperator... operators) {
+    protected void registerOperator(BuiltinFunctionName functionName, SqlOperator... operators) {
       for (SqlOperator operator : operators) {
         SqlOperandTypeChecker typeChecker;
         if (operator instanceof SqlUserDefinedFunction udfOperator) {
@@ -662,18 +595,24 @@ public void registerOperator(BuiltinFunctionName functionName, SqlOperator... op
         PPLTypeChecker pplTypeChecker =
             wrapSqlOperandTypeChecker(
                 typeChecker, operator.getName(), operator instanceof SqlUserDefinedFunction);
-        register(
-            functionName,
-            (RexBuilder builder, RexNode... args) -> builder.makeCall(operator, args),
-            pplTypeChecker);
+        registerOperator(functionName, operator, pplTypeChecker);
       }
     }
 
-    private static SqlOperandTypeChecker extractTypeCheckerFromUDF(
-        SqlUserDefinedFunction udfOperator) {
-      UDFOperandMetadata udfOperandMetadata =
-          (UDFOperandMetadata) udfOperator.getOperandTypeChecker();
-      return (udfOperandMetadata == null) ? null : udfOperandMetadata.getInnerTypeChecker();
+    /**
+     * Registers an operator for a built-in function name with a specified {@link PPLTypeChecker}.
+     * This allows custom type checking logic to be associated with the operator.
+     *
+     * @param functionName the built-in function name
+     * @param operator the SQL operator to register
+     * @param typeChecker the type checker to use for validating argument types
+     */
+    protected void registerOperator(
+        BuiltinFunctionName functionName, SqlOperator operator, PPLTypeChecker typeChecker) {
+      register(
+          functionName,
+          (RexBuilder builder, RexNode... args) -> builder.makeCall(operator, args),
+          typeChecker);
     }
 
     void populate() {
@@ -690,15 +629,6 @@ void populate() {
       registerOperator(OR, SqlStdOperatorTable.OR);
       registerOperator(NOT, SqlStdOperatorTable.NOT);
 
-      // Register ADD (+ symbol) for numeric addition
-      register(
-          ADD,
-          (RexBuilder builder, RexNode... args) -> builder.makeCall(SqlStdOperatorTable.PLUS, args),
-          new PPLTypeChecker.PPLFamilyTypeChecker(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
-
-      // Register ADD (+ symbol) for string concatenation
-      registerOperator(ADD, SqlStdOperatorTable.CONCAT);
-
       // Register ADDFUNCTION for numeric addition only
       registerOperator(ADDFUNCTION, SqlStdOperatorTable.PLUS);
       registerOperator(SUBTRACT, SqlStdOperatorTable.MINUS);
@@ -853,15 +783,16 @@ void populate() {
       registerOperator(WEEK, PPLBuiltinOperators.WEEK);
       registerOperator(WEEK_OF_YEAR, PPLBuiltinOperators.WEEK);
       registerOperator(WEEKOFYEAR, PPLBuiltinOperators.WEEK);
-      registerOperator(INTERNAL_PATTERN_PARSER, PPLBuiltinOperators.PATTERN_PARSER);
 
+      registerOperator(INTERNAL_PATTERN_PARSER, PPLBuiltinOperators.PATTERN_PARSER);
       registerOperator(ARRAY, PPLBuiltinOperators.ARRAY);
       registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
       registerOperator(FORALL, PPLBuiltinOperators.FORALL);
       registerOperator(EXISTS, PPLBuiltinOperators.EXISTS);
       registerOperator(FILTER, PPLBuiltinOperators.FILTER);
       registerOperator(TRANSFORM, PPLBuiltinOperators.TRANSFORM);
       registerOperator(REDUCE, PPLBuiltinOperators.REDUCE);
+
       // Register Json function
       register(
           JSON_ARRAY,
@@ -889,6 +820,53 @@ void populate() {
       registerOperator(JSON_APPEND, PPLBuiltinOperators.JSON_APPEND);
       registerOperator(JSON_EXTEND, PPLBuiltinOperators.JSON_EXTEND);
 
+      // Register operators with a different type checker
+
+      // Register ADD (+ symbol) for string concatenation
+      // Replaced type checker since CONCAT also supports array concatenation
+      registerOperator(
+          ADD,
+          SqlStdOperatorTable.CONCAT,
+          PPLTypeChecker.family(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER));
+      // Register ADD (+ symbol) for numeric addition
+      // Replace type checker since PLUS also supports binary addition
+      registerOperator(
+          ADD,
+          SqlStdOperatorTable.PLUS,
+          PPLTypeChecker.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
+      // Replace with a custom CompositeOperandTypeChecker to check both operands as
+      // SqlStdOperatorTable.ITEM.getOperandTypeChecker() checks only the first operand instead
+      // of all operands.
+      registerOperator(
+          INTERNAL_ITEM,
+          SqlStdOperatorTable.ITEM,
+          PPLTypeChecker.wrapComposite(
+              (CompositeOperandTypeChecker)
+                  OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.INTEGER)
+                      .or(OperandTypes.family(SqlTypeFamily.MAP, SqlTypeFamily.ANY)),
+              false));
+      registerOperator(
+          XOR,
+          SqlStdOperatorTable.NOT_EQUALS,
+          PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.BOOLEAN));
+      // SqlStdOperatorTable.CASE.getOperandTypeChecker is null. We manually create a type checker
+      // for it. The second and third operands are required to be of the same type. If not,
+      // it will throw an IllegalArgumentException with information Can't find leastRestrictive type
+      registerOperator(
+          IF,
+          SqlStdOperatorTable.CASE,
+          PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.ANY, SqlTypeFamily.ANY));
+      // Re-define the type checker for is not null, is present, and is null since their original
+      // type checker ANY isn't compatible with struct types.
+      registerOperator(
+          IS_NOT_NULL,
+          SqlStdOperatorTable.IS_NOT_NULL,
+          PPLTypeChecker.family(SqlTypeFamily.IGNORE));
+      registerOperator(
+          IS_PRESENT, SqlStdOperatorTable.IS_NOT_NULL, PPLTypeChecker.family(SqlTypeFamily.IGNORE));
+      registerOperator(
+          IS_NULL, SqlStdOperatorTable.IS_NULL, PPLTypeChecker.family(SqlTypeFamily.IGNORE));
+
       // Register implementation.
       // Note, make the implementation an individual class if too complex.
       register(
@@ -922,10 +900,9 @@ void populate() {
                       builder.makeLiteral(" "),
                       arg),
           PPLTypeChecker.family(SqlTypeFamily.CHARACTER));
-      register(
+      registerOperator(
           ATAN,
-          (FunctionImp2)
-              (builder, arg1, arg2) -> builder.makeCall(SqlStdOperatorTable.ATAN2, arg1, arg2),
+          SqlStdOperatorTable.ATAN2,
           PPLTypeChecker.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
       register(
           STRCMP,
@@ -960,17 +937,6 @@ void populate() {
                               SqlTypeFamily.INTEGER,
                               SqlTypeFamily.INTEGER)),
               false));
-      // SqlStdOperatorTable.ITEM.getOperandTypeChecker() checks only the first operand instead of
-      // all operands. Therefore, we wrap it with a custom CompositeOperandTypeChecker to check both
-      // operands.
-      register(
-          INTERNAL_ITEM,
-          (RexBuilder builder, RexNode... args) -> builder.makeCall(SqlStdOperatorTable.ITEM, args),
-          PPLTypeChecker.wrapComposite(
-              (CompositeOperandTypeChecker)
-                  OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.INTEGER)
-                      .or(OperandTypes.family(SqlTypeFamily.MAP, SqlTypeFamily.ANY)),
-              false));
       register(
           LOG,
           (FunctionImp2)
@@ -1002,18 +968,6 @@ void populate() {
               (builder, arg) ->
                   builder.makeLiteral(getLegacyTypeName(arg.getType(), QueryType.PPL)),
           null);
-      register(
-          XOR,
-          (FunctionImp2)
-              (builder, arg1, arg2) -> builder.makeCall(SqlStdOperatorTable.NOT_EQUALS, arg1, arg2),
-          PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.BOOLEAN));
-      // SqlStdOperatorTable.CASE.getOperandTypeChecker is null. We manually create a type checker
-      // for it. The second and third operands are required to be of the same type. If not,
-      // it will throw an IllegalArgumentException with information Can't find leastRestrictive type
-      register(
-          IF,
-          (RexBuilder builder, RexNode... args) -> builder.makeCall(SqlStdOperatorTable.CASE, args),
-          PPLTypeChecker.family(SqlTypeFamily.BOOLEAN, SqlTypeFamily.ANY, SqlTypeFamily.ANY));
       register(
           NULLIF,
           (FunctionImp2)
@@ -1048,20 +1002,6 @@ void populate() {
                               builder.makeLiteral(" "),
                               arg))),
           PPLTypeChecker.family(SqlTypeFamily.ANY));
-      // Re-define the type checker for is not null, is present, and is null since their original
-      // type checker ANY isn't compatible with struct types.
-      register(
-          IS_NOT_NULL,
-          (builder, args) -> builder.makeCall(SqlStdOperatorTable.IS_NOT_NULL, args),
-          PPLTypeChecker.family(SqlTypeFamily.IGNORE));
-      register(
-          IS_PRESENT,
-          (builder, args) -> builder.makeCall(SqlStdOperatorTable.IS_NOT_NULL, args),
-          PPLTypeChecker.family(SqlTypeFamily.IGNORE));
-      register(
-          IS_NULL,
-          (builder, args) -> builder.makeCall(SqlStdOperatorTable.IS_NULL, args),
-          PPLTypeChecker.family(SqlTypeFamily.IGNORE));
       register(
           LIKE,
           (FunctionImp2)
@@ -1075,6 +1015,13 @@ void populate() {
                       builder.makeLiteral("\\")),
           PPLTypeChecker.family(SqlTypeFamily.STRING, SqlTypeFamily.STRING));
     }
+
+    private static SqlOperandTypeChecker extractTypeCheckerFromUDF(
+        SqlUserDefinedFunction udfOperator) {
+      UDFOperandMetadata udfOperandMetadata =
+          (UDFOperandMetadata) udfOperator.getOperandTypeChecker();
+      return (udfOperandMetadata == null) ? null : udfOperandMetadata.getInnerTypeChecker();
+    }
   }
 
   private static class Builder extends AbstractBuilder {
@@ -1213,4 +1160,71 @@ void populate() {
           null);
     }
   }
+
+  /**
+   * Get a string representation of the argument types expressed in ExprType for error messages.
+   *
+   * @param argTypes the list of argument types as {@link RelDataType}
+   * @return a string in the format [type1,type2,...] representing the argument types
+   */
+  private static String getActualSignature(List<RelDataType> argTypes) {
+    return "["
+        + argTypes.stream()
+            .map(OpenSearchTypeFactory::convertRelDataTypeToExprType)
+            .map(Objects::toString)
+            .collect(Collectors.joining(","))
+        + "]";
+  }
+
+  /**
+   * Wraps a {@link SqlOperandTypeChecker} into a {@link PPLTypeChecker} for use in function
+   * signature validation.
+   *
+   * @param typeChecker the original SQL operand type checker
+   * @param functionName the name of the function for error reporting
+   * @param isUserDefinedFunction true if the function is user-defined, false otherwise
+   * @return a {@link PPLTypeChecker} that delegates to the provided {@code typeChecker}
+   */
+  private static PPLTypeChecker wrapSqlOperandTypeChecker(
+      SqlOperandTypeChecker typeChecker, String functionName, boolean isUserDefinedFunction) {
+    PPLTypeChecker pplTypeChecker;
+    // Only the composite operand type checker for UDFs are concerned here.
+    if (isUserDefinedFunction
+        && typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
+      // UDFs implement their own composite type checkers, which always use OR logic for
+      // argument
+      // types. Verifying the composition type would require accessing a protected field in
+      // CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
+      // be skipped, so we avoid checking the composition type here.
+      pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, false);
+    } else if (typeChecker instanceof ImplicitCastOperandTypeChecker implicitCastTypeChecker) {
+      pplTypeChecker = PPLTypeChecker.wrapFamily(implicitCastTypeChecker);
+    } else if (typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
+      // If compositeTypeChecker contains operand checkers other than family type checkers or
+      // other than OR compositions, the function with be registered with a null type checker,
+      // which means the function will not be type checked.
+      try {
+        pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, true);
+      } catch (IllegalArgumentException | UnsupportedOperationException e) {
+        logger.debug(
+            String.format(
+                "Failed to create composite type checker for operator: %s. Will skip its type"
+                    + " checking",
+                functionName),
+            e);
+        pplTypeChecker = null;
+      }
+    } else if (typeChecker instanceof SameOperandTypeChecker comparableTypeChecker) {
+      // Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
+      // SameOperandTypeCheckers like COALESCE, IFNULL, etc.
+      pplTypeChecker = PPLTypeChecker.wrapComparable(comparableTypeChecker);
+    } else if (typeChecker instanceof UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
+      pplTypeChecker = PPLTypeChecker.wrapUDT(udtOperandMetadata.allowedParamTypes());
+    } else {
+      logger.info(
+          "Cannot create type checker for function: {}. Will skip its type checking", functionName);
+      pplTypeChecker = null;
+    }
+    return pplTypeChecker;
+  }
 }