[PPL] Add support for fields command

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

@@ -16,13 +16,16 @@
 package com.amazon.opendistroforelasticsearch.sql.analysis;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Field;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Aggregation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Filter;
+import com.amazon.opendistroforelasticsearch.sql.ast.tree.Project;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Relation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.UnresolvedPlan;
+import com.amazon.opendistroforelasticsearch.sql.data.model.ExprMissingValue;
 import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException;
 import com.amazon.opendistroforelasticsearch.sql.expression.DSL;
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
@@ -31,12 +34,17 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalAggregation;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalFilter;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan;
+import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalProject;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRelation;
+import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRemove;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRename;
 import com.amazon.opendistroforelasticsearch.sql.storage.StorageEngine;
 import com.amazon.opendistroforelasticsearch.sql.storage.Table;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import java.util.List;
+import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
 
 /**
@@ -108,4 +116,29 @@ public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
         }
         return new LogicalAggregation(child, aggregatorBuilder.build(), groupbyBuilder.build());
     }
+
+    /**
+     * Build {@link LogicalProject} or {@link LogicalRemove} from {@link Field}.
+     *
+     * <p>Todo, the include/exclude fields should change the env definition. The cons of current
+     * implementation is even the query contain the field reference which has been excluded from fields command. There
+     * is no {@link SemanticCheckException} will be thrown. Instead, the during runtime evaluation, the not exist filed
+     * will be resolve to {@link ExprMissingValue} which will not impact the correctness.
+     *
+     * Postpone the implementation when finding more use case.
+     */
+    @Override
+    public LogicalPlan visitProject(Project node, AnalysisContext context) {
+        LogicalPlan child = node.getChild().get(0).accept(this, context);
+        List<ReferenceExpression> referenceExpressions = node.getProjectList().stream()
+                .map(expr -> (ReferenceExpression) expressionAnalyzer.analyze(expr, context))
+                .collect(Collectors.toList());
+        Argument argument = node.getArgExprList().get(0);
+        Boolean exclude = (Boolean) argument.getValue().getValue();
+        if (exclude) {
+            return new LogicalRemove(child, ImmutableSet.copyOf(referenceExpressions));
+        } else {
+            return new LogicalProject(child, referenceExpressions);
+        }
+    }
 }

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

@@ -57,12 +57,12 @@ public static UnresolvedPlan project(UnresolvedPlan input, UnresolvedExpression.
         return new Project(Arrays.asList(projectList)).attach(input);
     }
 
