Refine PPL head command syntax

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

@@ -29,7 +29,6 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Let;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Literal;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Map;
-import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedArgument;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.UnresolvedExpression;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Aggregation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Dedupe;
@@ -59,7 +58,6 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalDedupe;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalEval;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalFilter;
-import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalHead;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalLimit;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalProject;
@@ -358,20 +356,11 @@ public LogicalPlan visitDedupe(Dedupe node, AnalysisContext context) {
   }
 
   /**
-   * Build {@link LogicalHead}.
+   * Logical head is identical to {@link LogicalLimit}.
    */
   public LogicalPlan visitHead(Head node, AnalysisContext context) {
     LogicalPlan child = node.getChild().get(0).accept(this, context);
-    List<UnresolvedArgument> options = node.getOptions();
-    Boolean keeplast = (Boolean) getOptionAsLiteral(options, 0).getValue();
-    Expression whileExpr = expressionAnalyzer.analyze(options.get(1).getValue(), context);
-    Integer number = (Integer) getOptionAsLiteral(options, 2).getValue();
-
-    return new LogicalHead(child, keeplast, whileExpr, number);
-  }
-
-  private static Literal getOptionAsLiteral(List<UnresolvedArgument> options, int optionIdx) {
-    return (Literal) options.get(optionIdx).getValue();
+    return new LogicalLimit(child, node.getSize(), 0);
   }
 
   @Override

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

@@ -365,18 +365,8 @@ public static Dedupe dedupe(UnresolvedPlan input, List<Argument> options, Field.
     return new Dedupe(input, options, Arrays.asList(fields));
   }
 
-  public static Head head(UnresolvedPlan input, List<UnresolvedArgument> options) {
-    return new Head(input, options);
-  }
-
-  /**
-   * Default Head Command Args.
-   */
-  public static List<UnresolvedArgument> defaultHeadArgs() {
-    return unresolvedArgList(
-            unresolvedArg("keeplast", booleanLiteral(true)),
-            unresolvedArg("whileExpr", booleanLiteral(true)),
-            unresolvedArg("number", intLiteral(10)));
+  public static Head head(UnresolvedPlan input, Integer size) {
+    return new Head(input, size);
   }
 
   public static List<Argument> defaultTopArgs() {

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

@@ -38,7 +38,7 @@
 public class Head extends UnresolvedPlan {
 
   private UnresolvedPlan child;
-  private final List<UnresolvedArgument> options;
+  private final Integer size;
 
   @Override
   public Head attach(UnresolvedPlan child) {

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/Explain.java

@@ -24,7 +24,6 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.DedupeOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.EvalOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.FilterOperator;
-import com.amazon.opendistroforelasticsearch.sql.planner.physical.HeadOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.LimitOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanNodeVisitor;
@@ -141,15 +140,6 @@ public ExplainResponseNode visitRareTopN(RareTopNOperator node, Object context)
     )));
   }
 
