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,109 @@
+/*
+ * 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.Map;
+import java.util.Optional;
+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 Optional<BuiltinFunctionName> funcNameOpt = BuiltinFunctionName.of(name);
+    if (funcNameOpt.isEmpty()) {
+      return; // No such function
+    }
+    operators.get(funcNameOpt.get()).forEach(consumer);
+  }
+
+  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);
+  }
+}

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.validate;
+
+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.implicit.TypeCoercion;
+import org.apache.calcite.tools.FrameworkConfig;
+import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
+
+public class TypeChecker {
+  public static SqlValidator getValidator(
+      CalciteServerStatement statement, FrameworkConfig config) {
+    SchemaPlus defaultSchema = config.getDefaultSchema();
+
+    final CalcitePrepare.Context prepareContext = statement.createPrepareContext();
+    final CalciteSchema schema =
+        defaultSchema != null ? CalciteSchema.from(defaultSchema) : prepareContext.getRootSchema();
+    CalciteCatalogReader catalogReader =
+        new CalciteCatalogReader(
+            schema.root(),
+            schema.path(null),
+            OpenSearchTypeFactory.TYPE_FACTORY,
+            prepareContext.config());
+    SqlValidator.Config validatorConfig =
+        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();
+    // 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);
+  }
+}