-    public static UnresolvedPlan projectWithArg(UnresolvedPlan input, List<UnresolvedExpression> argList, UnresolvedExpression... projectList) {
+    public static UnresolvedPlan projectWithArg(UnresolvedPlan input, List<Argument> argList, UnresolvedExpression... projectList) {
         return new Project(Arrays.asList(projectList), argList).attach(input);
     }
 
     public static UnresolvedPlan agg(UnresolvedPlan input, List<UnresolvedExpression> aggList, List<UnresolvedExpression> sortList,
-                                     List<UnresolvedExpression> groupList, List<UnresolvedExpression> argList) {
+                                     List<UnresolvedExpression> groupList, List<Argument> argList) {
         return new Aggregation(aggList, sortList, groupList, argList).attach(input);
     }
 
@@ -82,27 +82,27 @@ public static UnresolvedExpression unresolvedAttr(String attr) {
         return new UnresolvedAttribute(attr);
     }
 
-    private static UnresolvedExpression literal(Object value, DataType type) {
+    private static Literal literal(Object value, DataType type) {
         return new Literal(value, type);
     }
 
-    public static UnresolvedExpression intLiteral(Integer value) {
+    public static Literal intLiteral(Integer value) {
         return literal(value, DataType.INTEGER);
     }
 
-    public static UnresolvedExpression doubleLiteral(Double value) {
+    public static Literal doubleLiteral(Double value) {
         return literal(value, DataType.DOUBLE);
     }
 
-    public static UnresolvedExpression stringLiteral(String value) {
+    public static Literal stringLiteral(String value) {
         return literal(value, DataType.STRING);
     }
 
-    public static UnresolvedExpression booleanLiteral(Boolean value) {
+    public static Literal booleanLiteral(Boolean value) {
         return literal(value, DataType.BOOLEAN);
     }
 
-    public static UnresolvedExpression nullLiteral() {
+    public static Literal nullLiteral() {
         return literal(null, DataType.NULL);
     }
 
@@ -146,7 +146,7 @@ public static UnresolvedExpression compare(String operator, UnresolvedExpression
         return new Compare(operator, left, right);
     }
 
-    public static UnresolvedExpression argument(String argName, UnresolvedExpression argValue) {
+    public static Argument argument(String argName, Literal argValue) {
         return new Argument(argName, argValue);
     }
 
@@ -158,33 +158,37 @@ public static UnresolvedExpression field(String field) {
         return new Field(field);
     }
 
-    public static UnresolvedExpression field(UnresolvedExpression field, UnresolvedExpression... fieldArgs) {
+    public static UnresolvedExpression field(UnresolvedExpression field, Argument... fieldArgs) {
         return new Field((QualifiedName) field, Arrays.asList(fieldArgs));
     }
 
-    public static UnresolvedExpression field(String field, UnresolvedExpression... fieldArgs) {
+    public static UnresolvedExpression field(String field, Argument... fieldArgs) {
         return new Field(field, Arrays.asList(fieldArgs));
     }
 
-    public static UnresolvedExpression field(UnresolvedExpression field, List<UnresolvedExpression> fieldArgs) {
+    public static UnresolvedExpression field(UnresolvedExpression field, List<Argument> fieldArgs) {
         return new Field((QualifiedName) field, fieldArgs);
     }
 
-    public static UnresolvedExpression field(String field, List<UnresolvedExpression> fieldArgs) {
+    public static UnresolvedExpression field(String field, List<Argument> fieldArgs) {
         return new Field(field, fieldArgs);
     }
 
     public static List<UnresolvedExpression> exprList(UnresolvedExpression... exprList) {
         return Arrays.asList(exprList);
     }
 
-    public static List<UnresolvedExpression> defaultFieldsArgs() {
+    public static List<Argument> exprList(Argument... exprList) {
+        return Arrays.asList(exprList);
+    }
+
+    public static List<Argument> defaultFieldsArgs() {
         return exprList(
                 argument("exclude", booleanLiteral(false))
         );
     }
 
-    public static List<UnresolvedExpression> defaultStatsArgs() {
+    public static List<Argument> defaultStatsArgs() {
         return exprList(
                 argument("partitions", intLiteral(1)),
                 argument("allnum", booleanLiteral(false)),
@@ -193,7 +197,7 @@ public static List<UnresolvedExpression> defaultStatsArgs() {
         );
     }
 
-    public static List<UnresolvedExpression> defaultDedupArgs() {
+    public static List<Argument> defaultDedupArgs() {
         return exprList(
                 argument("number", intLiteral(1)),
                 argument("keepevents", booleanLiteral(false)),
@@ -202,14 +206,14 @@ public static List<UnresolvedExpression> defaultDedupArgs() {
         );
     }
 
-    public static List<UnresolvedExpression> defaultSortArgs() {
+    public static List<Argument> defaultSortArgs() {
         return exprList(
                 argument("count", intLiteral(1000)),
                 argument("desc", booleanLiteral(false))
         );
     }
 
-    public static List<UnresolvedExpression> defaultSortFieldArgs() {
+    public static List<Argument> defaultSortFieldArgs() {
         return exprList(
                 argument("exclude", booleanLiteral(false)),
                 argument("type", nullLiteral())

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

@@ -29,7 +29,7 @@
 @EqualsAndHashCode(callSuper = false)
 public class Argument extends UnresolvedExpression {
     private final String argName;
-    private final UnresolvedExpression value;
+    private final Literal value;
     //    private final DataType valueType;
     @Override
     public List<UnresolvedExpression> getChild() {

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

@@ -29,7 +29,7 @@
 @AllArgsConstructor
 public class Field extends UnresolvedExpression {
     private QualifiedName field;
-    private List<UnresolvedExpression> fieldArgs;
+    private List<Argument> fieldArgs;
 
     public Field(QualifiedName field) {
         this.field = field;
@@ -40,7 +40,7 @@ public Field(String field) {
         this.field = new QualifiedName(field);
     }
 
-    public Field(String field, List<UnresolvedExpression> fieldArgs) {
+    public Field(String field, List<Argument> fieldArgs) {
         this.field = new QualifiedName(field);
         this.fieldArgs = fieldArgs;
     }

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

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.sql.ast.tree;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.google.common.collect.ImmutableList;
 import java.util.Collections;
@@ -36,7 +37,7 @@ public class Aggregation extends UnresolvedPlan {
     private List<UnresolvedExpression> aggExprList;
     private List<UnresolvedExpression> sortExprList;
     private List<UnresolvedExpression> groupExprList;
-    private List<UnresolvedExpression> argExprList;
+    private List<Argument> argExprList;
     private UnresolvedPlan child;
 
     public Aggregation(List<UnresolvedExpression> aggExprList,
@@ -51,7 +52,7 @@ public Aggregation(List<UnresolvedExpression> aggExprList,
     public Aggregation(List<UnresolvedExpression> aggExprList,
                        List<UnresolvedExpression> sortExprList,
                        List<UnresolvedExpression> groupExprList,
-                       List<UnresolvedExpression> argExprList) {
+                       List<Argument> argExprList) {
         this.aggExprList = aggExprList;
         this.sortExprList = sortExprList;
         this.groupExprList = groupExprList;

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

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.sql.ast.tree;
 
 import com.amazon.opendistroforelasticsearch.sql.ast.AbstractNodeVisitor;
+import com.amazon.opendistroforelasticsearch.sql.ast.expression.Argument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.google.common.collect.ImmutableList;
 import java.util.Collections;
@@ -34,15 +35,15 @@
 public class Project extends UnresolvedPlan {
     @Setter
     private List<UnresolvedExpression> projectList;
-    private List<UnresolvedExpression> argExprList;
+    private List<Argument> argExprList;
     private UnresolvedPlan child;
 
     public Project(List<UnresolvedExpression> projectList) {
         this.projectList = projectList;
         this.argExprList = Collections.emptyList();
     }
 
-    public Project(List<UnresolvedExpression> projectList, List<UnresolvedExpression> argExprList) {
+    public Project(List<UnresolvedExpression> projectList, List<Argument> argExprList) {
         this.projectList = projectList;
         this.argExprList = argExprList;
     }

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

@@ -18,30 +18,39 @@
 import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
 import com.amazon.opendistroforelasticsearch.sql.expression.ReferenceExpression;
 import com.amazon.opendistroforelasticsearch.sql.expression.aggregation.Aggregator;
+import com.google.common.collect.ImmutableSet;
+import java.util.Arrays;
 import lombok.experimental.UtilityClass;
 
 import java.util.List;
 import java.util.Map;
 
-/**
- * Logical Plan DSL.
- */
+/** Logical Plan DSL. */
 @UtilityClass
 public class LogicalPlanDSL {
-    public static LogicalPlan aggregation(LogicalPlan input, List<Aggregator> aggregatorList,
-                                          List<Expression> groupByList) {
-        return new LogicalAggregation(input, aggregatorList, groupByList);
-    }
-
-    public static LogicalPlan filter(LogicalPlan input, Expression expression) {
-        return new LogicalFilter(input, expression);
-    }
-
-    public static LogicalPlan relation(String tableName) {
-        return new LogicalRelation(tableName);
-    }
-
-    public static LogicalPlan rename(LogicalPlan input, Map<ReferenceExpression, ReferenceExpression> renameMap) {
-        return new LogicalRename(input, renameMap);
-    }
+  public static LogicalPlan aggregation(
+      LogicalPlan input, List<Aggregator> aggregatorList, List<Expression> groupByList) {
+    return new LogicalAggregation(input, aggregatorList, groupByList);
+  }
+
+  public static LogicalPlan filter(LogicalPlan input, Expression expression) {
+    return new LogicalFilter(input, expression);
+  }
+
+  public static LogicalPlan relation(String tableName) {
+    return new LogicalRelation(tableName);
+  }
+
+  public static LogicalPlan rename(
+      LogicalPlan input, Map<ReferenceExpression, ReferenceExpression> renameMap) {
+    return new LogicalRename(input, renameMap);
+  }
+
+  public static LogicalPlan project(LogicalPlan input, ReferenceExpression... fields) {
+    return new LogicalProject(input, Arrays.asList(fields));
+  }
+
+  public static LogicalPlan remove(LogicalPlan input, ReferenceExpression... fields) {
+    return new LogicalRemove(input, ImmutableSet.copyOf(fields));
+  }
 }

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

@@ -22,23 +22,31 @@
  * @param <C> context type.
  */
 public abstract class LogicalPlanNodeVisitor<R, C> {
-    protected R visitNode(LogicalPlan plan, C context) {
-        return null;
-    }
+  protected R visitNode(LogicalPlan plan, C context) {
+    return null;
+  }
 
-    public R visitRelation(LogicalRelation plan, C context) {
-        return visitNode(plan, context);
-    }
+  public R visitRelation(LogicalRelation plan, C context) {
+    return visitNode(plan, context);
+  }
 
-    public R visitFilter(LogicalFilter plan, C context) {
-        return visitNode(plan, context);
-    }
+  public R visitFilter(LogicalFilter plan, C context) {
+    return visitNode(plan, context);
+  }
 
-    public R visitAggregation(LogicalAggregation plan, C context) {
-        return visitNode(plan, context);
-    }
+  public R visitAggregation(LogicalAggregation plan, C context) {
+    return visitNode(plan, context);
+  }
 
-    public R visitRename(LogicalRename plan, C context) {
-        return visitNode(plan, context);
-    }
+  public R visitRename(LogicalRename plan, C context) {
+    return visitNode(plan, context);
+  }
+
+  public R visitProject(LogicalProject plan, C context) {
+    return visitNode(plan, context);
+  }
+
+  public R visitRemove(LogicalRemove plan, C context) {
+    return visitNode(plan, context);
+  }
 }