Merge remote-tracking branch 'origin/main' into issues/4136

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

Author: yuancu

async-query-core/src/main/java/org/opensearch/sql/spark/rest/model/LangType.java

@@ -8,7 +8,8 @@
 /** Language type accepted in async query apis. */
 public enum LangType {
   SQL("sql"),
-  PPL("ppl");
+  PPL("ppl"),
+  PROMQL("promql");
   private final String text;
 
   LangType(String text) {

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -86,6 +86,7 @@
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Rex;
+import org.opensearch.sql.ast.tree.SPath;
 import org.opensearch.sql.ast.tree.Search;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
@@ -95,6 +96,7 @@
 import org.opensearch.sql.ast.tree.Trendline;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
+import org.opensearch.sql.ast.tree.Window;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.data.model.ExprMissingValue;
 import org.opensearch.sql.data.type.ExprCoreType;
@@ -755,11 +757,21 @@ public LogicalPlan visitReverse(Reverse node, AnalysisContext context) {
     throw getOnlyForCalciteException("Reverse");
   }
 
+  @Override
+  public LogicalPlan visitSpath(SPath node, AnalysisContext context) {
+    throw getOnlyForCalciteException("Spath");
+  }
+
   @Override
   public LogicalPlan visitTimechart(Timechart node, AnalysisContext context) {
     throw getOnlyForCalciteException("Timechart");
   }
 
+  @Override
+  public LogicalPlan visitWindow(Window node, AnalysisContext context) {
+    throw getOnlyForCalciteException("Window");
+  }
+
   @Override
   public LogicalPlan visitRegex(Regex node, AnalysisContext context) {
     throw getOnlyForCalciteException("Regex");

core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

@@ -38,6 +38,7 @@
 import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.In;
 import org.opensearch.sql.ast.expression.Interval;
+import org.opensearch.sql.ast.expression.LambdaFunction;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
@@ -479,6 +480,11 @@ public Expression visitInSubquery(InSubquery node, AnalysisContext context) {
     throw getOnlyForCalciteException("Subsearch");
   }
 
+  @Override
+  public Expression visitLambdaFunction(LambdaFunction node, AnalysisContext context) {
+    throw getOnlyForCalciteException("Lambda function");
+  }
+
   /**
    * If QualifiedName is actually a reserved metadata field, return the expr type associated with
    * the metadata field.

core/src/main/java/org/opensearch/sql/ast/expression/Literal.java

@@ -7,6 +7,7 @@
 
 import com.google.common.collect.ImmutableList;
 import java.math.BigDecimal;
+import java.text.DecimalFormat;
 import java.util.List;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
@@ -25,7 +26,21 @@ public class Literal extends UnresolvedExpression {
 
   public Literal(Object value, DataType dataType) {
     if (dataType == DataType.DECIMAL && value instanceof Double) {
-      this.value = BigDecimal.valueOf((Double) value);
+      // This branch is only used in testing for backward compatibility:
+      // We accept decimal literal by Literal(double, DataType.DECIMAL).
+      // Some double values such as 0.0001 will be converted to string "1.0E-4" and finally
+      // generate decimal 0.00010. So here we parse a decimal text to Double then convert it
+      // to BigDecimal as well.
+      // In v2, a decimal literal will be converted back to double in resolving expression
+      // via ExprDoubleValue.
+      // In v3, a decimal literal will be kept in Calcite RexNode and converted back to double
+      // in runtime.
+      DecimalFormat df = new DecimalFormat();
+      df.setMaximumFractionDigits(38);
+      df.setMinimumFractionDigits(1);
+      df.setGroupingUsed(false);
+      String plain = df.format(value);
+      this.value = new BigDecimal(plain);
     } else {
       this.value = value;
     }

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.calcite;
 
 import static org.apache.calcite.sql.SqlKind.AS;
+import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.INFORMATION_SCHEMA_NAME;
 import static org.opensearch.sql.ast.tree.Join.JoinType.ANTI;
 import static org.opensearch.sql.ast.tree.Join.JoinType.SEMI;
 import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST;
@@ -20,6 +21,7 @@
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRelation;
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRexCall;
 import static org.opensearch.sql.calcite.utils.PlanUtils.transformPlanToAttachChild;
+import static org.opensearch.sql.utils.SystemIndexUtils.DATASOURCES_TABLE_NAME;
 
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -65,6 +67,7 @@
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.EmptySourcePropagateVisitor;
 import org.opensearch.sql.ast.Node;
@@ -136,6 +139,7 @@
 import org.opensearch.sql.calcite.utils.WildcardUtils;
 import org.opensearch.sql.common.patterns.PatternUtils;
 import org.opensearch.sql.common.utils.StringUtils;
+import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
@@ -148,10 +152,12 @@ public class CalciteRelNodeVisitor extends AbstractNodeVisitor<RelNode, CalciteP
 
   private final CalciteRexNodeVisitor rexVisitor;
   private final CalciteAggCallVisitor aggVisitor;
+  private final DataSourceService dataSourceService;
 
-  public CalciteRelNodeVisitor() {
+  public CalciteRelNodeVisitor(DataSourceService dataSourceService) {
     this.rexVisitor = new CalciteRexNodeVisitor(this);
     this.aggVisitor = new CalciteAggCallVisitor(rexVisitor);
+    this.dataSourceService = dataSourceService;
   }
 
   public RelNode analyze(UnresolvedPlan unresolved, CalcitePlanContext context) {
@@ -160,6 +166,21 @@ public RelNode analyze(UnresolvedPlan unresolved, CalcitePlanContext context) {
 
   @Override
   public RelNode visitRelation(Relation node, CalcitePlanContext context) {
+    DataSourceSchemaIdentifierNameResolver nameResolver =
+        new DataSourceSchemaIdentifierNameResolver(
+            dataSourceService, node.getTableQualifiedName().getParts());
+    if (!nameResolver
+        .getDataSourceName()
+        .equals(DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME)) {
+      throw new CalciteUnsupportedException(
+          "Datasource " + nameResolver.getDataSourceName() + " is unsupported in Calcite");
+    }
+    if (nameResolver.getIdentifierName().equals(DATASOURCES_TABLE_NAME)) {
+      throw new CalciteUnsupportedException("SHOW DATASOURCES is unsupported in Calcite");
+    }
+    if (nameResolver.getSchemaName().equals(INFORMATION_SCHEMA_NAME)) {
+      throw new CalciteUnsupportedException("information_schema is unsupported in Calcite");
+    }
     context.relBuilder.scan(node.getTableQualifiedName().getParts());
     return context.relBuilder.peek();
   }
@@ -1364,7 +1385,7 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
       throw new IllegalArgumentException("Number of duplicate events must be greater than 0");
     }
     if (consecutive) {
-      throw new UnsupportedOperationException("Consecutive deduplication is not supported");
+      throw new CalciteUnsupportedException("Consecutive deduplication is unsupported in Calcite");
     }
     // Columns to deduplicate
     List<RexNode> dedupeFields =
@@ -2330,13 +2351,15 @@ private void buildParseRelNode(Parse node, CalcitePlanContext context) {
         };
     if (ParseMethod.PATTERNS.equals(parseMethod)) {
       rexNodeList = ArrayUtils.add(rexNodeList, context.relBuilder.literal("<*>"));
+    } else {
+      rexNodeList = ArrayUtils.add(rexNodeList, context.relBuilder.literal(parseMethod.getName()));
     }
     List<RexNode> newFields = new ArrayList<>();
     for (String groupCandidate : groupCandidates) {
       RexNode innerRex =
           PPLFuncImpTable.INSTANCE.resolve(
               context.rexBuilder, ParseUtils.BUILTIN_FUNCTION_MAP.get(parseMethod), rexNodeList);
-      if (ParseMethod.GROK.equals(parseMethod)) {
+      if (!ParseMethod.PATTERNS.equals(parseMethod)) {
         newFields.add(
             PPLFuncImpTable.INSTANCE.resolve(
                 context.rexBuilder,

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

@@ -47,6 +47,7 @@
 import org.opensearch.sql.ast.expression.Compare;
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.In;
 import org.opensearch.sql.ast.expression.Interval;
 import org.opensearch.sql.ast.expression.LambdaFunction;
@@ -56,6 +57,7 @@
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.RelevanceFieldList;
+import org.opensearch.sql.ast.expression.ScoreFunction;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
@@ -649,6 +651,16 @@ public RexNode visitWhen(When node, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("CastWhen function is unsupported in Calcite");
   }
 
+  @Override
+  public RexNode visitHighlightFunction(HighlightFunction node, CalcitePlanContext context) {
+    throw new CalciteUnsupportedException("Highlight function is unsupported in Calcite");
+  }
+
+  @Override
+  public RexNode visitScoreFunction(ScoreFunction node, CalcitePlanContext context) {
+    throw new CalciteUnsupportedException("Score function is unsupported in Calcite");
+  }
+
   @Override
   public RexNode visitRelevanceFieldList(RelevanceFieldList node, CalcitePlanContext context) {
     List<RexNode> varArgRexNodeList = new ArrayList<>();

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

@@ -92,6 +92,10 @@ private PPLOperandTypes() {}
       UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.INTEGER_INTEGER);
   public static final UDFOperandMetadata STRING_STRING =
       UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.CHARACTER_CHARACTER);
+  public static final UDFOperandMetadata STRING_STRING_STRING =
+      UDFOperandMetadata.wrap(
+          OperandTypes.family(
+              SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER));
   public static final UDFOperandMetadata NUMERIC_NUMERIC =
       UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC_NUMERIC);
   public static final UDFOperandMetadata STRING_INTEGER =

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

@@ -17,9 +17,23 @@
 
 public class OpenSearchTypeSystem extends RelDataTypeSystemImpl {
   public static final RelDataTypeSystem INSTANCE = new OpenSearchTypeSystem();
+  // same with Spark DecimalType.MAX_PRECISION
+  public static int MAX_PRECISION = 38;
+  // same with Spark DecimalType.MAX_SCALE
+  public static int MAX_SCALE = 38;
 
   private OpenSearchTypeSystem() {}
 
+  @Override
+  public int getMaxNumericPrecision() {
+    return MAX_PRECISION;
+  }
+
+  @Override
+  public int getMaxNumericScale() {
+    return MAX_SCALE;
+  }
+
   @Override
   public RelDataType deriveAvgAggType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
     if (SqlTypeName.DECIMAL == argumentType.getSqlTypeName()) {

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

@@ -9,6 +9,7 @@
 import java.security.PrivilegedAction;
 import java.util.List;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -54,13 +55,12 @@ public class QueryService {
   private final Analyzer analyzer;
   private final ExecutionEngine executionEngine;
   private final Planner planner;
-
-  @Getter(lazy = true)
-  private final CalciteRelNodeVisitor relNodeVisitor = new CalciteRelNodeVisitor();
-
   private DataSourceService dataSourceService;
   private Settings settings;
 
+  @Getter(lazy = true)
+  private final CalciteRelNodeVisitor relNodeVisitor = new CalciteRelNodeVisitor(dataSourceService);
+
   /** Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.<br> */
   public void execute(
       UnresolvedPlan plan,
@@ -104,7 +104,7 @@ public void executeWithCalcite(
                 return null;
               });
     } catch (Throwable t) {
-      if (isCalciteFallbackAllowed() && !(t instanceof NonFallbackCalciteException)) {
+      if (isCalciteFallbackAllowed(t) && !(t instanceof NonFallbackCalciteException)) {
         log.warn("Fallback to V2 query engine since got exception", t);
         executeWithLegacy(plan, queryType, listener, Optional.of(t));
       } else {
@@ -140,7 +140,7 @@ public void explainWithCalcite(
                 return null;
               });
     } catch (Throwable t) {
-      if (isCalciteFallbackAllowed()) {
+      if (isCalciteFallbackAllowed(t)) {
         log.warn("Fallback to V2 query engine since got exception", t);
         explainWithLegacy(plan, queryType, listener, format, Optional.of(t));
       } else {
@@ -162,7 +162,7 @@ public void executeWithLegacy(
     try {
       executePlan(analyze(plan, queryType), PlanContext.emptyPlanContext(), listener);
     } catch (Exception e) {
-      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed()) {
+      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed(null)) {
         // if there is a failure thrown from Calcite and execution after fallback V2
         // keeps failure, we should throw the failure from Calcite.
         calciteFailure.ifPresentOrElse(
@@ -195,7 +195,7 @@ public void explainWithLegacy(
       }
       executionEngine.explain(plan(analyze(plan, queryType)), listener);
     } catch (Exception e) {
-      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed()) {
+      if (shouldUseCalcite(queryType) && isCalciteFallbackAllowed(null)) {
         // if there is a failure thrown from Calcite and execution after fallback V2
         // keeps failure, we should throw the failure from Calcite.
         calciteFailure.ifPresentOrElse(
@@ -260,15 +260,21 @@ public RelNode optimize(RelNode plan, CalcitePlanContext context) {
         SystemLimitType.QUERY_SIZE_LIMIT, plan, context.relBuilder.literal(context.querySizeLimit));
   }
 
-  private boolean isCalciteFallbackAllowed() {
-    if (settings != null) {
-      Boolean fallback_allowed = settings.getSettingValue(Settings.Key.CALCITE_FALLBACK_ALLOWED);
-      if (fallback_allowed == null) {
-        return false;
-      }
-      return fallback_allowed;
-    } else {
+  private boolean isCalciteFallbackAllowed(@Nullable Throwable t) {
+    // We always allow fallback the query failed with CalciteUnsupportedException.
+    // This is for avoiding breaking changes when enable Calcite by default.
+    if (t instanceof CalciteUnsupportedException) {
       return true;
+    } else {
+      if (settings != null) {
+        Boolean fallback_allowed = settings.getSettingValue(Settings.Key.CALCITE_FALLBACK_ALLOWED);
+        if (fallback_allowed == null) {
+          return false;
+        }
+        return fallback_allowed;
+      } else {
+        return true;
+      }
     }
   }
 

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

@@ -323,8 +323,8 @@ public enum BuiltinFunctionName {
   INTERNAL_PATTERN_PARSER(FunctionName.of("pattern_parser")),
   INTERNAL_PATTERN(FunctionName.of("pattern")),
   INTERNAL_UNCOLLECT_PATTERNS(FunctionName.of("uncollect_patterns")),
-  INTERNAL_REGEXP_EXTRACT(FunctionName.of("regexp_extract"), true),
   INTERNAL_GROK(FunctionName.of("grok"), true),
+  INTERNAL_PARSE(FunctionName.of("parse"), true),
   INTERNAL_REGEXP_REPLACE_3(FunctionName.of("regexp_replace_3"), true),
   INTERNAL_REGEXP_REPLACE_PG_4(FunctionName.of("regexp_replace_pg_4"), true),
   INTERNAL_REGEXP_REPLACE_5(FunctionName.of("regexp_replace_5"), true),