Support table alias in new engine (#695)

* Change grammar and ast builder

* Define index symbol and resolve it when visiting qualified name

* Raise error if qualifier is not index name or alias

* Fix index not found exception by adding another get name or alias method

* Remove qualifier if select item is a qualified field

* Prepare PR

* Prepare PR

* Add doctest

* Add bug fix test

* Add more UT for where after where PR merged

Author: dai-chen

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/Analyzer.java

@@ -15,6 +15,8 @@
 
 package com.amazon.opendistroforelasticsearch.sql.analysis;
 
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRUCT;
+
 import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Namespace;
 import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Symbol;
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
@@ -99,6 +101,11 @@ public LogicalPlan visitRelation(Relation node, AnalysisContext context) {
     TypeEnvironment curEnv = context.peek();
     Table table = storageEngine.getTable(node.getTableName());
     table.getFieldTypes().forEach((k, v) -> curEnv.define(new Symbol(Namespace.FIELD_NAME, k), v));
+
+    // Put index name or its alias in index namespace on type environment so qualifier
+    // can be removed when analyzing qualified name. The value (expr type) here doesn't matter.
+    curEnv.define(new Symbol(Namespace.INDEX_NAME, node.getTableNameOrAlias()), STRUCT);
+
     return new LogicalRelation(node.getTableName());
   }
 

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/ExpressionAnalyzer.java

@@ -163,13 +163,8 @@ public Expression visitField(Field node, AnalysisContext context) {
 
   @Override
   public Expression visitQualifiedName(QualifiedName node, AnalysisContext context) {
-    // Name with qualifier (index.field, index_alias.field, object/nested.inner_field
-    //  text.keyword) is not supported for now
-    if (node.getParts().size() > 1) {
-      throw new SyntaxCheckException(String.format(
-          "Qualified name [%s] is not supported yet", node));
-    }
-    return visitIdentifier(node.toString(), context);
+    QualifierAnalyzer qualifierAnalyzer = new QualifierAnalyzer(context);
+    return visitIdentifier(qualifierAnalyzer.unqualified(node), context);
   }
 
   private Expression visitIdentifier(String ident, AnalysisContext context) {

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/QualifierAnalyzer.java

@@ -0,0 +1,72 @@
+/*
+ *    Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License").
+ *    You may not use this file except in compliance with the License.
+ *    A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    or in the "license" file accompanying this file. This file is distributed
+ *    on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *    express or implied. See the License for the specific language governing
+ *    permissions and limitations under the License.
+ *
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.analysis;
+
+import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Namespace;
+import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Symbol;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.QualifiedName;
+import com.amazon.opendistroforelasticsearch.sql.common.antlr.SyntaxCheckException;
+import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException;
+import java.util.Arrays;
+import java.util.Optional;
+import lombok.RequiredArgsConstructor;
+
+/**
+ * Analyzer that analyzes qualifier(s) in a full field name.
+ */
+@RequiredArgsConstructor
+public class QualifierAnalyzer {
+
+  private final AnalysisContext context;
+
+  public String unqualified(String... parts) {
+    return unqualified(QualifiedName.of(Arrays.asList(parts)));
+  }
+
+  /**
+   * Get unqualified name if its qualifier symbol found is in index namespace
+   * on type environment. Unqualified name means name with qualifier removed.
+   * For example, unqualified name of "accounts.age" or "acc.age" is "age".
+   *
+   * @return  unqualified name if criteria met above, otherwise original name
+   */
+  public String unqualified(QualifiedName fullName) {
+    return isQualifierIndexOrAlias(fullName) ? fullName.rest().toString() : fullName.toString();
+  }
+
+  private boolean isQualifierIndexOrAlias(QualifiedName fullName) {
+    Optional<String> qualifier = fullName.first();
+    if (qualifier.isPresent()) {
+      resolveQualifierSymbol(fullName, qualifier.get());
+      return true;
+    }
+    return false;
+  }
+
+  private void resolveQualifierSymbol(QualifiedName fullName, String qualifier) {
+    try {
+      context.peek().resolve(new Symbol(Namespace.INDEX_NAME, qualifier));
+    } catch (SemanticCheckException e) {
+      // Throw syntax check intentionally to indicate fall back to old engine.
+      // Need change to semantic check exception in future.
+      throw new SyntaxCheckException(String.format(
+          "The qualifier [%s] of qualified name [%s] must be an index name or its alias",
+              qualifier, fullName));
+    }
+  }
+
+}

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/SelectExpressionAnalyzer.java

@@ -22,6 +22,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AllFields;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Field;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.QualifiedName;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
@@ -63,7 +64,8 @@ public List<NamedExpression> visitField(Field node, AnalysisContext context) {
 
   @Override
   public List<NamedExpression> visitAlias(Alias node, AnalysisContext context) {
-    return Collections.singletonList(DSL.named(node.getName(),
+    return Collections.singletonList(DSL.named(
+        unqualifiedNameIfFieldOnly(node, context),
         node.getDelegated().accept(expressionAnalyzer, context),
         node.getAlias()));
   }
@@ -76,4 +78,23 @@ public List<NamedExpression> visitAllFields(AllFields node,
     return lookupAllFields.entrySet().stream().map(entry -> DSL.named(entry.getKey(),
         new ReferenceExpression(entry.getKey(), entry.getValue()))).collect(Collectors.toList());
   }
+
+  /**
+   * Get unqualified name if select item is just a field. For example, suppose an index
+   * named "accounts", return "age" for "SELECT accounts.age". But do nothing for expression
+   * in "SELECT ABS(accounts.age)".
+   * Note that an assumption is made implicitly that original name field in Alias must be
+   * the same as the values in QualifiedName. This is true because AST builder does this.
+   * Otherwise, what unqualified() returns will override Alias's name as NamedExpression's name
+   * even though the QualifiedName doesn't have qualifier.
+   */
+  private String unqualifiedNameIfFieldOnly(Alias node, AnalysisContext context) {
+    UnresolvedExpression selectItem = node.getDelegated();
+    if (selectItem instanceof QualifiedName) {
+      QualifierAnalyzer qualifierAnalyzer = new QualifierAnalyzer(context);
+      return qualifierAnalyzer.unqualified((QualifiedName) selectItem);
+    }
+    return node.getName();
+  }
+
 }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/analysis/symbol/Namespace.java

@@ -20,6 +20,7 @@
  */
 public enum Namespace {
 
+  INDEX_NAME("Index"),
   FIELD_NAME("Field"),
   FUNCTION_NAME("Function");
 

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/dsl/AstDSL.java

@@ -17,7 +17,6 @@
 
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
-import com.amazon.opendistroforelasticsearch.sql.ast.expression.AllFields;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.And;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Compare;
@@ -61,10 +60,14 @@ public static UnresolvedPlan filter(UnresolvedPlan input, UnresolvedExpression e
     return new Filter(expression).attach(input);
   }
 
-  public static UnresolvedPlan relation(String tableName) {
+  public UnresolvedPlan relation(String tableName) {
     return new Relation(qualifiedName(tableName));
   }
 
+  public UnresolvedPlan relation(String tableName, String alias) {
+    return new Relation(qualifiedName(tableName), alias);
+  }
+
   public static UnresolvedPlan project(UnresolvedPlan input, UnresolvedExpression... projectList) {
     return new Project(Arrays.asList(projectList)).attach(input);
   }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/expression/QualifiedName.java

@@ -60,7 +60,7 @@ public static QualifiedName of(String first, String... rest) {
     return new QualifiedName(parts);
   }
 
-  private static QualifiedName of(Iterable<String> parts) {
+  public static QualifiedName of(Iterable<String> parts) {
     return new QualifiedName(parts);
   }
 
@@ -78,6 +78,34 @@ public String getSuffix() {
     return parts.get(parts.size() - 1);
   }
 
+  /**
+   * Get first part of the qualified name.
+   * @return  first part
+   */
+  public Optional<String> first() {
+    if (parts.size() == 1) {
+      return Optional.empty();
+    }
+    return Optional.of(parts.get(0));
+  }
+
+  /**
+   * Get rest parts of the qualified name. Assume that there must be remaining parts
+   * so caller is responsible for the check (first() or size() must be called first).
+   * For example:
+   * {@code
+   *   QualifiedName name = ...
+   *   Optional<String> first = name.first();
+   *   if (first.isPresent()) {
+   *     name.rest() ...
+   *   }
+   * }
+   * @return  rest part(s)
+   */
+  public QualifiedName rest() {
+    return QualifiedName.of(parts.subList(1, parts.size()));
+  }
+
   public String toString() {
     return String.join(".", this.parts);
   }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/tree/Relation.java

@@ -19,23 +19,43 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 
 /**
  * Logical plan node of Relation, the interface for building the searching sources.
  */
+@AllArgsConstructor
 @ToString
 @EqualsAndHashCode(callSuper = false)
 @RequiredArgsConstructor
 public class Relation extends UnresolvedPlan {
   private final UnresolvedExpression tableName;
 
+  /**
+   * Optional alias name for the relation.
+   */
+  private String alias;
+
+  /**
+   * Get original table name. Unwrap and get name if table name expression
+   * is actually an Alias.
+   * @return    table name
+   */
   public String getTableName() {
     return tableName.toString();
   }
 
+  /**
+   * Get original table name or its alias if present in Alias.
+   * @return    table name or its alias
+   */
+  public String getTableNameOrAlias() {
+    return (alias == null) ? getTableName() : alias;
+  }
+
   @Override
   public List<UnresolvedPlan> getChild() {
     return ImmutableList.of();

core/src/test/java/com/amazon/opendistroforelasticsearch/sql/analysis/ExpressionAnalyzerTest.java

@@ -20,10 +20,12 @@
 import static com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils.integerValue;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.BOOLEAN;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER;
-import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTERVAL;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRUCT;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Namespace;
+import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Symbol;
 import com.amazon.opendistroforelasticsearch.sql.ast.dsl.AstDSL;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.DataType;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
@@ -92,22 +94,31 @@ public void qualified_name() {
   }
 
   @Test
-  public void interval() {
+  public void qualified_name_with_qualifier() {
+    analysisContext.push();
+    analysisContext.peek().define(new Symbol(Namespace.INDEX_NAME, "index_alias"), STRUCT);
     assertAnalyzeEqual(
-        dsl.interval(DSL.literal(1L), DSL.literal("DAY")),
-        AstDSL.intervalLiteral(1L, DataType.LONG, "DAY"));
-  }
+        DSL.ref("integer_value", INTEGER),
+        AstDSL.qualifiedName("index_alias", "integer_value")
+    );
 
-  @Test
-  public void skip_identifier_with_qualifier() {
+    analysisContext.peek().define(new Symbol(Namespace.FIELD_NAME, "nested_field"), STRUCT);
     SyntaxCheckException exception =
         assertThrows(SyntaxCheckException.class,
-            () -> analyze(AstDSL.qualifiedName("index_alias", "integer_value")));
-
+            () -> analyze(AstDSL.qualifiedName("nested_field", "integer_value")));
     assertEquals(
-        "Qualified name [index_alias.integer_value] is not supported yet",
+        "The qualifier [nested_field] of qualified name [nested_field.integer_value] "
+            + "must be an index name or its alias",
         exception.getMessage()
     );
+    analysisContext.pop();
+  }
+
+  @Test
+  public void interval() {
+    assertAnalyzeEqual(
+        dsl.interval(DSL.literal(1L), DSL.literal("DAY")),
+        AstDSL.intervalLiteral(1L, DataType.LONG, "DAY"));
   }
 
   @Test