Merge remote-tracking branch 'upstream/main' into issues/3614

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,29 @@ public List<UnresolvedExpression> getChild() {
   public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
     return nodeVisitor.visitArgument(this, context);
   }
+
+  /** ArgumentMap is a helper class to get argument value by name. */
+  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);
+    }
+
+    /**
+     * Get argument value by name.
+     *
+     * @param name argument name
+     * @return argument value
+     */
+    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/CalciteAggCallVisitor.java

@@ -37,7 +37,8 @@ public AggCall visitAlias(Alias node, CalcitePlanContext context) {
 
   @Override
   public AggCall visitAggregateFunction(AggregateFunction node, CalcitePlanContext context) {
-    RexNode field = rexNodeVisitor.analyze(node.getField(), context);
+    RexNode field =
+        node.getField() == null ? null : rexNodeVisitor.analyze(node.getField(), context);
     List<RexNode> argList = new ArrayList<>();
     for (UnresolvedExpression arg : node.getArgList()) {
       argList.add(rexNodeVisitor.analyze(arg, context));

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

@@ -31,6 +31,7 @@ public class CalcitePlanContext {
   public final ExtendedRexBuilder rexBuilder;
   public final FunctionProperties functionProperties;
   public final QueryType queryType;
+  public final Integer querySizeLimit;
 
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
@@ -46,8 +47,9 @@ public class CalcitePlanContext {
   private final Stack<RexCorrelVariable> correlVar = new Stack<>();
   private final Stack<List<RexNode>> windowPartitions = new Stack<>();
 
-  private CalcitePlanContext(FrameworkConfig config, QueryType queryType) {
+  private CalcitePlanContext(FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
     this.config = config;
+    this.querySizeLimit = querySizeLimit;
     this.queryType = queryType;
     this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
     this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
@@ -84,7 +86,8 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
     }
   }
 
-  public static CalcitePlanContext create(FrameworkConfig config, QueryType queryType) {
-    return new CalcitePlanContext(config, queryType);
+  public static CalcitePlanContext create(
+      FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
+    return new CalcitePlanContext(config, querySizeLimit, queryType);
   }
 }

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;
@@ -26,6 +27,7 @@
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.ViewExpanders;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexCorrelVariable;
@@ -42,14 +44,17 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.Node;
+import org.opensearch.sql.ast.dsl.AstDSL;
 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;
@@ -83,6 +88,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;
 
@@ -374,70 +380,102 @@ private void projectPlusOverriding(
     context.relBuilder.rename(expectedRenameFields);
   }
 
-  private Pair<List<AggCall>, List<RexNode>> resolveAggCallAndGroupBy(
-      Aggregation node, CalcitePlanContext context) {
+  /**
+   * Resolve the aggregation with trimming unused fields to avoid bugs in {@link
+   * org.apache.calcite.sql2rel.RelDecorrelator#decorrelateRel(Aggregate, boolean)}
+   *
+   * @param groupExprList group by expression list
+   * @param aggExprList aggregate expression list
+   * @param context CalcitePlanContext
+   * @return Pair of (group-by list, field list, aggregate list)
+   */
+  private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
+      List<UnresolvedExpression> groupExprList,
+      List<UnresolvedExpression> aggExprList,
+      CalcitePlanContext context) {
+    // Example 1: source=t | where a > 1 | stats avg(b + 1) by c
+    // Before: Aggregate(avg(b + 1))
+    //         \- Filter(a > 1)
+    //            \- Scan t
+    // After: Aggregate(avg(b + 1))
+    //        \- Project([c, b])
+    //           \- Filter(a > 1)
+    //              \- Scan t
+    //
+    // Example 2: source=t | where a > 1 | top b by c
+    // Before: Aggregate(count)
+    //         \-Filter(a > 1)
+    //           \- Scan t
+    // After: Aggregate(count)
+    //        \- Project([c, b])
+    //           \- Filter(a > 1)
+    //              \- Scan t
+    Pair<List<RexNode>, List<AggCall>> resolved =
+        resolveAttributesForAggregation(groupExprList, aggExprList, context);
+    List<RexInputRef> trimmedRefs = new ArrayList<>();
+    trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.getLeft())); // group-by keys first
+    trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.getRight()));
+    context.relBuilder.project(trimmedRefs);
+
+    // Re-resolve all attributes based on adding trimmed Project.
+    // Using re-resolving rather than Calcite Mapping (ref Calcite ProjectTableScanRule)
+    // because that Mapping only works for RexNode, but we need both AggCall and RexNode list.
+    Pair<List<RexNode>, List<AggCall>> reResolved =
+        resolveAttributesForAggregation(groupExprList, aggExprList, context);
+    context.relBuilder.aggregate(
+        context.relBuilder.groupKey(reResolved.getLeft()), reResolved.getRight());
+    return Pair.of(reResolved.getLeft(), reResolved.getRight());
+  }
+
+  /**
+   * Resolve attributes for aggregation.
+   *
+   * @param groupExprList group by expression list
+   * @param aggExprList aggregate expression list
+   * @param context CalcitePlanContext
+   * @return Pair of (group-by list, aggregate list)
+   */
+  private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
+      List<UnresolvedExpression> groupExprList,
+      List<UnresolvedExpression> aggExprList,
+      CalcitePlanContext context) {
     List<AggCall> aggCallList =
-        node.getAggExprList().stream()
-            .map(expr -> aggVisitor.analyze(expr, context))
-            .collect(Collectors.toList());
-    // The span column is always the first column in result whatever
-    // the order of span in query is first or last one
-    List<RexNode> groupByList = new ArrayList<>();
-    UnresolvedExpression span = node.getSpan();
-    if (!Objects.isNull(span)) {
-      RexNode spanRex = rexVisitor.analyze(span, context);
-      groupByList.add(spanRex);
-      // add span's group alias field (most recent added expression)
-    }
-    groupByList.addAll(
-        node.getGroupExprList().stream().map(expr -> rexVisitor.analyze(expr, context)).toList());
-    return Pair.of(aggCallList, groupByList);
+        aggExprList.stream().map(expr -> aggVisitor.analyze(expr, context)).toList();
+    List<RexNode> groupByList =
+        groupExprList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
+    return Pair.of(groupByList, aggCallList);
   }
 
   @Override
   public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     visitChildren(node, context);
