Support NULL and MISSING value in response

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

@@ -15,8 +15,6 @@
 
 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;
@@ -65,19 +63,32 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
-import lombok.RequiredArgsConstructor;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 
 /**
  * Analyze the {@link UnresolvedPlan} in the {@link AnalysisContext} to construct the {@link
  * LogicalPlan}.
  */
-@RequiredArgsConstructor
 public class Analyzer extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
+
   private final ExpressionAnalyzer expressionAnalyzer;
+
+  private final SelectExpressionAnalyzer selectExpressionAnalyzer;
+
   private final StorageEngine storageEngine;
 
+  /**
+   * Constructor.
+   */
+  public Analyzer(
+      ExpressionAnalyzer expressionAnalyzer,
+      StorageEngine storageEngine) {
+    this.expressionAnalyzer = expressionAnalyzer;
+    this.storageEngine = storageEngine;
+    this.selectExpressionAnalyzer = new SelectExpressionAnalyzer(expressionAnalyzer);
+  }
+
   public LogicalPlan analyze(UnresolvedPlan unresolved, AnalysisContext context) {
     return unresolved.accept(this, context);
   }
@@ -113,8 +124,11 @@ public LogicalPlan visitRename(Rename node, AnalysisContext context) {
         ReferenceExpression target =
             new ReferenceExpression(((Field) renameMap.getTarget()).getField().toString(),
                 origin.type());
-        context.peek().define(target);
-        renameMapBuilder.put(DSL.ref(origin.toString(), origin.type()), target);
+        ReferenceExpression originExpr = DSL.ref(origin.toString(), origin.type());
+        TypeEnvironment curEnv = context.peek();
+        curEnv.remove(originExpr);
+        curEnv.define(target);
+        renameMapBuilder.put(originExpr, target);
       } else {
         throw new SemanticCheckException(
             String.format("the target expected to be field, but is %s", renameMap.getTarget()));
@@ -129,17 +143,27 @@ public LogicalPlan visitRename(Rename node, AnalysisContext context) {
    */
   @Override
   public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
-    LogicalPlan child = node.getChild().get(0).accept(this, context);
+    final LogicalPlan child = node.getChild().get(0).accept(this, context);
     ImmutableList.Builder<Aggregator> aggregatorBuilder = new ImmutableList.Builder<>();
     for (UnresolvedExpression expr : node.getAggExprList()) {
       aggregatorBuilder.add((Aggregator) expressionAnalyzer.analyze(expr, context));
     }
+    ImmutableList<Aggregator> aggregators = aggregatorBuilder.build();
 
     ImmutableList.Builder<Expression> groupbyBuilder = new ImmutableList.Builder<>();
     for (UnresolvedExpression expr : node.getGroupExprList()) {
       groupbyBuilder.add(expressionAnalyzer.analyze(expr, context));
     }
-    return new LogicalAggregation(child, aggregatorBuilder.build(), groupbyBuilder.build());
+    ImmutableList<Expression> groupBys = groupbyBuilder.build();
+
+    // new context
+    context.push();
+    TypeEnvironment newEnv = context.peek();
+    aggregators.forEach(aggregator -> newEnv.define(new Symbol(Namespace.FIELD_NAME,
+        aggregator.toString()), aggregator.type()));
+    groupBys.forEach(group -> newEnv.define(new Symbol(Namespace.FIELD_NAME,
+        group.toString()), group.type()));
+    return new LogicalAggregation(child, aggregators, groupBys);
   }
 
   /**
@@ -161,18 +185,24 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
       Argument argument = node.getArgExprList().get(0);
       Boolean exclude = (Boolean) argument.getValue().getValue();
       if (exclude) {
+        TypeEnvironment curEnv = context.peek();
         List<ReferenceExpression> referenceExpressions =
             node.getProjectList().stream()
                 .map(expr -> (ReferenceExpression) expressionAnalyzer.analyze(expr, context))
                 .collect(Collectors.toList());
+        referenceExpressions.forEach(ref -> curEnv.remove(ref));
         return new LogicalRemove(child, ImmutableSet.copyOf(referenceExpressions));
       }
     }
 
-    List<NamedExpression> expressions = node.getProjectList().stream()
-        .map(expr -> named(expressionAnalyzer.analyze(expr, context)))
-        .collect(Collectors.toList());
-    return new LogicalProject(child, expressions);
+    List<NamedExpression> namedExpressions =
+        selectExpressionAnalyzer.analyze(node.getProjectList(), context);
+    // new context
+    context.push();
+    TypeEnvironment newEnv = context.peek();
+    namedExpressions.forEach(expr -> newEnv.define(new Symbol(Namespace.FIELD_NAME,
+        expr.getName()), expr.type()));
+    return new LogicalProject(child, namedExpressions);
   }
 
   /**

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

@@ -164,13 +164,6 @@ public Expression visitQualifiedName(QualifiedName node, AnalysisContext context
     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,

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

@@ -0,0 +1,79 @@
+/*
+ *
+ *    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.ast.AbstractNodeVisitor;
+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.UnresolvedExpression;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
+import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
+import com.amazon.opendistroforelasticsearch.sql.expression.NamedExpression;
+import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
+import com.google.common.collect.ImmutableList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+
+/**
+ * Analyze the select list in the {@link AnalysisContext} to construct the list of
+ * {@link NamedExpression}.
+ */
+@RequiredArgsConstructor
+public class SelectExpressionAnalyzer
+    extends
+    AbstractNodeVisitor<List<NamedExpression>, AnalysisContext> {
+  private final ExpressionAnalyzer expressionAnalyzer;
+
+  /**
+   * Analyze Select fields.
+   */
+  public List<NamedExpression> analyze(List<UnresolvedExpression> selectList,
+                                       AnalysisContext analysisContext) {
+    ImmutableList.Builder<NamedExpression> builder = new ImmutableList.Builder<>();
+    for (UnresolvedExpression unresolvedExpression : selectList) {
+      builder.addAll(unresolvedExpression.accept(this, analysisContext));
+    }
+    return builder.build();
+  }
+
+  @Override
+  public List<NamedExpression> visitField(Field node, AnalysisContext context) {
+    return Collections.singletonList(DSL.named(node.accept(expressionAnalyzer, context)));
+  }
+
+  @Override
+  public List<NamedExpression> visitAlias(Alias node, AnalysisContext context) {
+    return Collections.singletonList(DSL.named(node.getName(),
+        node.getDelegated().accept(expressionAnalyzer, context),
+        node.getAlias()));
+  }
+
+  @Override
+  public List<NamedExpression> visitAllFields(AllFields node,
+                                              AnalysisContext context) {
+    TypeEnvironment environment = context.peek();
+    Map<String, ExprType> lookupAllFields = environment.lookupAllFields(Namespace.FIELD_NAME);
+    return lookupAllFields.entrySet().stream().map(entry -> DSL.named(entry.getKey(),
+        new ReferenceExpression(entry.getKey(), entry.getValue()))).collect(Collectors.toList());
+  }
+}

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

@@ -23,6 +23,8 @@
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.env.Environment;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 import lombok.Getter;
 
@@ -62,6 +64,17 @@ public ExprType resolve(Symbol symbol) {
         String.format("can't resolve %s in type env", symbol));
   }
 
