Experiment blacklist certain coercion by extending TypeCoercionImpl

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

Author: yuancu

core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java

@@ -92,7 +92,7 @@
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.calcite.plan.Scannable;
 import org.opensearch.sql.calcite.udf.udaf.NullableSqlAvgAggFunction;
-import org.opensearch.sql.calcite.validate.PPLOpTable;
+import org.opensearch.sql.calcite.validate.PplOpTable;
 
 /**
  * Calcite Tools Helper. This class is used to create customized: 1. Connection 2. JavaTypeFactory
@@ -343,7 +343,7 @@ public Type getElementType() {
     protected SqlValidator createSqlValidator(CalciteCatalogReader catalogReader) {
       return SqlValidatorUtil.newValidator(
           // this is different from the original implementation
-          PPLOpTable.getInstance(),
+          PplOpTable.getInstance(),
           catalogReader,
           context.getTypeFactory(),
           // this may be customized in the future

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

@@ -27,7 +27,7 @@
  * PPLOpTable is a custom implementation of {@link SqlOperatorTable} that provides a way to register
  * and look up PPL operators.
  */
-public class PPLOpTable implements SqlOperatorTable {
+public class PplOpTable implements SqlOperatorTable {
   // Implementation notes:
   //  - Did not extend ListSqlOperatorTable because it does not support registering multiple
   // SqlOperator to one name.
@@ -38,13 +38,13 @@ public class PPLOpTable implements SqlOperatorTable {
 
   protected Map<BuiltinFunctionName, ArrayList<SqlOperator>> operators;
 
-  private static final PPLOpTable INSTANCE = new PPLOpTable();
+  private static final PplOpTable INSTANCE = new PplOpTable();
 
-  public static PPLOpTable getInstance() {
+  public static PplOpTable getInstance() {
     return INSTANCE;
   }
 
-  private PPLOpTable() {
+  private PplOpTable() {
     this.operators = new HashMap<>();
   }
 

…rch/sql/calcite/validate/PPLOpTable.java …rch/sql/calcite/validate/PplOpTable.javacore/src/main/java/org/opensearch/sql/calcite/validate/PPLOpTable.java renamed to core/src/main/java/org/opensearch/sql/calcite/validate/PplOpTable.java core/src/main/java/org/opensearch/sql/calcite/validate/PPLOpTable.java renamed to core/src/main/java/org/opensearch/sql/calcite/validate/PplOpTable.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.validate;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.IntStream;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql.validate.implicit.TypeCoercionImpl;
+
+public class PplTypeCoercion extends TypeCoercionImpl {
+  // A blacklist of coercions that are not allowed in PPL.
+  // key cannot be cast from values
+  private static final Map<SqlTypeFamily, Set<SqlTypeFamily>> BLACKLISTED_COERCIONS;
+
+  static {
+    // Initialize the blacklist for coercions that are not allowed in PPL.
+    BLACKLISTED_COERCIONS =
+        Map.of(
+            SqlTypeFamily.CHARACTER,
+            Set.of(SqlTypeFamily.NUMERIC),
+            SqlTypeFamily.STRING,
+            Set.of(SqlTypeFamily.NUMERIC),
+            SqlTypeFamily.NUMERIC,
+            Set.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.STRING));
+  }
+
+  public PplTypeCoercion(RelDataTypeFactory typeFactory, SqlValidator validator) {
+    super(typeFactory, validator);
+  }
+
+  @Override
+  public boolean builtinFunctionCoercion(
+      SqlCallBinding binding,
+      List<RelDataType> operandTypes,
+      List<SqlTypeFamily> expectedFamilies) {
+    assert binding.getOperandCount() == operandTypes.size();
+    if (IntStream.range(0, operandTypes.size())
+        .anyMatch(i -> isBlacklistedCoercion(operandTypes.get(i), expectedFamilies.get(i)))) {
+      return false;
+    }
+    return super.builtinFunctionCoercion(binding, operandTypes, expectedFamilies);
+  }
+
+  // This method tries to blacklist coercions that are not allowed in PPL.
+  private boolean isBlacklistedCoercion(RelDataType operandType, SqlTypeFamily expectedFamily) {
+    if (BLACKLISTED_COERCIONS.containsKey(expectedFamily)) {
+      Set<SqlTypeFamily> blacklistedFamilies = BLACKLISTED_COERCIONS.get(expectedFamily);
+      if (blacklistedFamilies.contains(operandType.getSqlTypeName().getFamily())) {
+        return true;
+      }
+    }
+    return false;
+  }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.validate;
+
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.sql.SqlOperatorTable;
+import org.apache.calcite.sql.validate.SqlValidatorCatalogReader;
+import org.apache.calcite.sql.validate.SqlValidatorImpl;
+
+public class PplValidator extends SqlValidatorImpl {
+  /**
+   * Creates a validator.
+   *
+   * @param opTab Operator table
+   * @param catalogReader Catalog reader
+   * @param typeFactory Type factory
+   * @param config Config
+   */
+  protected PplValidator(
+      SqlOperatorTable opTab,
+      SqlValidatorCatalogReader catalogReader,
+      RelDataTypeFactory typeFactory,
+      Config config) {
+    super(opTab, catalogReader, typeFactory, config);
+  }
+}

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

@@ -5,15 +5,15 @@
 
 package org.opensearch.sql.calcite.validate;
 
-import java.util.Map;
 import org.apache.calcite.jdbc.CalcitePrepare;
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.prepare.CalciteCatalogReader;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.schema.SchemaPlus;
 import org.apache.calcite.server.CalciteServerStatement;
 import org.apache.calcite.sql.type.SqlTypeCoercionRule;
 import org.apache.calcite.sql.validate.SqlValidator;
-import org.apache.calcite.sql.validate.SqlValidatorUtil;
+import org.apache.calcite.sql.validate.implicit.TypeCoercion;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 
@@ -32,16 +32,25 @@ public static SqlValidator getValidator(
             OpenSearchTypeFactory.TYPE_FACTORY,
             prepareContext.config());
     SqlValidator.Config validatorConfig =
-        SqlValidator.Config.DEFAULT.withTypeCoercionRules(getTypeCoercionRule());
-    return SqlValidatorUtil.newValidator(
-        PPLOpTable.getInstance(),
+        SqlValidator.Config.DEFAULT
+            .withTypeCoercionRules(getTypeCoercionRule())
+            .withTypeCoercionFactory(TypeChecker::createTypeCoercion);
+    return new PplValidator(
+        PplOpTable.getInstance(),
         catalogReader,
         OpenSearchTypeFactory.TYPE_FACTORY,
         validatorConfig);
   }
 
   public static SqlTypeCoercionRule getTypeCoercionRule() {
     var defaultMapping = SqlTypeCoercionRule.instance().getTypeMapping();
-    return SqlTypeCoercionRule.instance(Map.of());
+    // try deleting all coercion rules
+    return SqlTypeCoercionRule.instance(defaultMapping);
+  }
+
+  /** Creates a custom TypeCoercion instance for PPL. This can be used as a TypeCoercionFactory. */
+  public static TypeCoercion createTypeCoercion(
+      RelDataTypeFactory typeFactory, SqlValidator validator) {
+    return new PplTypeCoercion(typeFactory, validator);
   }
 }

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

@@ -153,6 +153,7 @@ public void explainWithCalcite(
                     CalcitePlanContext.create(
                         buildFrameworkConfig(), getQuerySizeLimit(), queryType);
                 RelNode relNode = analyze(plan, context);
+                RelNode validated = validate(relNode, context);
                 RelNode optimized = optimize(relNode);
                 RelNode calcitePlan = convertToCalcitePlan(optimized);
                 executionEngine.explain(calcitePlan, format, context, listener);
@@ -284,15 +285,24 @@ public SqlNode visit(SqlIdentifier id) {
               }
             });
     SqlValidator validator = context.getValidator();
+    SqlNode validated;
     if (rewritten != null) {
       try {
-        SqlNode validated = validator.validate(rewritten);
-        log.debug("After validation [{}]", validated);
+        String before = rewritten.toString();
+        // rewritten will be modified in-place
+        validator.validate(rewritten);
+        log.debug("After validation [{}]", rewritten);
+        String after = rewritten.toString();
+        if (before.equals(after)) {
+          // If the rewritten SQL node is not modified, we can return the original RelNode as is
+          return relNode;
+        }
       } catch (CalciteContextException e) {
         throw new ExpressionEvaluationException(e.getMessage(), e);
       }
     } else {
       log.debug("Failed to rewrite the SQL node before validation: {}", root);
+      return relNode;
     }
 
     // Convert the validated SqlNode to RelNode

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -255,7 +255,7 @@
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 import org.opensearch.sql.calcite.utils.PlanUtils;
 import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
-import org.opensearch.sql.calcite.validate.PPLOpTable;
+import org.opensearch.sql.calcite.validate.PplOpTable;
 import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.executor.QueryType;
 
@@ -521,7 +521,7 @@ void registerOperator(BuiltinFunctionName functionName, SqlOperator operator) {
       }
 
       // Currently, only functions registered via registerOperator is added to PPLOpTable
-      PPLOpTable.getInstance().add(functionName, operator);
+      PplOpTable.getInstance().add(functionName, operator);
     }
 
     private static SqlOperandTypeChecker extractTypeCheckerFromUDF(