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

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/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 =

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/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/utils/YamlFormatter.java

@@ -7,7 +7,6 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.SerializationFeature;
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator;
 
@@ -27,7 +26,6 @@ public class YamlFormatter {
         YAMLGenerator.Feature.ALWAYS_QUOTE_NUMBERS_AS_STRINGS); // Quote numeric strings
     yamlFactory.enable(YAMLGenerator.Feature.INDENT_ARRAYS_WITH_INDICATOR);
     YAML_MAPPER = new ObjectMapper(yamlFactory);
-    YAML_MAPPER.enable(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS);
   }
 
   /**

core/src/test/java/org/opensearch/sql/calcite/CalciteRelNodeVisitorSearchSimpleTest.java

@@ -11,9 +11,11 @@
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
 import org.opensearch.sql.ast.dsl.AstDSL;
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.Search;
+import org.opensearch.sql.datasource.DataSourceService;
 
 /**
  * Simple tests for CalciteRelNodeVisitor.visitSearch method. Tests basic functionality without
@@ -22,10 +24,11 @@
 public class CalciteRelNodeVisitorSearchSimpleTest {
 
   private CalciteRelNodeVisitor visitor;
+  @Mock DataSourceService dataSourceService;
 
   @BeforeEach
   public void setUp() {
-    visitor = new CalciteRelNodeVisitor();
+    visitor = new CalciteRelNodeVisitor(dataSourceService);
   }
 
   @Test

core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java

@@ -29,6 +29,7 @@
 import org.opensearch.sql.ast.expression.LambdaFunction;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
+import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.executor.QueryType;
 
 @ExtendWith(MockitoExtension.class)
@@ -44,6 +45,7 @@ public class CalciteRexNodeVisitorTest {
   @Mock RelDataType accType;
   @Mock QualifiedName functionArg1;
   @Mock QualifiedName functionArg2;
+  @Mock DataSourceService dataSourceService;
 
   static CalciteRexNodeVisitor visitor;
   static CalciteRelNodeVisitor relNodeVisitor;
@@ -57,7 +59,7 @@ public class CalciteRexNodeVisitorTest {
 
   @BeforeEach
   public void setUpContext() {
-    relNodeVisitor = new CalciteRelNodeVisitor();
+    relNodeVisitor = new CalciteRelNodeVisitor(dataSourceService);
     visitor = new CalciteRexNodeVisitor(relNodeVisitor);
     when(relBuilder.getRexBuilder()).thenReturn(rexBuilder);
     when(rexBuilder.getTypeFactory()).thenReturn(TYPE_FACTORY);