-  @Override
-  public ExplainResponseNode visitHead(HeadOperator node, Object context) {
-    return explain(node, context, explainNode -> explainNode.setDescription(ImmutableMap.of(
-        "keepLast", node.getKeepLast(),
-        "whileExpr", node.getWhileExpr().toString(),
-        "number", node.getNumber()
-    )));
-  }
-
   @Override
   public ExplainResponseNode visitValues(ValuesOperator node, Object context) {
     return explain(node, context, explainNode -> explainNode.setDescription(ImmutableMap.of(

core/src/main/java/com/amazon/opendistroforelasticsearch/sql/planner/DefaultImplementor.java

@@ -20,7 +20,6 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalDedupe;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalEval;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalFilter;
-import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalHead;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalLimit;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlanNodeVisitor;
@@ -36,7 +35,6 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.DedupeOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.EvalOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.FilterOperator;
-import com.amazon.opendistroforelasticsearch.sql.planner.physical.HeadOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.LimitOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.ProjectOperator;
@@ -80,15 +78,6 @@ public PhysicalPlan visitDedupe(LogicalDedupe node, C context) {
         node.getConsecutive());
   }
 
-  @Override
-  public PhysicalPlan visitHead(LogicalHead node, C context) {
-    return new HeadOperator(
-            visitChild(node, context),
-            node.getKeeplast(),
-            node.getWhileExpr(),
-            node.getNumber());
-  }
-
   @Override
   public PhysicalPlan visitProject(LogicalProject node, C context) {
     return new ProjectOperator(visitChild(node, context), node.getProjectList());

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

@@ -90,11 +90,6 @@ public static LogicalPlan dedupe(
     return new LogicalDedupe(
         input, Arrays.asList(fields), allowedDuplication, keepEmpty, consecutive);
   }
-
-  public static LogicalPlan head(
-      LogicalPlan input, boolean keeplast, Expression whileExpr, int number) {
-    return new LogicalHead(input, keeplast, whileExpr, number);
-  }
 
   public static LogicalPlan rareTopN(LogicalPlan input, CommandType commandType,
       List<Expression> groupByList, Expression... fields) {

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

@@ -43,10 +43,6 @@ public R visitDedupe(LogicalDedupe plan, C context) {
     return visitNode(plan, context);
   }
 
-  public R visitHead(LogicalHead plan, C context) {
-    return visitNode(plan, context);
-  }
-
   public R visitRename(LogicalRename plan, C context) {
     return visitNode(plan, context);
   }

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

@@ -88,11 +88,6 @@ public WindowOperator window(PhysicalPlan input,
     return new WindowOperator(input, windowFunction, windowDefinition);
   }
 
-  public static HeadOperator head(PhysicalPlan input, boolean keepLast, Expression whileExpr,
-      int number) {
-    return new HeadOperator(input, keepLast, whileExpr, number);
-  }
-
   public static RareTopNOperator rareTopN(PhysicalPlan input, CommandType commandType,
       List<Expression> groups, Expression... expressions) {
     return new RareTopNOperator(input, commandType, Arrays.asList(expressions), groups);

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

@@ -72,10 +72,6 @@ public R visitValues(ValuesOperator node, C context) {
   public R visitSort(SortOperator node, C context) {
     return visitNode(node, context);
   }
-
-  public R visitHead(HeadOperator node, C context) {
-    return visitNode(node, context);
-  }
 
   public R visitRareTopN(RareTopNOperator node, C context) {
     return visitNode(node, context);

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

@@ -80,15 +80,8 @@ public void filter_relation() {
   @Test
   public void head_relation() {
     assertAnalyzeEqual(
-        LogicalPlanDSL.head(
-            LogicalPlanDSL.relation("schema"),
-            false, dsl.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(1))), 10),
-        AstDSL.head(
-            AstDSL.relation("schema"),
-            unresolvedArgList(
-                unresolvedArg("keeplast", booleanLiteral(false)),
-                unresolvedArg("whileExpr", compare("=", field("integer_value"), intLiteral(1))),
-                unresolvedArg("number", intLiteral(10)))));
+        LogicalPlanDSL.limit(LogicalPlanDSL.relation("schema"),10, 0),
+        AstDSL.head(AstDSL.relation("schema"), 10));
   }
 
   @Test

core/src/test/java/com/amazon/opendistroforelasticsearch/sql/executor/ExplainTest.java

@@ -28,7 +28,6 @@
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.dedupe;
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.eval;
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.filter;
-import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.head;
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.limit;
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.project;
 import static com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlanDSL.rareTopN;
@@ -140,28 +139,6 @@ void can_explain_rare_top_n() {
         explain.apply(plan));
   }
 
-  @Test
-  void can_explain_head() {
-    Boolean keepLast = false;
-    Expression whileExpr = dsl.and(
-        dsl.equal(ref("balance", INTEGER), literal(10000)),
-        dsl.greater(ref("age", INTEGER), literal(30)));
-    Integer number = 5;
-
-    PhysicalPlan plan = head(tableScan, keepLast, whileExpr, number);
-
-    assertEquals(
-        new ExplainResponse(
-            new ExplainResponseNode(
-                "HeadOperator",
-                ImmutableMap.of(
-                    "keepLast", false,
-                    "whileExpr", "and(=(balance, 10000), >(age, 30))",
-                    "number", 5),
-                singletonList(tableScan.explainNode()))),
-        explain.apply(plan));
-  }
-
   @Test
   void can_explain_window() {
     List<Expression> partitionByList = ImmutableList.of(DSL.ref("state", STRING));