Implement implict cast and type validation with Calcite

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -8,6 +8,7 @@
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
 
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -19,10 +20,13 @@
 import org.apache.calcite.rex.RexCorrelVariable;
 import org.apache.calcite.rex.RexLambdaRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.server.CalciteServerStatement;
+import org.apache.calcite.sql.validate.SqlValidator;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.RelBuilder;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
+import org.opensearch.sql.calcite.validate.TypeChecker;
 import org.opensearch.sql.executor.QueryType;
 import org.opensearch.sql.expression.function.FunctionProperties;
 
@@ -35,6 +39,7 @@ public class CalcitePlanContext {
   public final FunctionProperties functionProperties;
   public final QueryType queryType;
   public final Integer querySizeLimit;
+  @Getter public final SqlValidator validator;
 
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
@@ -61,6 +66,13 @@ private CalcitePlanContext(FrameworkConfig config, Integer querySizeLimit, Query
     this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
     this.functionProperties = new FunctionProperties(QueryType.PPL);
     this.rexLambdaRefMap = new HashMap<>();
+    final CalciteServerStatement statement;
+    try {
+      statement = connection.createStatement().unwrap(CalciteServerStatement.class);
+    } catch (SQLException e) {
+      throw new RuntimeException(e);
+    }
+    this.validator = TypeChecker.getValidator(statement, config);
   }
 
   public RexNode resolveJoinCondition(

core/src/main/java/org/opensearch/sql/calcite/PplRelToSqlConverter.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite;
+
+import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
+import org.apache.calcite.sql.SqlDialect;
+
+/**
+ * An extension of {@link RelToSqlConverter} to convert a relation algebra tree, translated from of
+ * PPL query, into a SQL statement.
+ *
+ * <p>Currently, we haven't implemented any specific change to it, just leaving it for future
+ * extension.
+ */
+public class PplRelToSqlConverter extends RelToSqlConverter {
+  /**
+   * Creates a RelToSqlConverter.
+   *
+   * @param dialect the SQL dialect to use
+   */
+  public PplRelToSqlConverter(SqlDialect dialect) {
+    super(dialect);
+  }
+}

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

@@ -80,6 +80,8 @@
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql.validate.SqlValidatorUtil;
 import org.apache.calcite.sql2rel.SqlRexConvertletTable;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.Frameworks;
@@ -90,6 +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;
 
 /**
  * Calcite Tools Helper. This class is used to create customized: 1. Connection 2. JavaTypeFactory
@@ -240,7 +243,7 @@ public <R> R perform(
      * return {@link OpenSearchCalcitePreparingStmt}
      */
     @Override
-    protected CalcitePrepareImpl.CalcitePreparingStmt getPreparingStmt(
+    public CalcitePrepareImpl.CalcitePreparingStmt getPreparingStmt(
         CalcitePrepare.Context context,
         Type elementType,
         CalciteCatalogReader catalogReader,
@@ -332,6 +335,25 @@ public Type getElementType() {
       }
       return super.implement(root);
     }
+
+    /**
+     * Imitated {@link org.apache.calcite.prepare.CalcitePrepareImpl}#createSqlValidator to create a
+     * SqlValidator
+     */
+    protected SqlValidator createSqlValidator(CalciteCatalogReader catalogReader) {
+      return SqlValidatorUtil.newValidator(
+          // this is different from the original implementation
+          PplOpTable.getInstance(),
+          catalogReader,
+          context.getTypeFactory(),
+          // this may be customized in the future
+          SqlValidator.Config.DEFAULT);
+    }
+
+    @Override
+    public SqlValidator getSqlValidator() {
+      return super.getSqlValidator();
+    }
   }
 
   public static class OpenSearchRelRunners {

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

@@ -0,0 +1,129 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.validate;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlIdentifier;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlOperatorTable;
+import org.apache.calcite.sql.SqlSyntax;
+import org.apache.calcite.sql.validate.SqlNameMatcher;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+
+/**
+ * PPLOpTable is a custom implementation of {@link SqlOperatorTable} that provides a way to register
+ * and look up PPL operators.
+ */
+public class PplOpTable implements SqlOperatorTable {
+  // Implementation notes:
+  //  - Did not extend ListSqlOperatorTable because it does not support registering multiple
+  // SqlOperator to one name.
+  //  - Did not extend ReflectiveSqlOperatorTable because it relies on reflectively looking for
+  // member fields of
+  //    SqlOperator type, which is not suitable for our use case.
+  //  - Did not add SqlOperatorTable to PPLFuncImpTable to reduce chaos with existing implementation
+
+  protected Map<BuiltinFunctionName, ArrayList<SqlOperator>> operators;
+
+  private static final PplOpTable INSTANCE = new PplOpTable();
+
+  public static PplOpTable getInstance() {
+    return INSTANCE;
+  }
+
+  private PplOpTable() {
+    this.operators = new HashMap<>();
+  }
+
+  @Override
+  public void lookupOperatorOverloads(
+      SqlIdentifier opName,
+      @Nullable SqlFunctionCategory category,
+      SqlSyntax syntax,
+      List<SqlOperator> operatorList,
+      SqlNameMatcher nameMatcher) {
+    if (!opName.isSimple()) {
+      return;
+    }
+    final String simpleName = opName.getSimple();
+    lookUpOperators(
+        simpleName,
+        op -> {
+          if (op.getSyntax() != syntax && op.getSyntax().family != syntax.family) {
+            // Allow retrieval on exact syntax or family; for example,
+            // CURRENT_DATETIME has FUNCTION_ID syntax but can also be called with
+            // both FUNCTION_ID and FUNCTION syntax (e.g. SELECT CURRENT_DATETIME,
+            // CURRENT_DATETIME('UTC')).
+            return;
+          }
+          if (category != null
+              && category != category(op)
+              && !category.isUserDefinedNotSpecificFunction()) {
+            return;
+          }
+          operatorList.add(op);
+        });
+  }
+
+  protected void lookUpOperators(String name, Consumer<SqlOperator> consumer) {
+    final BuiltinFunctionName funcNameOpt = sqlFunctionNameToPplFunctionName(name);
+    if (funcNameOpt == null) {
+      return; // No function with this name
+    }
+    if (!operators.containsKey(funcNameOpt)) {
+      return; // The function is not registered
+    }
+    operators.get(funcNameOpt).forEach(consumer);
+  }
+
+  /**
+   * At this stage, the function name of a Calcite's builtin operator is acquired via
+   * `sqlFunction.getSqlIdentifier()`
+   *
+   * <p>This will return the name in Calcite, instead of that registered in PPL. We use this method
+   * to convert the Calcite function name to the PPL function name.
+   */
+  private BuiltinFunctionName sqlFunctionNameToPplFunctionName(String name) {
+    return switch (name.toUpperCase(Locale.ROOT)) {
+      case "CONVERT" -> BuiltinFunctionName.CONV;
+      case "ILIKE" -> BuiltinFunctionName.LIKE;
+      case "CHAR_LENGTH" -> BuiltinFunctionName.LENGTH;
+      case "NOT_EQUALS", "<>" -> BuiltinFunctionName.XOR;
+      default -> BuiltinFunctionName.of(name).orElse(null);
+    };
+  }
+
+  protected static SqlFunctionCategory category(SqlOperator operator) {
+    if (operator instanceof SqlFunction) {
+      return ((SqlFunction) operator).getFunctionType();
+    } else {
+      return SqlFunctionCategory.SYSTEM;
+    }
+  }
+
+  @Override
+  public List<SqlOperator> getOperatorList() {
+    return operators.values().stream()
+        .flatMap(iterable -> StreamSupport.stream(iterable.spliterator(), false))
+        .collect(Collectors.toList());
+  }
+
+  public void add(BuiltinFunctionName name, SqlOperator operator) {
+    ArrayList<SqlOperator> list = operators.getOrDefault(name, new ArrayList<>());
+    list.add(operator);
+    operators.put(name, list);
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/validate/PplTypeCoercion.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/PplValidator.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);
+  }
+}