[feat](nereids)push Limit to local agg

fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java

@@ -294,7 +294,8 @@ public PlanFragment visitPhysicalDistribute(PhysicalDistribute<? extends Plan> d
                 && context.getFirstAggregateInFragment(inputFragment) == child) {
             PhysicalHashAggregate<?> hashAggregate = (PhysicalHashAggregate<?>) child;
             if (hashAggregate.getAggPhase() == AggPhase.LOCAL
-                    && hashAggregate.getAggMode() == AggMode.INPUT_TO_BUFFER) {
+                    && hashAggregate.getAggMode() == AggMode.INPUT_TO_BUFFER
+                    && hashAggregate.getTopnPushInfo() == null) {
                 AggregationNode aggregationNode = (AggregationNode) inputFragment.getPlanRoot();
                 aggregationNode.setUseStreamingPreagg(hashAggregate.isMaybeUsingStream());
             }
@@ -1058,6 +1059,23 @@ public PlanFragment visitPhysicalHashAggregate(
             // local exchanger will be used.
             aggregationNode.setColocate(true);
         }
+        if (aggregate.getTopnPushInfo() != null) {
+            List<Expr> orderingExprs = Lists.newArrayList();
+            List<Boolean> ascOrders = Lists.newArrayList();
+            List<Boolean> nullsFirstParams = Lists.newArrayList();
+            aggregate.getTopnPushInfo().orderkeys.forEach(k -> {
+                orderingExprs.add(ExpressionTranslator.translate(k.getExpr(), context));
+                ascOrders.add(k.isAsc());
+                nullsFirstParams.add(k.isNullFirst());
+            });
+            SortInfo sortInfo = new SortInfo(orderingExprs, ascOrders, nullsFirstParams, outputTupleDesc);
+            aggregationNode.setSortByGroupKey(sortInfo);
+            if (aggregationNode.getLimit() == -1) {
+                aggregationNode.setLimit(aggregate.getTopnPushInfo().limit);
+            }
+        } else {
+            aggregationNode.setSortByGroupKey(null);
+        }
         setPlanRoot(inputPlanFragment, aggregationNode, aggregate);
         if (aggregate.getStats() != null) {
             aggregationNode.setCardinality((long) aggregate.getStats().getRowCount());

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -84,6 +84,7 @@
 import org.apache.doris.nereids.rules.rewrite.InferPredicates;
 import org.apache.doris.nereids.rules.rewrite.InferSetOperatorDistinct;
 import org.apache.doris.nereids.rules.rewrite.InlineLogicalView;
+import org.apache.doris.nereids.rules.rewrite.LimitAggToTopNAgg;
 import org.apache.doris.nereids.rules.rewrite.LimitSortToTopN;
 import org.apache.doris.nereids.rules.rewrite.LogicalResultSinkToShortCircuitPointQuery;
 import org.apache.doris.nereids.rules.rewrite.MergeAggregate;
@@ -366,6 +367,7 @@ public class Rewriter extends AbstractBatchJobExecutor {
                         //       generate one PhysicalLimit if current distribution is gather or two
                         //       PhysicalLimits with gather exchange
                         topDown(new LimitSortToTopN()),
+                        topDown(new LimitAggToTopNAgg()),
                         topDown(new MergeTopNs()),
                         topDown(new SplitLimit()),
                         topDown(

fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/PlanPostProcessors.java

@@ -65,6 +65,9 @@ public List<PlanPostProcessor> getProcessors() {
         builder.add(new AddOffsetIntoDistribute());
         builder.add(new CommonSubExpressionOpt());
         // DO NOT replace PLAN NODE from here
+        if (cascadesContext.getConnectContext().getSessionVariable().pushTopnToAgg) {
+            builder.add(new PushTopnToAgg());
+        }
         builder.add(new TopNScanOpt());
         builder.add(new FragmentProcessor());
         if (!cascadesContext.getConnectContext().getSessionVariable().getRuntimeFilterMode()

fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/PushTopnToAgg.java

@@ -0,0 +1,177 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+// This file is copied from
+// https://github.com/apache/impala/blob/branch-2.9.0/fe/src/main/java/org/apache/impala/AggregationNode.java
+// and modified by Doris
+
+package org.apache.doris.nereids.processor.post;
+
+import org.apache.doris.nereids.CascadesContext;
+import org.apache.doris.nereids.properties.DistributionSpecGather;
+import org.apache.doris.nereids.properties.OrderKey;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.plans.AggMode;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalDistribute;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate.TopnPushInfo;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalLimit;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalProject;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalTopN;
+import org.apache.doris.qe.ConnectContext;
+
+import org.apache.hadoop.util.Lists;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Add SortInfo to Agg. This SortInfo is used as boundary, not used to sort elements.
+ * example
+ * sql: select count(*) from orders group by o_clerk order by o_clerk limit 1;
+ * plan: topn(1) -> aggGlobal -> shuffle -> aggLocal -> scan
+ * optimization: aggLocal and aggGlobal only need to generate the smallest row with respect to o_clerk.
+ *
+ * TODO: the following case is not covered:
+ * sql: select sum(o_shippriority) from orders group by o_clerk limit 1;
+ * plan: limit -> aggGlobal -> shuffle -> aggLocal -> scan
+ * aggGlobal may receive partial aggregate results, and hence is not supported now
+ * instance1: input (key=2, v=1) => localAgg => (2, 1) => aggGlobal inst1 => (2, 1)
+ * instance2: input (key=1, v=1), (key=2, v=2) => localAgg inst2 => (1, 1)
+ * (2,1),(1,1) => limit => may output (2, 1), which is not complete, missing (2, 2) in instance2
+ *
+ *TOPN:
+ *  Precondition: topn orderkeys are the prefix of group keys
+ *  TODO: topnKeys could be subset of groupKeys. This will be implemented in future
+ *  Pattern 2-phase agg:
+ *     topn -> aggGlobal -> distribute -> aggLocal
+ *     =>
+ *     topn(n) -> aggGlobal(topn=n) -> distribute -> aggLocal(topn=n)
+ *  Pattern 1-phase agg:
+ *     topn->agg->Any(not agg) -> topn -> agg(topn=n) -> any
+ *
+ *  LIMIT:
+ *  Pattern 1: limit->agg(1phase)->any
+ *  Pattern 2: limit->agg(global)->gather->agg(local)
+ */
+public class PushTopnToAgg extends PlanPostProcessor {
+    @Override
+    public Plan visitPhysicalTopN(PhysicalTopN<? extends Plan> topN, CascadesContext ctx) {
+        topN.child().accept(this, ctx);
+        if (ConnectContext.get().getSessionVariable().topnOptLimitThreshold <= topN.getLimit() + topN.getOffset()) {
+            return topN;
+        }
+        Plan topnChild = topN.child();
+        if (topnChild instanceof PhysicalProject) {
+            topnChild = topnChild.child(0);
+        }
+        if (topnChild instanceof PhysicalHashAggregate) {
+            PhysicalHashAggregate<? extends Plan> upperAgg = (PhysicalHashAggregate<? extends Plan>) topnChild;
+            List<OrderKey> orderKeys = tryGenerateOrderKeyByGroupKeyAndTopnKey(topN, upperAgg);
+            if (!orderKeys.isEmpty()) {
+
+                if (upperAgg.getAggPhase().isGlobal() && upperAgg.getAggMode() == AggMode.BUFFER_TO_RESULT) {
+                    upperAgg.setTopnPushInfo(new TopnPushInfo(
+                            orderKeys,
+                            topN.getLimit() + topN.getOffset()));
+                    if (upperAgg.child() instanceof PhysicalDistribute
+                            && upperAgg.child().child(0) instanceof PhysicalHashAggregate) {
+                        PhysicalHashAggregate<? extends Plan> bottomAgg =
+                                (PhysicalHashAggregate<? extends Plan>) upperAgg.child().child(0);
+                        bottomAgg.setTopnPushInfo(new TopnPushInfo(
+                                orderKeys,
+                                topN.getLimit() + topN.getOffset()));
+                    }
+                } else if (upperAgg.getAggPhase().isLocal() && upperAgg.getAggMode() == AggMode.INPUT_TO_RESULT) {
+                    // one phase agg
+                    upperAgg.setTopnPushInfo(new TopnPushInfo(
+                            orderKeys,
+                            topN.getLimit() + topN.getOffset()));
+                }
+            }
+        }
+        return topN;
+    }
+
+    /**
+     return true, if topn order-key is prefix of agg group-key, ignore asc/desc and null_first
+     TODO order-key can be subset of group-key. BE does not support now.
+     */
+    private List<OrderKey> tryGenerateOrderKeyByGroupKeyAndTopnKey(PhysicalTopN<? extends Plan> topN,
+            PhysicalHashAggregate<? extends Plan> agg) {
+        List<OrderKey> orderKeys = Lists.newArrayListWithCapacity(agg.getGroupByExpressions().size());
+        if (topN.getOrderKeys().size() > agg.getGroupByExpressions().size()) {
+            return orderKeys;
+        }
+        List<Expression> topnKeys = topN.getOrderKeys().stream()
+                .map(OrderKey::getExpr).collect(Collectors.toList());
+        for (int i = 0; i < topN.getOrderKeys().size(); i++) {
+            // prefix check
+            if (!topnKeys.get(i).equals(agg.getGroupByExpressions().get(i))) {
+                return Lists.newArrayList();
+            }
+            orderKeys.add(topN.getOrderKeys().get(i));
+        }
+        for (int i = topN.getOrderKeys().size(); i < agg.getGroupByExpressions().size(); i++) {
+            orderKeys.add(new OrderKey(agg.getGroupByExpressions().get(i), true, false));
+        }
+        return orderKeys;
+    }
+
+    @Override
+    public Plan visitPhysicalLimit(PhysicalLimit<? extends Plan> limit, CascadesContext ctx) {
+        limit.child().accept(this, ctx);
+        if (ConnectContext.get().getSessionVariable().topnOptLimitThreshold <= limit.getLimit() + limit.getOffset()) {
+            return limit;
+        }
+        Plan limitChild = limit.child();
+        if (limitChild instanceof PhysicalProject) {
+            limitChild = limitChild.child(0);
+        }
+        if (limitChild instanceof PhysicalHashAggregate) {
+            PhysicalHashAggregate<? extends Plan> upperAgg = (PhysicalHashAggregate<? extends Plan>) limitChild;
+            if (upperAgg.getAggPhase().isGlobal() && upperAgg.getAggMode() == AggMode.BUFFER_TO_RESULT) {
+                Plan child = upperAgg.child();
+                Plan grandChild = child.child(0);
+                if (child instanceof PhysicalDistribute
+                        && ((PhysicalDistribute<?>) child).getDistributionSpec() instanceof DistributionSpecGather
+                        && grandChild instanceof PhysicalHashAggregate) {
+                    upperAgg.setTopnPushInfo(new TopnPushInfo(
+                            generateOrderKeyByGroupKey(upperAgg),
+                            limit.getLimit() + limit.getOffset()));
+                    PhysicalHashAggregate<? extends Plan> bottomAgg =
+                            (PhysicalHashAggregate<? extends Plan>) grandChild;
+                    bottomAgg.setTopnPushInfo(new TopnPushInfo(
+                            generateOrderKeyByGroupKey(bottomAgg),
+                            limit.getLimit() + limit.getOffset()));
+                }
+            } else if (upperAgg.getAggMode() == AggMode.INPUT_TO_RESULT) {
+                // 1-phase agg
+                upperAgg.setTopnPushInfo(new TopnPushInfo(
+                        generateOrderKeyByGroupKey(upperAgg),
+                        limit.getLimit() + limit.getOffset()));
+            }
+        }
+        return limit;
+    }
+
+    private List<OrderKey> generateOrderKeyByGroupKey(PhysicalHashAggregate<? extends Plan> agg) {
+        return agg.getGroupByExpressions().stream()
+                .map(key -> new OrderKey(key, true, false))
+                .collect(Collectors.toList());
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java

@@ -304,6 +304,7 @@ public enum RuleType {
     PUSH_LIMIT_THROUGH_UNION(RuleTypeClass.REWRITE),
     PUSH_LIMIT_THROUGH_WINDOW(RuleTypeClass.REWRITE),
     LIMIT_SORT_TO_TOP_N(RuleTypeClass.REWRITE),
+    LIMIT_AGG_TO_TOPN_AGG(RuleTypeClass.REWRITE),
     // topN push down
     PUSH_DOWN_TOP_N_THROUGH_JOIN(RuleTypeClass.REWRITE),
     PUSH_DOWN_TOP_N_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LimitAggToTopNAgg.java

@@ -0,0 +1,119 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.properties.OrderKey;
+import org.apache.doris.nereids.rules.Rule;
+import org.apache.doris.nereids.rules.RuleType;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
+import org.apache.doris.nereids.trees.plans.logical.LogicalLimit;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
+import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
+import org.apache.doris.qe.ConnectContext;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * convert limit->agg to topn->agg
+ * if all group keys are in limit.output
+ * to enable
+ * 1. topn-filter
+ * 2. push limit to local agg
+ */
+public class LimitAggToTopNAgg implements RewriteRuleFactory {
+    @Override
+    public List<Rule> buildRules() {
+        return ImmutableList.of(
+                // limit -> agg to topn->agg
+                logicalLimit(logicalAggregate())
+                        .when(limit -> ConnectContext.get() != null
+                                && ConnectContext.get().getSessionVariable().pushTopnToAgg
+                                && ConnectContext.get().getSessionVariable().topnOptLimitThreshold
+                                >= limit.getLimit() + limit.getOffset())
+                        .then(limit -> {
+                            LogicalAggregate<? extends Plan> agg = limit.child();
+                            List<OrderKey> orderKeys = generateOrderKeyByGroupKey(agg);
+                            return new LogicalTopN<>(orderKeys, limit.getLimit(), limit.getOffset(), agg);
+                        }).toRule(RuleType.LIMIT_AGG_TO_TOPN_AGG),
+                //limit->project->agg to topn->project->agg
+                logicalLimit(logicalProject(logicalAggregate()))
+                        .when(limit -> ConnectContext.get() != null
+                                && ConnectContext.get().getSessionVariable().pushTopnToAgg
+                                && ConnectContext.get().getSessionVariable().topnOptLimitThreshold
+                                >= limit.getLimit() + limit.getOffset())
+                        .when(limit -> outputAllGroupKeys(limit, limit.child().child()))
+                        .then(limit -> {
+                            LogicalProject<? extends Plan> project = limit.child();
+                            LogicalAggregate<? extends Plan> agg = (LogicalAggregate<? extends Plan>) project.child();
+                            List<OrderKey> orderKeys = generateOrderKeyByGroupKey(agg);
+                            return new LogicalTopN<>(orderKeys, limit.getLimit(), limit.getOffset(), project);
+                        }).toRule(RuleType.LIMIT_AGG_TO_TOPN_AGG),
+                // topn -> agg: add all group key to sort key, if sort key is prefix of group key
+                logicalTopN(logicalAggregate())
+                        .when(topn -> ConnectContext.get() != null
+                                && ConnectContext.get().getSessionVariable().pushTopnToAgg
+                                && ConnectContext.get().getSessionVariable().topnOptLimitThreshold
+                                >= topn.getLimit() + topn.getOffset())
+                        .then(topn -> {
+                            LogicalAggregate<? extends Plan> agg = (LogicalAggregate<? extends Plan>) topn.child();
+                            List<OrderKey> newOrders = tryGenerateOrderKeyByGroupKeyAndTopnKey(topn, agg);
+                            if (newOrders.isEmpty()) {
+                                return topn;
+                            } else {
+                                return topn.withOrderKeys(newOrders);
+                            }
+                        }).toRule(RuleType.LIMIT_AGG_TO_TOPN_AGG));
+    }
+
+    private List<OrderKey> tryGenerateOrderKeyByGroupKeyAndTopnKey(LogicalTopN<? extends Plan> topN,
+            LogicalAggregate<? extends Plan> agg) {
+        List<OrderKey> orderKeys = Lists.newArrayListWithCapacity(agg.getGroupByExpressions().size());
+        if (topN.getOrderKeys().size() > agg.getGroupByExpressions().size()) {
+            return orderKeys;
+        }
+        List<Expression> topnKeys = topN.getOrderKeys().stream()
+                .map(OrderKey::getExpr).collect(Collectors.toList());
+        for (int i = 0; i < topN.getOrderKeys().size(); i++) {
+            // prefix check
+            if (!topnKeys.get(i).equals(agg.getGroupByExpressions().get(i))) {
+                return Lists.newArrayList();
+            }
+            orderKeys.add(topN.getOrderKeys().get(i));
+        }
+        for (int i = topN.getOrderKeys().size(); i < agg.getGroupByExpressions().size(); i++) {
+            orderKeys.add(new OrderKey(agg.getGroupByExpressions().get(i), true, false));
+        }
+        return orderKeys;
+    }
+
+    private boolean outputAllGroupKeys(LogicalLimit limit, LogicalAggregate agg) {
+        return limit.getOutputSet().containsAll(agg.getGroupByExpressions());
+    }
+
+    private List<OrderKey> generateOrderKeyByGroupKey(LogicalAggregate<? extends Plan> agg) {
+        return agg.getGroupByExpressions().stream()
+                .map(key -> new OrderKey(key, true, false))
+                .collect(Collectors.toList());
+    }
+}