Support top, rare commands with Calcite

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -382,7 +382,7 @@ public LogicalPlan visitRareTopN(RareTopN node, AnalysisContext context) {
     fields.forEach(
         field -> newEnv.define(new Symbol(Namespace.FIELD_NAME, field.toString()), field.type()));
 
-    List<Argument> options = node.getNoOfResults();
+    List<Argument> options = node.getArguments();
     Integer noOfResults = (Integer) options.get(0).getValue().getValue();
 
     return new LogicalRareTopN(child, node.getCommandType(), noOfResults, fields, groupBys);

core/src/main/java/org/opensearch/sql/ast/expression/Argument.java

@@ -7,6 +7,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -32,4 +33,22 @@ public List<UnresolvedExpression> getChild() {
   public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
     return nodeVisitor.visitArgument(this, context);
   }
+
+  public static class ArgumentMap {
+    private final Map<String, Literal> map;
+
+    public ArgumentMap(List<Argument> arguments) {
+      this.map =
+          arguments.stream()
+              .collect(java.util.stream.Collectors.toMap(Argument::getArgName, Argument::getValue));
+    }
+
+    public static ArgumentMap of(List<Argument> arguments) {
+      return new ArgumentMap(arguments);
+    }
+
+    public Literal get(String name) {
+      return map.get(name);
+    }
+  }
 }

core/src/main/java/org/opensearch/sql/ast/expression/WindowFrame.java

@@ -10,6 +10,7 @@
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
+import org.opensearch.sql.ast.dsl.AstDSL;
 
 @EqualsAndHashCode(callSuper = false)
 @Getter
@@ -25,12 +26,25 @@ public enum FrameType {
     ROWS
   }
 
