Support Top/Rare Command In PPL

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

@@ -31,6 +31,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Eval;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Filter;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Project;
+import com.amazon.opendistroforelasticsearch.sql.ast.tree.RareTopN;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Relation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort;
@@ -51,6 +52,7 @@
 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.LogicalRareTopN;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRelation;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRemove;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRename;
@@ -173,6 +175,39 @@ public LogicalPlan visitAggregation(Aggregation node, AnalysisContext context) {
     return new LogicalAggregation(child, aggregators, groupBys);
   }
 
+  /**
+   * Build {@link LogicalRareTopN}.
+   */
+  @Override
+  public LogicalPlan visitRareTopN(RareTopN node, AnalysisContext context) {
+    final LogicalPlan child = node.getChild().get(0).accept(this, context);
+
+    ImmutableList.Builder<Expression> groupbyBuilder = new ImmutableList.Builder<>();
+    for (UnresolvedExpression expr : node.getGroupExprList()) {
+      groupbyBuilder.add(expressionAnalyzer.analyze(expr, context));
+    }
+    ImmutableList<Expression> groupBys = groupbyBuilder.build();
+
+    ImmutableList.Builder<Expression> fieldsBuilder = new ImmutableList.Builder<>();
+    for (Field f : node.getFields()) {
+      fieldsBuilder.add(expressionAnalyzer.analyze(f, context));
+    }
+    ImmutableList<Expression> fields = fieldsBuilder.build();
+
+    // new context
+    context.push();
+    TypeEnvironment newEnv = context.peek();
+    groupBys.forEach(group -> newEnv.define(new Symbol(Namespace.FIELD_NAME,
+        group.toString()), group.type()));
+    fields.forEach(field -> newEnv.define(new Symbol(Namespace.FIELD_NAME,
+        field.toString()), field.type()));
+
+    List<Argument> options = node.getNoOfResults();
+    Integer noOfResults = (Integer) options.get(0).getValue().getValue();
+
+    return new LogicalRareTopN(child, node.getRareTopFlag(), noOfResults, fields, groupBys);
+  }
+
   /**
    * Build {@link LogicalProject} or {@link LogicalRemove} from {@link Field}.
    *

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

@@ -40,6 +40,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Eval;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Filter;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Project;
+import com.amazon.opendistroforelasticsearch.sql.ast.tree.RareTopN;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Relation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort;
@@ -178,6 +179,10 @@ public T visitDedupe(Dedupe node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitRareTopN(RareTopN node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitValues(Values node, C context) {
     return visitChildren(node, context);
   }

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

@@ -26,7 +26,6 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Function;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.In;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Interval;
-import com.amazon.opendistroforelasticsearch.sql.ast.expression.IntervalUnit;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Let;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Literal;
 import com.amazon.opendistroforelasticsearch.sql.ast.expression.Map;
@@ -41,6 +40,7 @@
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Eval;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Filter;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Project;
+import com.amazon.opendistroforelasticsearch.sql.ast.tree.RareTopN;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Relation;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Rename;
 import com.amazon.opendistroforelasticsearch.sql.ast.tree.Sort;
@@ -296,4 +296,15 @@ public static Sort sort(UnresolvedPlan input, List<Argument> options, Field... s
   public static Dedupe dedupe(UnresolvedPlan input, List<Argument> options, Field... fields) {
     return new Dedupe(input, options, Arrays.asList(fields));
   }
+
+  public static List<Argument> defaultTopArgs() {
+    return exprList(
+        argument("noOfResults", intLiteral(10)));
+  }
+
+  public static RareTopN rareTopN(UnresolvedPlan input, Boolean rareTopFlag,
+      List<Argument> noOfResults, List<UnresolvedExpression> groupList, Field... fields) {
+    return new RareTopN(input, rareTopFlag, noOfResults, Arrays.asList(fields), groupList)
+        .attach(input);
+  }
 }

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

@@ -0,0 +1,49 @@
+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.Field;
+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.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import lombok.ToString;
+
+/**
+ * AST node represent RareTopN operation.
+ */
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+@AllArgsConstructor
+public class RareTopN extends UnresolvedPlan {
+
+  private UnresolvedPlan child;
+  private final Boolean rareTopFlag;
+  private final List<Argument> noOfResults;
+  private final List<Field> fields;
+  private final List<UnresolvedExpression> groupExprList;
+
+  @Override
+  public RareTopN attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return ImmutableList.of(this.child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitRareTopN(this, context);
+  }
+}
+

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

