Check composition type for only calcite's built-in operators

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

Author: yuancu

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

@@ -186,23 +186,18 @@ void registerOperator(BuiltinFunctionName functionName, SqlOperator operator) {
       // Only the composite operand type checker for UDFs are concerned here.
       if (operator instanceof SqlUserDefinedFunction
           && typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
-        register(functionName, createCompositeFunctionImp(operator, 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.
+        register(functionName, createCompositeFunctionImp(operator, compositeTypeChecker, false));
       } else if (typeChecker instanceof ImplicitCastOperandTypeChecker implicitCastTypeChecker) {
         register(functionName, createImplicitCastFunctionImp(operator, implicitCastTypeChecker));
       } else if (typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
-        try {
-          // If compositeTypeChecker contains operand checkers other than family type checkers or
-          // other than OR compositions, this will throw an IllegalArgumentException.
-          register(functionName, createCompositeFunctionImp(operator, compositeTypeChecker));
-        } catch (IllegalArgumentException e) {
-          // register without type checker
-          logger.debug(
-              "Cannot create composite type checker for function: {}. Will skip its type checking",
-              functionName);
-          register(
-              functionName,
-              (RexBuilder builder, RexNode... node) -> builder.makeCall(operator, node));
-        }
+        // 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.
+        register(functionName, createCompositeFunctionImp(operator, compositeTypeChecker, true));
       } else {
         logger.debug(
             "Cannot create type checker for function: {}. Will skip its type checking",
@@ -221,7 +216,9 @@ private static SqlOperandTypeChecker extractTypeCheckerFromUDF(
     }
 
     private static FunctionImp createCompositeFunctionImp(
-        SqlOperator operator, CompositeOperandTypeChecker typeChecker) {
+        SqlOperator operator,
+        CompositeOperandTypeChecker typeChecker,
+        boolean checkCompositionType) {
       return new FunctionImp() {
         @Override
         public RexNode resolve(RexBuilder builder, RexNode... args) {
@@ -230,7 +227,17 @@ public RexNode resolve(RexBuilder builder, RexNode... args) {
 
         @Override
         public PPLTypeChecker getTypeChecker() {
-          return wrapComposite(typeChecker);
+          try {
+            return wrapComposite(typeChecker, checkCompositionType);
+          } catch (IllegalArgumentException | UnsupportedOperationException e) {
+            logger.debug(
+                String.format(
+                    "Failed to create composite type checker for operator: %s. Will skip its type"
+                        + " checking",
+                    operator.getName()),
+                e);
+            return null;
+          }
         }
       };
     }
@@ -250,6 +257,36 @@ public PPLTypeChecker getTypeChecker() {
       };
     }
 
+    private static FunctionImp createFunctionImpWithTypeChecker(
+        BiFunction<RexBuilder, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
+      return new FunctionImp1() {
+        @Override
+        public RexNode resolve(RexBuilder builder, RexNode arg1) {
+          return resolver.apply(builder, arg1);
+        }
+
+        @Override
+        public PPLTypeChecker getTypeChecker() {
+          return typeChecker;
+        }
+      };
+    }
+
+    private static FunctionImp createFunctionImpWithTypeChecker(
+        TriFunction<RexBuilder, RexNode, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
+      return new FunctionImp2() {
+        @Override
+        public RexNode resolve(RexBuilder builder, RexNode arg1, RexNode arg2) {
+          return resolver.apply(builder, arg1, arg2);
+        }
+
+        @Override
+        public PPLTypeChecker getTypeChecker() {
+          return typeChecker;
+        }
+      };
+    }
+
     void populate() {
       // Register std operator
       registerOperator(AND, SqlStdOperatorTable.AND);
@@ -537,34 +574,4 @@ public PPLTypeChecker getTypeChecker() {
       return family(booleanFamily, booleanFamily);
     }
   }
-
-  private static FunctionImp createFunctionImpWithTypeChecker(
-      BiFunction<RexBuilder, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
-    return new FunctionImp1() {
-      @Override
-      public RexNode resolve(RexBuilder builder, RexNode arg1) {
-        return resolver.apply(builder, arg1);
-      }
-
-      @Override
-      public PPLTypeChecker getTypeChecker() {
-        return typeChecker;
-      }
-    };
-  }
-
-  private static FunctionImp createFunctionImpWithTypeChecker(
-      TriFunction<RexBuilder, RexNode, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
-    return new FunctionImp2() {
-      @Override
-      public RexNode resolve(RexBuilder builder, RexNode arg1, RexNode arg2) {
-        return resolver.apply(builder, arg1, arg2);
-      }
-
-      @Override
-      public PPLTypeChecker getTypeChecker() {
-        return typeChecker;
-      }
-    };
-  }
 }

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.expression.function;
 
 import java.lang.reflect.Field;
+import java.lang.reflect.InaccessibleObjectException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -17,7 +18,6 @@
 import org.apache.calcite.sql.type.SqlOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
-import org.apache.logging.log4j.LogManager;
 import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
 
 /**
@@ -225,30 +225,29 @@ static PPLFamilyTypeCheckerWrapper wrapFamily(ImplicitCastOperandTypeChecker typ
    * are instances of {@link ImplicitCastOperandTypeChecker}. If any rule does not meet this
    * requirement, an {@link IllegalArgumentException} is thrown.
    *
+   * <p>Additionally, if {@code checkCompositionType} is true, the method checks if the composition
+   * type of the provided {@code CompositeOperandTypeChecker} is OR via reflection. If it is not, an
+   * {@link IllegalArgumentException} is thrown. If the reflective access to the composition field
+   * of CompositeOperandTypeChecker fails, an {@link UnsupportedOperationException} is thrown.
+   *
    * @param typeChecker the Calcite {@link CompositeOperandTypeChecker} to wrap
+   * @param checkCompositionType if true, checks if the composition type is OR.
    * @return a {@link PPLCompositeTypeChecker} that delegates type checking to the wrapped rules
    * @throws IllegalArgumentException if any rule is not an {@link ImplicitCastOperandTypeChecker}
    */
-  static PPLCompositeTypeChecker wrapComposite(CompositeOperandTypeChecker typeChecker) {
-    try {
-      Field compositionField = CompositeOperandTypeChecker.class.getDeclaredField("composition");
-      compositionField.setAccessible(true);
-      CompositeOperandTypeChecker.Composition composition =
-          (CompositeOperandTypeChecker.Composition) compositionField.get(typeChecker);
-      if (composition != CompositeOperandTypeChecker.Composition.OR) {
-        throw new IllegalArgumentException(
-            String.format(
-                "Currently only OR compositions of ImplicitCastOperandTypeChecker are supported,"
-                    + " but got %s composition",
-                composition.name()));
+  static PPLCompositeTypeChecker wrapComposite(
+      CompositeOperandTypeChecker typeChecker, boolean checkCompositionType)
+      throws IllegalArgumentException, UnsupportedOperationException {
+    if (checkCompositionType) {
+      try {
+        if (!isCompositionOr(typeChecker)) {
+          throw new IllegalArgumentException(
+              "Currently only support CompositeOperandTypeChecker with a OR composition");
+        }
+      } catch (ReflectiveOperationException | InaccessibleObjectException | SecurityException e) {
+        throw new UnsupportedOperationException(
+            String.format("Failed to check composition type of %s", typeChecker), e);
       }
-    } catch (IllegalAccessException e) {
-      LogManager.getLogger(PPLTypeChecker.class).error(e);
-    } catch (NoSuchFieldException e) {
-      LogManager.getLogger(PPLTypeChecker.class)
-          .error(
-              "Failed to access the composition field of CompositeOperandTypeChecker. "
-                  + "This may indicate a change in the Calcite library.");
     }
 
     for (SqlOperandTypeChecker rule : typeChecker.getRules()) {
@@ -308,4 +307,25 @@ private static String getFamilySignature(List<SqlTypeFamily> families) {
         + families.stream().map(SqlTypeFamily::toString).collect(Collectors.joining(","))
         + "]";
   }
+
+  /**
+   * Checks if the provided {@link CompositeOperandTypeChecker} is of type OR composition.
+   *
+   * <p>This method uses reflection to access the protected "composition" field of the
+   * CompositeOperandTypeChecker class.
+   *
+   * @param typeChecker the CompositeOperandTypeChecker to check
+   * @return true if the composition is OR, false otherwise
+   */
+  private static boolean isCompositionOr(CompositeOperandTypeChecker typeChecker)
+      throws NoSuchFieldException,
+          IllegalAccessException,
+          InaccessibleObjectException,
+          SecurityException {
+    Field compositionField = CompositeOperandTypeChecker.class.getDeclaredField("composition");
+    compositionField.setAccessible(true);
+    CompositeOperandTypeChecker.Composition composition =
+        (CompositeOperandTypeChecker.Composition) compositionField.get(typeChecker);
+    return composition == CompositeOperandTypeChecker.Composition.OR;
+  }
 }