-  public static WindowFrame defaultFrame() {
-    return new WindowFrame(
-        FrameType.ROWS, createBound("UNBOUNDED PRECEDING"), createBound("UNBOUNDED FOLLOWING"));
+  public static WindowFrame rowsUnbounded() {
+    return WindowFrame.of(
+        FrameType.ROWS,
+        AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
+        AstDSL.stringLiteral("UNBOUNDED FOLLOWING"));
   }
 
-  public static WindowFrame create(FrameType type, Literal lower, Literal upper) {
+  public static WindowFrame toCurrentRow() {
+    return WindowFrame.of(
+        FrameType.ROWS,
+        AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
+        AstDSL.stringLiteral("CURRENT ROW"));
+  }
+
+  public static WindowFrame of(FrameType type, String lower, String upper) {
+    return WindowFrame.of(type, AstDSL.stringLiteral(lower), AstDSL.stringLiteral(upper));
+  }
+
+  public static WindowFrame of(FrameType type, Literal lower, Literal upper) {
     WindowBound lowerBound = null;
     WindowBound upperBound = null;
     if (lower != null) {

core/src/main/java/org/opensearch/sql/ast/expression/WindowFunction.java

@@ -26,7 +26,7 @@ public class WindowFunction extends UnresolvedExpression {
   private final UnresolvedExpression function;
   @Setter private List<UnresolvedExpression> partitionByList = new ArrayList<>();
   @Setter private List<Pair<SortOption, UnresolvedExpression>> sortList = new ArrayList<>();
-  @Setter private WindowFrame windowFrame = WindowFrame.defaultFrame();
+  @Setter private WindowFrame windowFrame = WindowFrame.rowsUnbounded();
 
   public WindowFunction(
       UnresolvedExpression function,

core/src/main/java/org/opensearch/sql/ast/tree/RareTopN.java

@@ -29,7 +29,8 @@ public class RareTopN extends UnresolvedPlan {
 
   private UnresolvedPlan child;
   private final CommandType commandType;
-  private final List<Argument> noOfResults;
+  // arguments: noOfResults: Integer, countField: String, showCount: Boolean
+  private final List<Argument> arguments;
   private final List<Field> fields;
   private final List<UnresolvedExpression> groupExprList;
 

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -11,6 +11,7 @@
 import static org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
+import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME;
 
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
@@ -44,11 +45,13 @@
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta;
 import org.opensearch.sql.ast.expression.Argument;
+import org.opensearch.sql.ast.expression.Argument.ArgumentMap;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.ParseMethod;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.ast.expression.WindowFrame;
 import org.opensearch.sql.ast.expression.subquery.SubqueryExpression;
 import org.opensearch.sql.ast.tree.AD;
 import org.opensearch.sql.ast.tree.Aggregation;
@@ -82,6 +85,7 @@
 import org.opensearch.sql.calcite.utils.PlanUtils;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
 import org.opensearch.sql.utils.ParseUtils;
 
@@ -708,9 +712,77 @@ public RelNode visitFillNull(FillNull fillNull, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("FillNull command is unsupported in Calcite");
   }
 
+  List<RexNode> resolveGroupByPlusFieldList(RareTopN node, CalcitePlanContext context) {
+    List<RexNode> groupByList = resolveGroupByList(node, context);
+    List<RexNode> filedsList =
+        node.getFields().stream().map(g -> rexVisitor.analyze(g, context)).toList();
+    List<RexNode> all = new ArrayList<>(groupByList);
+    all.addAll(filedsList);
+    return all;
+  }
+
+  List<RexNode> resolveGroupByList(RareTopN node, CalcitePlanContext context) {
+    return node.getGroupExprList().stream().map(g -> rexVisitor.analyze(g, context)).toList();
+  }
+
   @Override
   public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
-    throw new CalciteUnsupportedException("Rare and Top commands are unsupported in Calcite");
+    visitChildren(node, context);
+
+    // 1. before aggregating, add a trim project
+    List<RexInputRef> trimmedRefs =
+        PlanUtils.getInputRefs(resolveGroupByPlusFieldList(node, context));
+    context.relBuilder.project(trimmedRefs);
+
+    ArgumentMap arguments = ArgumentMap.of(node.getArguments());
+    // 2. group the group-by list + field list and add a count() aggregation
+    List<RexNode> firstGroupBy = resolveGroupByPlusFieldList(node, context);
+    String countFieldName = (String) arguments.get("countField").getValue();
+    if (context.relBuilder.peek().getRowType().getFieldNames().contains(countFieldName)) {
+      throw new IllegalArgumentException(
+          "Field `"
+              + countFieldName
+              + "` is existed, change the count field by setting countfield='xyz'");
+    }
+    AggCall aggCall =
+        PlanUtils.makeAggCall(context, BuiltinFunctionName.COUNT, false, null, List.of())
+            .as(countFieldName);
+    context.relBuilder.aggregate(context.relBuilder.groupKey(firstGroupBy), aggCall);
+
+    // 3. add a window column
+    List<RexNode> partitionKeys = new ArrayList<>(resolveGroupByList(node, context));
+    RexNode countField;
+    if (node.getCommandType() == RareTopN.CommandType.TOP) {
+      countField = context.relBuilder.desc(context.relBuilder.field(countFieldName));
+    } else {
+      countField = context.relBuilder.field(countFieldName);
+    }
+    RexNode rowNumberWindowOver =
+        PlanUtils.makeOver(
+            context,
+            BuiltinFunctionName.ROW_NUMBER,
+            null,
+            List.of(),
+            partitionKeys,
+            List.of(countField),
+            WindowFrame.toCurrentRow());
+    context.relBuilder.projectPlus(
+        context.relBuilder.alias(rowNumberWindowOver, ROW_NUMBER_COLUMN_NAME));
+
+    // 4. filter row_number() <= k in each partition
+    Integer N = (Integer) arguments.get("noOfResults").getValue();
+    context.relBuilder.filter(
+        context.relBuilder.lessThanOrEqual(
+            context.relBuilder.field(ROW_NUMBER_COLUMN_NAME), context.relBuilder.literal(N)));
+
+    // 5. project final output. the default output is group by list + field list
+    List<RexNode> finalProjectList = new ArrayList<>(resolveGroupByPlusFieldList(node, context));
+    Boolean showCount = (Boolean) arguments.get("showCount").getValue();
+    if (showCount) {
+      finalProjectList.add(context.relBuilder.field(countFieldName));
+    }
+    context.relBuilder.project(finalProjectList);
+    return context.relBuilder.peek();
   }
 
   @Override

core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java

@@ -365,7 +365,7 @@ public RexNode visitWindowFunction(WindowFunction node, CalcitePlanContext conte
                       ? Collections.emptyList()
                       : arguments.subList(1, arguments.size());
               return PlanUtils.makeOver(
-                  context, functionName, field, args, partitions, node.getWindowFrame());
+                  context, functionName, field, args, partitions, List.of(), node.getWindowFrame());
             })
         .orElseThrow(
             () ->

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -39,6 +39,8 @@
 
 public interface PlanUtils {
 
+  String ROW_NUMBER_COLUMN_NAME = "_row_number_";
+
   static SpanUnit intervalUnitToSpanUnit(IntervalUnit unit) {
     return switch (unit) {
       case MICROSECOND -> SpanUnit.MILLISECOND;
@@ -61,9 +63,10 @@ static RexNode makeOver(
       RexNode field,
       List<RexNode> argList,
       List<RexNode> partitions,
+      List<RexNode> orderKeys,
       @Nullable WindowFrame windowFrame) {
     if (windowFrame == null) {
-      windowFrame = WindowFrame.defaultFrame();
+      windowFrame = WindowFrame.rowsUnbounded();
     }
     boolean rows = windowFrame.getType() == WindowFrame.FrameType.ROWS;
     RexWindowBound lowerBound = convert(context, windowFrame.getLower());
@@ -99,6 +102,14 @@ static RexNode makeOver(
         return variance(context, field, partitions, rows, lowerBound, upperBound, true, false);
       case VARSAMP:
         return variance(context, field, partitions, rows, lowerBound, upperBound, false, false);
+      case ROW_NUMBER:
+        return withOver(
+            context.relBuilder.aggregateCall(SqlStdOperatorTable.ROW_NUMBER),
+            partitions,
+            orderKeys,
+            true,
+            lowerBound,
+            upperBound);
       default:
         return withOver(
             makeAggCall(context, functionName, false, field, argList),
@@ -151,6 +162,25 @@ private static RexNode withOver(
         .toRex();
   }
 
+  private static RexNode withOver(
+      RelBuilder.AggCall aggCall,
+      List<RexNode> partitions,
+      List<RexNode> orderKeys,
+      boolean rows,
+      RexWindowBound lowerBound,
+      RexWindowBound upperBound) {
+    return aggCall
+        .over()
+        .partitionBy(partitions)
+        .orderBy(orderKeys)
+        .let(
+            c ->
+                rows
+                    ? c.rowsBetween(lowerBound, upperBound)
+                    : c.rangeBetween(lowerBound, upperBound))
+        .toRex();
+  }
+
   private static RexNode variance(
       CalcitePlanContext ctx,
       RexNode operator,

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -330,6 +330,7 @@ public enum BuiltinFunctionName {
           .put("stddev", BuiltinFunctionName.STDDEV_POP)
           .put("stddev_pop", BuiltinFunctionName.STDDEV_POP)
           .put("stddev_samp", BuiltinFunctionName.STDDEV_SAMP)
+          .put("rank", BuiltinFunctionName.RANK)
           .build();
 
   public static Optional<BuiltinFunctionName> of(String str) {

docs/user/ppl/cmd/rare.rst

@@ -19,8 +19,13 @@ Syntax
 ============
 rare <field-list> [by-clause]
 
+rare <field-list> [rare-options] [by-clause] ``(available from 3.1.0+)``
+
 * field-list: mandatory. comma-delimited list of field names.
 * by-clause: optional. one or more fields to group the results by.
+* rare-options: optional. options for the rare command. Supported syntax is [countfield=<string>] [showcount=<bool>].
+* showcount=<bool>: optional. whether to create a field in output that represent a count of the tuple of values. Default value is ``true``.
+* countfield=<string>: optional. the name of the field that contains count. Default value is ``'count'``.
 
 
 Example 1: Find the least common values in a field
@@ -58,6 +63,39 @@ PPL query::
     | M      | 36  |
     +--------+-----+
 
+Example 3: Rare command with Calcite enabled
+============================================
+
+The example finds least common gender of all the accounts when ``plugins.calcite.enabled`` is true.
+
+PPL query::
+
+    PPL> source=accounts | rare gender;
+    fetched row
+    +--------+-------+
+    | gender | count |
+    |--------+-------|
+    | F      | 1     |
+    | M      | 3     |
+    +--------+-------+
+
+
+Example 4: Specify the count field option
+=========================================
+
+The example specifies the count field when ``plugins.calcite.enabled`` is true.
+
+PPL query::
+
+    PPL> source=accounts | rare countfield='cnt' gender;
+    fetched row
+    +--------+-----+
+    | gender | cnt |
+    |--------+-----|
+    | F      | 1   |
+    | M      | 3   |
+    +--------+-----+
+
 Limitation
 ==========
 The ``rare`` command is not rewritten to OpenSearch DSL, it is only executed on the coordination node.