Update the type checker of transform function to allow arbitrary additional arguments (1931/2069)

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

Author: yuancu

core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/TransformFunctionImpl.java

@@ -6,10 +6,12 @@
 package org.opensearch.sql.expression.function.CollectionUDF;
 
 import static org.apache.calcite.sql.type.SqlTypeUtil.createArrayType;
+import static org.apache.calcite.util.Static.RESOURCE;
 import static org.opensearch.sql.expression.function.CollectionUDF.LambdaUtils.transferLambdaOutputToTargetType;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.IntStream;
 import org.apache.calcite.adapter.enumerable.NotNullImplementor;
 import org.apache.calcite.adapter.enumerable.NullPolicy;
 import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
@@ -19,11 +21,19 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlOperandCountRange;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlUtil;
 import org.apache.calcite.sql.type.ArraySqlType;
-import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.sql.type.SqlTypeUtil;
 import org.opensearch.sql.expression.function.ImplementorUDF;
 import org.opensearch.sql.expression.function.UDFOperandMetadata;
 
@@ -50,8 +60,94 @@ public SqlReturnTypeInference getReturnTypeInference() {
 
   @Override
   public UDFOperandMetadata getOperandMetadata() {
+    // Only checks the first two arguments as it allows arbitrary number of arguments to follow them
     return UDFOperandMetadata.wrap(
-        OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.FUNCTION));
+        new SqlSingleOperandTypeChecker() {
+          private static final List<SqlTypeFamily> families =
+              List.of(SqlTypeFamily.ARRAY, SqlTypeFamily.FUNCTION);
+
+          /**
+           * Copied from {@link FamilyOperandTypeChecker#checkSingleOperandType(SqlCallBinding
+           * callBinding, SqlNode node, int iFormalOperand, boolean throwOnFailure)}
+           */
+          @Override
+          public boolean checkSingleOperandType(
+              SqlCallBinding callBinding,
+              SqlNode operand,
+              int iFormalOperand,
+              boolean throwOnFailure) {
+            // Do not check types after the second operands
+            if (iFormalOperand > 1) {
+              return true;
+            }
+            SqlTypeFamily family = families.get(iFormalOperand);
+            switch (family) {
+              case ANY:
+                final RelDataType type = SqlTypeUtil.deriveType(callBinding, operand);
+                SqlTypeName typeName = type.getSqlTypeName();
+
+                if (typeName == SqlTypeName.CURSOR) {
+                  // We do not allow CURSOR operands, even for ANY
+                  if (throwOnFailure) {
+                    throw callBinding.newValidationSignatureError();
+                  }
+                  return false;
+                }
+              // fall through
+              case IGNORE:
+                // no need to check
+                return true;
+              default:
+                break;
+            }
+            if (SqlUtil.isNullLiteral(operand, false)) {
+              if (callBinding.isTypeCoercionEnabled()) {
+                return true;
+              } else if (throwOnFailure) {
+                throw callBinding
+                    .getValidator()
+                    .newValidationError(operand, RESOURCE.nullIllegal());
+              } else {
+                return false;
+              }
+            }
+            RelDataType type = SqlTypeUtil.deriveType(callBinding, operand);
+            SqlTypeName typeName = type.getSqlTypeName();
+
+            // Pass type checking for operators if it's of type 'ANY'.
+            if (typeName.getFamily() == SqlTypeFamily.ANY) {
+              return true;
+            }
+
+            if (!family.getTypeNames().contains(typeName)) {
+              if (throwOnFailure) {
+                throw callBinding.newValidationSignatureError();
+              }
+              return false;
+            }
+            return true;
+          }
+
+          @Override
+          public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
+            if (!getOperandCountRange().isValidCount(callBinding.getOperandCount())) {
+              return false;
+            }
+            return IntStream.range(0, 2)
+                .allMatch(
+                    i -> checkSingleOperandType(callBinding, callBinding.operand(i), i, false));
+          }
+
+          @Override
+          public SqlOperandCountRange getOperandCountRange() {
+            return SqlOperandCountRanges.from(2);
+          }
+
+          @Override
+          public String getAllowedSignatures(SqlOperator op, String opName) {
+            return "<ARRAY, FUNCTION, ...ANY OTHER ARGUMENTS>";
+          }
+        });
   }
 
   public static class TransformImplementor implements NotNullImplementor {