[Multi-stage] Support is_enable_group_trim agg option

pinot-common/src/main/proto/plan.proto

@@ -69,6 +69,8 @@ message AggregateNode {
   repeated int32 groupKeys = 3;
   AggType aggType = 4;
   bool leafReturnFinalResult = 5;
+  repeated Collation collations = 6;
+  int32 limit = 7;
 }
 
 message FilterNode {

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java

@@ -42,7 +42,8 @@ private PinotHintOptions() {
   public static class AggregateOptions {
     public static final String IS_PARTITIONED_BY_GROUP_BY_KEYS = "is_partitioned_by_group_by_keys";
     public static final String IS_LEAF_RETURN_FINAL_RESULT = "is_leaf_return_final_result";
-    public static final String SKIP_LEAF_STAGE_GROUP_BY_AGGREGATION = "is_skip_leaf_stage_group_by";
+    public static final String IS_SKIP_LEAF_STAGE_GROUP_BY = "is_skip_leaf_stage_group_by";
+    public static final String IS_ENABLE_GROUP_TRIM = "is_enable_group_trim";
 
     public static final String NUM_GROUPS_LIMIT = "num_groups_limit";
     public static final String MAX_INITIAL_RESULT_HOLDER_CAPACITY = "max_initial_result_holder_capacity";

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalAggregate.java

@@ -22,6 +22,7 @@
 import javax.annotation.Nullable;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelWriter;
 import org.apache.calcite.rel.core.Aggregate;
@@ -35,39 +36,36 @@ public class PinotLogicalAggregate extends Aggregate {
   private final AggType _aggType;
   private final boolean _leafReturnFinalResult;
 
+  // The following fields are set when group trim is enabled, and are extracted from the Sort on top of this Aggregate.
+  private final List<RelFieldCollation> _collations;
+  private final int _limit;
+
   public PinotLogicalAggregate(RelOptCluster cluster, RelTraitSet traitSet, List<RelHint> hints, RelNode input,
       ImmutableBitSet groupSet, @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls,
-      AggType aggType, boolean leafReturnFinalResult) {
+      AggType aggType, boolean leafReturnFinalResult, @Nullable List<RelFieldCollation> collations, int limit) {
     super(cluster, traitSet, hints, input, groupSet, groupSets, aggCalls);
     _aggType = aggType;
     _leafReturnFinalResult = leafReturnFinalResult;
+    _collations = collations;
+    _limit = limit;
   }
 
-  public PinotLogicalAggregate(RelOptCluster cluster, RelTraitSet traitSet, List<RelHint> hints, RelNode input,
-      ImmutableBitSet groupSet, @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls,
-      AggType aggType) {
-    this(cluster, traitSet, hints, input, groupSet, groupSets, aggCalls, aggType, false);
-  }
-
-  public PinotLogicalAggregate(Aggregate aggRel, List<AggregateCall> aggCalls, AggType aggType,
-      boolean leafReturnFinalResult) {
-    this(aggRel.getCluster(), aggRel.getTraitSet(), aggRel.getHints(), aggRel.getInput(), aggRel.getGroupSet(),
-        aggRel.getGroupSets(), aggCalls, aggType, leafReturnFinalResult);
+  public PinotLogicalAggregate(Aggregate aggRel, RelNode input, ImmutableBitSet groupSet,
+      @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls, AggType aggType,
+      boolean leafReturnFinalResult, @Nullable List<RelFieldCollation> collations, int limit) {
+    this(aggRel.getCluster(), aggRel.getTraitSet(), aggRel.getHints(), input, groupSet, groupSets, aggCalls, aggType,
+        leafReturnFinalResult, collations, limit);
   }
 
-  public PinotLogicalAggregate(Aggregate aggRel, List<AggregateCall> aggCalls, AggType aggType) {
-    this(aggRel, aggCalls, aggType, false);
-  }
-
-  public PinotLogicalAggregate(Aggregate aggRel, RelNode input, List<AggregateCall> aggCalls, AggType aggType) {
-    this(aggRel.getCluster(), aggRel.getTraitSet(), aggRel.getHints(), input, aggRel.getGroupSet(),
-        aggRel.getGroupSets(), aggCalls, aggType);
+  public PinotLogicalAggregate(Aggregate aggRel, RelNode input, List<AggregateCall> aggCalls, AggType aggType,
+      boolean leafReturnFinalResult, @Nullable List<RelFieldCollation> collations, int limit) {
+    this(aggRel, input, aggRel.getGroupSet(), aggRel.getGroupSets(), aggCalls, aggType,
+        leafReturnFinalResult, collations, limit);
   }
 
   public PinotLogicalAggregate(Aggregate aggRel, RelNode input, ImmutableBitSet groupSet, List<AggregateCall> aggCalls,
-      AggType aggType, boolean leafReturnFinalResult) {
-    this(aggRel.getCluster(), aggRel.getTraitSet(), aggRel.getHints(), input, groupSet, null, aggCalls, aggType,
-        leafReturnFinalResult);
+      AggType aggType, boolean leafReturnFinalResult, @Nullable List<RelFieldCollation> collations, int limit) {
+    this(aggRel, input, groupSet, null, aggCalls, aggType, leafReturnFinalResult, collations, limit);
   }
 
   public AggType getAggType() {
@@ -78,24 +76,35 @@ public boolean isLeafReturnFinalResult() {
     return _leafReturnFinalResult;
   }
 
+  @Nullable
+  public List<RelFieldCollation> getCollations() {
+    return _collations;
+  }
+
+  public int getLimit() {
+    return _limit;
+  }
+
   @Override
   public PinotLogicalAggregate copy(RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet,
       @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
     return new PinotLogicalAggregate(getCluster(), traitSet, hints, input, groupSet, groupSets, aggCalls, _aggType,
-        _leafReturnFinalResult);
+        _leafReturnFinalResult, _collations, _limit);
   }
 
   @Override
   public RelWriter explainTerms(RelWriter pw) {
     RelWriter relWriter = super.explainTerms(pw);
     relWriter.item("aggType", _aggType);
     relWriter.itemIf("leafReturnFinalResult", true, _leafReturnFinalResult);
+    relWriter.itemIf("collations", _collations, _collations != null);
+    relWriter.itemIf("limit", _limit, _limit > 0);
     return relWriter;
   }
 
   @Override
   public RelNode withHints(List<RelHint> hintList) {
     return new PinotLogicalAggregate(getCluster(), traitSet, hintList, input, groupSet, groupSets, aggCalls, _aggType,
-        _leafReturnFinalResult);
+        _leafReturnFinalResult, _collations, _limit);
   }
 }