Support select fields and alias in new query engine (#636)

* Change grammar

* Add UT

* Change grammar and add UT for field alias

* Support alias by named expression

* Add UT

* Add javadoc

* Make query field required after test

* Change java doc

* Fix comparison test data type issue

* Fix broken UT

* Fix jacoco coverage

* Fix broken UTs

* Fix PPL ast builder

* Add doctest

* Fix broken legacy IT

* Fix jacoco

* Fix all IT and preserve original name with alias

* Fix remove command failure

* Fix remove command failure

* Fix doctest

* Fix broken UT

* Fix broken UT

* Revert to make Alias optional for PPL

* Prepare PR

* Prepare PR

* Prepare PR

* Dont remove single quotes for column and alias name

common/src/main/java/com/amazon/opendistroforelasticsearch/sql/common/utils/StringUtils.java

@@ -19,13 +19,13 @@
 
 public class StringUtils {
   /**
-   * Unquote Identifier with mark.
+   * Unquote any string with mark specified.
    * @param text string
    * @param mark quotation mark
    * @return An unquoted string whose outer pair of (single/double/back-tick) quotes have been
    *     removed
    */
-  public static String unquoteIdentifier(String text, String mark) {
+  public static String unquote(String text, String mark) {
     if (isQuoted(text, mark)) {
       return text.substring(mark.length(), text.length() - mark.length());
     }
@@ -38,14 +38,28 @@ public static String unquoteIdentifier(String text, String mark) {
    * @return An unquoted string whose outer pair of (single/double/back-tick) quotes have been
    *     removed
    */
-  public static String unquoteIdentifier(String text) {
+  public static String unquoteText(String text) {
     if (isQuoted(text, "\"") || isQuoted(text, "'") || isQuoted(text, "`")) {
       return text.substring(1, text.length() - 1);
     } else {
       return text;
     }
   }
 
+  /**
+   * Unquote Identifier which has " or ` as mark.
+   * @param identifier identifier that possibly enclosed by double quotes or back ticks
+   * @return An unquoted string whose outer pair of (double/back-tick) quotes have been
+   *     removed
+   */
+  public static String unquoteIdentifier(String identifier) {
+    if (isQuoted(identifier, "\"") || isQuoted(identifier, "`")) {
+      return identifier.substring(1, identifier.length() - 1);
+    } else {
+      return identifier;
+    }
+  }
+
   private static boolean isQuoted(String text, String mark) {
     return !Strings.isNullOrEmpty(text) && text.startsWith(mark) && text.endsWith(mark);
   }

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.expression.DSL.named;
+
 import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Namespace;
 import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Symbol;
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
@@ -40,6 +42,7 @@
 import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.LiteralExpression;
+import com.amazon.opendistroforelasticsearch.sql.expression.NamedExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalAggregation;
@@ -166,8 +169,8 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
       }
     }
 
-    List<Expression> expressions = node.getProjectList().stream()
-        .map(expr -> expressionAnalyzer.analyze(expr, context))
+    List<NamedExpression> expressions = node.getProjectList().stream()
+        .map(expr -> named(expressionAnalyzer.analyze(expr, context)))
         .collect(Collectors.toList());
     return new LogicalProject(child, expressions);
   }

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

@@ -19,6 +19,7 @@
 import com.amazon.opendistroforelasticsearch.sql.analysis.symbol.Symbol;
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.And;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Compare;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.EqualTo;
@@ -27,10 +28,13 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Literal;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Not;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Or;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.QualifiedName;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedAttribute;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Xor;
+import com.amazon.opendistroforelasticsearch.sql.common.antlr.SyntaxCheckException;
 import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException;
 import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
@@ -149,10 +153,40 @@ public Expression visitField(Field node, AnalysisContext context) {
     return visitIdentifier(attr, 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);
+  }
+
+  @Override
+  public Expression visitAlias(Alias node, AnalysisContext context) {
+    return DSL.named(node.getName(),
+                     node.getDelegated().accept(this, context),
+                     node.getAlias());
+  }
+
   private Expression visitIdentifier(String ident, AnalysisContext context) {
     TypeEnvironment typeEnv = context.peek();
     ReferenceExpression ref = DSL.ref(ident,
         typeEnv.resolve(new Symbol(Namespace.FIELD_NAME, ident)));
+
+    // Fall back to old engine too if type is not supported semantically
+    if (isTypeNotSupported(ref.type())) {
+      throw new SyntaxCheckException(String.format(
+          "Identifier [%s] of type [%s] is not supported yet", ident, ref.type()));
+    }
     return ref;
   }
+
+  private boolean isTypeNotSupported(ExprType type) {
+    return "struct".equalsIgnoreCase(type.typeName())
+        || "array".equalsIgnoreCase(type.typeName());
+  }
+
 }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/ast/AbstractNodeVisitor.java

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.sql.ast;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.And;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AttributeList;
@@ -178,4 +179,8 @@ public T visitDedupe(Dedupe node, C context) {
   public T visitValues(Values node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitAlias(Alias node, C context) {
+    return visitChildren(node, context);
+  }
 }

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

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.sql.ast.dsl;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.AggregateFunction;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.And;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Compare;
@@ -43,7 +44,6 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.UnresolvedPlan;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Values;
-import com.google.common.collect.ImmutableList;
 import java.util.Arrays;
 import java.util.List;
 import lombok.experimental.UtilityClass;
@@ -226,6 +226,14 @@ public static Field field(String field, List<Argument> fieldArgs) {
     return new Field(field, fieldArgs);
   }
 
+  public Alias alias(String name, UnresolvedExpression expr) {
+    return new Alias(name, expr);
+  }
+
+  public Alias alias(String name, UnresolvedExpression expr, String alias) {
+    return new Alias(name, expr, alias);
+  }
+
   public static List<UnresolvedExpression> exprList(UnresolvedExpression... exprList) {
     return Arrays.asList(exprList);
   }

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

@@ -0,0 +1,58 @@
+/*
+ *    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.ast.expression;
+
+import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/**
+ * Alias abstraction that associate an unnamed expression with a name and an optional alias.
+ * The name and alias information preserved is useful for semantic analysis and response
+ * formatting eventually. This can avoid restoring the info in toString() method which is
+ * inaccurate because original info is already lost.
+ */
+@AllArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@Getter
+@RequiredArgsConstructor
+@ToString
+public class Alias extends UnresolvedExpression {
+
+  /**
+   * Original field name.
+   */
+  private final String name;
+
+  /**
+   * Expression aliased.
+   */
+  private final UnresolvedExpression delegated;
+
+  /**
+   * Optional field alias.
+   */
+  private String alias;
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitAlias(this, context);
+  }
+}

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/DSL.java

@@ -75,6 +75,30 @@ public static ReferenceExpression ref(String ref, ExprType type) {
     return new ReferenceExpression(ref, type);
   }
 
+  /**
+   * Wrap a named expression if not yet. The intent is that different languages may use
+   * Alias or not when building AST. This caused either named or unnamed expression
+   * is resolved by analyzer. To make unnamed expression acceptable for logical project,
+   * it is required to wrap it by named expression here before passing to logical project.
+   *
+   * @param expression  expression
+   * @return            expression if named already or expression wrapped by named expression.
+   */
+  public static NamedExpression named(Expression expression) {
+    if (expression instanceof NamedExpression) {
+      return (NamedExpression) expression;
+    }
+    return named(expression.toString(), expression);
+  }
+
+  public static NamedExpression named(String name, Expression expression) {
+    return new NamedExpression(name, expression);
+  }
+
+  public static NamedExpression named(String name, Expression expression, String alias) {
+    return new NamedExpression(name, expression, alias);
+  }
+
   public FunctionExpression abs(Expression... expressions) {
     return function(BuiltinFunctionName.ABS, expressions);
   }

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/NamedExpression.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.expression;
+
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
+import com.amazon.opendistroforelasticsearch.sql.expression.env.Environment;
+import com.google.common.base.Strings;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/**
+ * Named expression that represents expression with name.
+ * Please see more details in associated unresolved expression operator
+ * {@link com.amazon.opendistroforelasticsearch.sql.ast.expression.Alias}.
+ */
+@AllArgsConstructor
+@EqualsAndHashCode
+@RequiredArgsConstructor
+@ToString
+public class NamedExpression implements Expression {
+
+  /**
+   * Expression name.
+   */
+  private final String name;
+
+  /**
+   * Expression that being named.
+   */
+  private final Expression delegated;
+
+  /**
+   * Optional alias.
+   */
+  private String alias;
+
+  @Override
+  public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+    return delegated.valueOf(valueEnv);
+  }
+
+  @Override
+  public ExprType type() {
+    return delegated.type();
+  }
+
+  /**
+   * Get expression name using name or its alias (if it's present).
+   * @return  expression name
+   */
+  public String getName() {
+    return Strings.isNullOrEmpty(alias) ? name : alias;
+  }
+
+}

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/planner/logical/LogicalPlanDSL.java

@@ -18,6 +18,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.LiteralExpression;
+import com.amazon.opendistroforelasticsearch.sql.expression.NamedExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
 import com.google.common.collect.ImmutableSet;
@@ -50,7 +51,7 @@ public static LogicalPlan rename(
     return new LogicalRename(input, renameMap);
   }
 
-  public static LogicalPlan project(LogicalPlan input, Expression... fields) {
+  public static LogicalPlan project(LogicalPlan input, NamedExpression... fields) {
     return new LogicalProject(input, Arrays.asList(fields));
   }
 

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/planner/logical/LogicalProject.java

@@ -15,7 +15,7 @@
 
 package com.amazon.opendistroforelasticsearch.sql.planner.logical;
 
-import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
+import com.amazon.opendistroforelasticsearch.sql.expression.NamedExpression;
 import java.util.Arrays;
 import java.util.List;
 import lombok.EqualsAndHashCode;
@@ -32,7 +32,7 @@
 public class LogicalProject extends LogicalPlan {
   private final LogicalPlan child;
   @Getter
-  private final List<Expression> projectList;
+  private final List<NamedExpression> projectList;
 
   @Override
   public List<LogicalPlan> getChild() {

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/planner/physical/PhysicalPlanDSL.java

@@ -18,6 +18,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort.SortOption;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.LiteralExpression;
+import com.amazon.opendistroforelasticsearch.sql.expression.NamedExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
 import com.google.common.collect.ImmutableSet;
@@ -47,7 +48,7 @@ public static RenameOperator rename(
     return new RenameOperator(input, renameMap);
   }
 
-  public static ProjectOperator project(PhysicalPlan input, Expression... fields) {
+  public static ProjectOperator project(PhysicalPlan input, NamedExpression... fields) {
     return new ProjectOperator(input, Arrays.asList(fields));
   }