address comments

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

Author: LantaoJin

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/CalciteRelNodeVisitor.java

@@ -27,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;
@@ -43,6 +44,7 @@
 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;
@@ -378,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))
@@ -744,45 +778,30 @@ public RelNode visitKmeans(Kmeans node, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("Kmeans 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) {
     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));
+    // 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));
@@ -801,19 +820,21 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
     context.relBuilder.projectPlus(
         context.relBuilder.alias(rowNumberWindowOver, ROW_NUMBER_COLUMN_NAME));
 
-    // 4. filter row_number() <= k in each partition
+    // 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)));
 
-    // 5. project final output. the default output is group by list + field list
-    List<RexNode> finalProjectList = new ArrayList<>(resolveGroupByPlusFieldList(node, context));
+    // 4. project final output. the default output is group by list + field list
     Boolean showCount = (Boolean) arguments.get("showCount").getValue();
     if (showCount) {
-      finalProjectList.add(context.relBuilder.field(countFieldName));
+      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));
     }
-    context.relBuilder.project(finalProjectList);
     return context.relBuilder.peek();
   }
 

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

@@ -74,6 +74,10 @@ public RexNode analyze(UnresolvedExpression unresolved, CalcitePlanContext conte
     return unresolved.accept(this, context);
   }
 
+  public List<RexNode> analyze(List<UnresolvedExpression> list, CalcitePlanContext context) {
+    return list.stream().map(u -> u.accept(this, context)).toList();
+  }
+
   public RexNode analyzeJoinCondition(UnresolvedExpression unresolved, CalcitePlanContext context) {
     return context.resolveJoinCondition(unresolved, this::analyze);
   }

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

@@ -114,6 +114,7 @@ static RexNode makeOver(
         return withOver(
             makeAggCall(context, functionName, false, field, argList),
             partitions,
+            orderKeys,
             rows,
             lowerBound,
             upperBound);
@@ -127,7 +128,8 @@ private static RexNode sumOver(
       boolean rows,
       RexWindowBound lowerBound,
       RexWindowBound upperBound) {
-    return withOver(ctx.relBuilder.sum(operation), partitions, rows, lowerBound, upperBound);
+    return withOver(
+        ctx.relBuilder.sum(operation), partitions, List.of(), rows, lowerBound, upperBound);
   }
 
   private static RexNode countOver(
@@ -140,28 +142,12 @@ private static RexNode countOver(
     return withOver(
         ctx.relBuilder.count(ImmutableList.of(operation)),
         partitions,
+        List.of(),
         rows,
         lowerBound,
         upperBound);
   }
 
-  private static RexNode withOver(
-      RelBuilder.AggCall aggCall,
-      List<RexNode> partitions,
-      boolean rows,
-      RexWindowBound lowerBound,
-      RexWindowBound upperBound) {
-    return aggCall
-        .over()
-        .partitionBy(partitions)
-        .let(
-            c ->
-                rows
-                    ? c.rowsBetween(lowerBound, upperBound)
-                    : c.rangeBetween(lowerBound, upperBound))
-        .toRex();
-  }
-
   private static RexNode withOver(
       RelBuilder.AggCall aggCall,
       List<RexNode> partitions,