-    // Add a trimmed Project before Aggregate.
-    // to avoid bugs in RelDecorrelator.decorrelateRel(Aggregate rel)
-    // For example:
-    // source=t | where a > 1 | stats avg(b+1) by c
-    // Before:
-    // Aggregate
-    //     \- Filter(a>1)
-    //            \- Scan t
-    // After:
-    // Aggregate
-    //     \- Project([c,b])
-    //            \- Filter(a>1)
-    //                   \- Scan t
-    Pair<List<AggCall>, List<RexNode>> resolved = resolveAggCallAndGroupBy(node, context);
-    List<RexInputRef> trimmedRefs = new ArrayList<>();
-    trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.getRight())); // group-by keys first
-    trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.getLeft()));
-    context.relBuilder.project(trimmedRefs);
 
-    // Re-resolve aggCalls and group-by list based on adding trimmed Project.
-    // Using re-resolving rather than Calcite Mapping (ref Calcite ProjectTableScanRule)
-    // because that Mapping only works for RexNode, but we need both AggCall and RexNode list.
-    Pair<List<AggCall>, List<RexNode>> reResolved = resolveAggCallAndGroupBy(node, context);
-    List<AggCall> aggList = reResolved.getLeft();
-    List<RexNode> groupByList = reResolved.getRight();
-    context.relBuilder.aggregate(context.relBuilder.groupKey(groupByList), aggList);
+    List<UnresolvedExpression> aggExprList = node.getAggExprList();
+    List<UnresolvedExpression> groupExprList = new ArrayList<>();
+    // The span column is always the first column in result whatever
+    // the order of span in query is first or last one
+    UnresolvedExpression span = node.getSpan();
+    if (!Objects.isNull(span)) {
+      groupExprList.add(span);
+    }
+    groupExprList.addAll(node.getGroupExprList());
+    Pair<List<RexNode>, List<AggCall>> aggregationAttributes =
+        aggregateWithTrimming(groupExprList, aggExprList, context);
 
     // schema reordering
     // As an example, in command `stats count() by colA, colB`,
     // the sequence of output schema is "count, colA, colB".
     List<RexNode> outputFields = context.relBuilder.fields();
     int numOfOutputFields = outputFields.size();
-    int numOfAggList = aggList.size();
+    int numOfAggList = aggExprList.size();
     List<RexNode> reordered = new ArrayList<>(numOfOutputFields);
     // Add aggregation results first
     List<RexNode> aggRexList =
         outputFields.subList(numOfOutputFields - numOfAggList, numOfOutputFields);
     reordered.addAll(aggRexList);
     // Add group by columns
     List<RexNode> aliasedGroupByList =
-        groupByList.stream()
+        aggregationAttributes.getLeft().stream()
             .map(this::extractAliasLiteral)
             .flatMap(Optional::stream)
             .map(ref -> ((RexLiteral) ref).getValueAs(String.class))
@@ -742,7 +780,62 @@ public RelNode visitKmeans(Kmeans node, CalcitePlanContext context) {
 
   @Override
   public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
-    throw new CalciteUnsupportedException("Rare and Top commands are unsupported in Calcite");
+    visitChildren(node, context);
+
+    ArgumentMap arguments = ArgumentMap.of(node.getArguments());
+    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'");
+    }
+
+    // 1. group the group-by list + field list and add a count() aggregation
+    List<UnresolvedExpression> groupExprList = new ArrayList<>(node.getGroupExprList());
+    List<UnresolvedExpression> fieldList =
+        node.getFields().stream().map(f -> (UnresolvedExpression) f).toList();
+    groupExprList.addAll(fieldList);
+    List<UnresolvedExpression> aggExprList =
+        List.of(AstDSL.alias(countFieldName, AstDSL.aggregate("count", null)));
+    aggregateWithTrimming(groupExprList, aggExprList, context);
+
+    // 2. add a window column
+    List<RexNode> partitionKeys = rexVisitor.analyze(node.getGroupExprList(), 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));
+
+    // 3. 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)));
+
+    // 4. project final output. the default output is group by list + field list
+    Boolean showCount = (Boolean) arguments.get("showCount").getValue();
+    if (showCount) {
+      context.relBuilder.projectExcept(context.relBuilder.field(ROW_NUMBER_COLUMN_NAME));
+    } else {
+      context.relBuilder.projectExcept(
+          context.relBuilder.field(ROW_NUMBER_COLUMN_NAME),
+          context.relBuilder.field(countFieldName));
+    }
+    return context.relBuilder.peek();
   }
 
   @Override