+  /**
+   * Resolve all fields in the current environment.
+   * @param namespace     a namespace
+   * @return              all symbols in the namespace
+   */
+  public Map<String, ExprType> lookupAllFields(Namespace namespace) {
+    Map<String, ExprType> result = new HashMap<>();
+    symbolTable.lookupAllFields(namespace).forEach(result::putIfAbsent);
+    return result;
+  }
+
   /**
    * Define symbol with the type.
    *
@@ -81,4 +94,14 @@ public void define(ReferenceExpression ref) {
     define(new Symbol(Namespace.FIELD_NAME, ref.getAttr()), ref.type());
   }
 
+  public void remove(Symbol symbol) {
+    symbolTable.remove(symbol);
+  }
+
+  /**
+   * Remove ref.
+   */
+  public void remove(ReferenceExpression ref) {
+    remove(new Symbol(Namespace.FIELD_NAME, ref.getAttr()));
+  }
 }

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

@@ -24,6 +24,7 @@
 import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.TreeMap;
+import java.util.stream.Collectors;
 
 /**
  * Symbol table for symbol definition and resolution.
@@ -49,6 +50,19 @@ public void store(Symbol symbol, ExprType type) {
     ).put(symbol.getName(), type);
   }
 
+  /**
+   * Remove a symbol from SymbolTable.
+   */
+  public void remove(Symbol symbol) {
+    tableByNamespace.computeIfPresent(
+        symbol.getNamespace(),
+        (k, v) -> {
+          v.remove(symbol.getName());
+          return v;
+        }
+    );
+  }
+
   /**
    * Look up symbol in the namespace map.
    *
@@ -78,6 +92,21 @@ public Map<String, ExprType> lookupByPrefix(Symbol prefix) {
     return emptyMap();
   }
 
+  /**
+   * Look up all top level symbols in the namespace.
+   * this function is mainly used by SELECT * use case to get the top level fields
+   * Todo. currently, the top level fields is the field which doesn't include "." in the name.
+   *
+   * @param namespace     a namespace
+   * @return              all symbols in the namespace map
+   */
+  public Map<String, ExprType> lookupAllFields(Namespace namespace) {
+    return tableByNamespace.getOrDefault(namespace, emptyNavigableMap())
+        .entrySet().stream()
+        .filter(entry -> !entry.getKey().contains("."))
+        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+  }
+
   /**
    * Check if namespace map in empty (none definition).
    *

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

@@ -17,6 +17,7 @@
 
 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.AttributeList;
@@ -183,4 +184,8 @@ public T visitValues(Values node, C context) {
   public T visitAlias(Alias node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitAllFields(AllFields node, C context) {
+    return visitChildren(node, context);
+  }
 }

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

@@ -17,6 +17,7 @@
 
 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;

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

@@ -0,0 +1,43 @@
+/*
+ *
+ *    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.EqualsAndHashCode;
+import lombok.ToString;
+
+/**
+ * Represent the All fields which is been used in SELECT *.
+ */
+@ToString
+@EqualsAndHashCode(callSuper = false)
+public class AllFields extends UnresolvedExpression {
+  public static final AllFields INSTANCE = new AllFields();
+
+  private AllFields() {
+  }
+
+  public static AllFields of() {
+    return INSTANCE;
+  }
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitAllFields(this, context);
+  }
+}

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

@@ -52,6 +52,17 @@ public boolean hasArgument() {
     return !argExprList.isEmpty();
   }
 
+  /**
+   * The Project could been used to exclude fields from the source.
+   */
+  public boolean isExcluded() {
+    if (hasArgument()) {
+      Argument argument = argExprList.get(0);
+      return (Boolean) argument.getValue().getValue();
+    }
+    return false;
+  }
+
   @Override
   public Project attach(UnresolvedPlan child) {
     this.child = child;

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/data/model/ExprMissingValue.java

@@ -35,7 +35,7 @@ public static ExprMissingValue of() {
 
   @Override
   public Object value() {
-    throw new ExpressionEvaluationException("invalid to call value operation on missing value");
+    return null;
   }
 
   @Override