@@ -23,6 +23,7 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalPlanNodeVisitor;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalProject;
+import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRareTopN;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRelation;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRemove;
 import com.amazon.opendistroforelasticsearch.sql.planner.logical.LogicalRename;
@@ -34,6 +35,7 @@
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.FilterOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.ProjectOperator;
+import com.amazon.opendistroforelasticsearch.sql.planner.physical.RareTopNOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.RemoveOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.RenameOperator;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.SortOperator;
@@ -51,6 +53,17 @@
  */
 public class DefaultImplementor<C> extends LogicalPlanNodeVisitor<PhysicalPlan, C> {
 
+  @Override
+  public PhysicalPlan visitRareTopN(LogicalRareTopN node, C context) {
+    return new RareTopNOperator(
+        visitChild(node, context),
+        node.getRareTopFlag(),
+        node.getNoOfResults(),
+        node.getFieldList(),
+        node.getGroupByList()
+    );
+  }
+
   @Override
   public PhysicalPlan visitDedupe(LogicalDedupe node, C context) {
     return new DedupeOperator(

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

@@ -33,6 +33,7 @@
  */
 @UtilityClass
 public class LogicalPlanDSL {
+
   public static LogicalPlan aggregation(
       LogicalPlan input, List<Aggregator> aggregatorList, List<Expression> groupByList) {
     return new LogicalAggregation(input, aggregatorList, groupByList);
@@ -83,6 +84,16 @@ public static LogicalPlan dedupe(
         input, Arrays.asList(fields), allowedDuplication, keepEmpty, consecutive);
   }
 
+  public static LogicalPlan rareTopN(LogicalPlan input, Boolean rareTopFlag,
+      List<Expression> groupByList, Expression... fields) {
+    return rareTopN(input, rareTopFlag, 10, groupByList, fields);
+  }
+
+  public static LogicalPlan rareTopN(LogicalPlan input, boolean rareTopFlag, int noOfResults,
+      List<Expression> groupByList, Expression... fields) {
+    return new LogicalRareTopN(input, rareTopFlag, noOfResults, Arrays.asList(fields), groupByList);
+  }
+
   @SafeVarargs
   public LogicalPlan values(List<LiteralExpression>... values) {
     return new LogicalValues(Arrays.asList(values));

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

@@ -22,6 +22,7 @@
  * @param <C> context type.
  */
 public abstract class LogicalPlanNodeVisitor<R, C> {
+
   protected R visitNode(LogicalPlan plan, C context) {
     return null;
   }
@@ -65,4 +66,9 @@ public R visitSort(LogicalSort plan, C context) {
   public R visitValues(LogicalValues plan, C context) {
     return visitNode(plan, context);
   }
+
+  public R visitRareTopN(LogicalRareTopN plan, C context) {
+    return visitNode(plan, context);
+  }
+
 }

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

@@ -0,0 +1,36 @@
+package com.amazon.opendistroforelasticsearch.sql.planner.logical;
+
+import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+
+/**
+ * Logical Rare and TopN Plan.
+ */
+@Getter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class LogicalRareTopN extends LogicalPlan {
+
+  private final LogicalPlan child;
+  private final Boolean rareTopFlag;
+  private final Integer noOfResults;
+  private final List<Expression> fieldList;
+  private final List<Expression> groupByList;
+
+  @Override
+  public List<LogicalPlan> getChild() {
+    return Collections.singletonList(child);
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitRareTopN(this, context);
+  }
+}

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

@@ -80,6 +80,17 @@ public static DedupeOperator dedupe(
         input, Arrays.asList(expressions), allowedDuplication, keepEmpty, consecutive);
   }
 
+  public static RareTopNOperator rareTopN(PhysicalPlan input, Boolean rareTopFlag,
+      List<Expression> groups, Expression... expressions) {
+    return rareTopN(input, rareTopFlag, 10, groups, expressions);
+  }
+
+  public static RareTopNOperator rareTopN(PhysicalPlan input, Boolean rareTopFlag, int noOfResults,
+      List<Expression> groups, Expression... expressions) {
+    return new RareTopNOperator(input, rareTopFlag, noOfResults, Arrays.asList(expressions),
+        groups);
+  }
+
   @SafeVarargs
   public ValuesOperator values(List<LiteralExpression>... values) {
     return new ValuesOperator(Arrays.asList(values));

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

@@ -68,4 +68,8 @@ public R visitValues(ValuesOperator node, C context) {
   public R visitSort(SortOperator node, C context) {
     return visitNode(node, context);
   }
+
+  public R visitRareTopN(RareTopNOperator node, C context) {
+    return visitNode(node, context);
+  }
 }