Skip to content

Commit 8c1ead2

Browse files
engleflyenglefly
committed
support limit->proj, support topn-agg
1 parent ec8782b commit 8c1ead2

File tree

8 files changed

+296
-9
lines changed

8 files changed

+296
-9
lines changed

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,8 @@ public PlanFragment visitPhysicalDistribute(PhysicalDistribute<? extends Plan> d
291291
&& context.getFirstAggregateInFragment(inputFragment) == child) {
292292
PhysicalHashAggregate<?> hashAggregate = (PhysicalHashAggregate<?>) child;
293293
if (hashAggregate.getAggPhase() == AggPhase.LOCAL
294-
&& hashAggregate.getAggMode() == AggMode.INPUT_TO_BUFFER) {
294+
&& hashAggregate.getAggMode() == AggMode.INPUT_TO_BUFFER
295+
&& hashAggregate.getTopnPushInfo() == null) {
295296
AggregationNode aggregationNode = (AggregationNode) inputFragment.getPlanRoot();
296297
aggregationNode.setUseStreamingPreagg(hashAggregate.isMaybeUsingStream());
297298
}
@@ -1035,6 +1036,23 @@ public PlanFragment visitPhysicalHashAggregate(
10351036
// local exchanger will be used.
10361037
aggregationNode.setColocate(true);
10371038
}
1039+
if (aggregate.getTopnPushInfo() != null) {
1040+
List<Expr> orderingExprs = Lists.newArrayList();
1041+
List<Boolean> ascOrders = Lists.newArrayList();
1042+
List<Boolean> nullsFirstParams = Lists.newArrayList();
1043+
aggregate.getTopnPushInfo().orderkeys.forEach(k -> {
1044+
orderingExprs.add(ExpressionTranslator.translate(k.getExpr(), context));
1045+
ascOrders.add(k.isAsc());
1046+
nullsFirstParams.add(k.isNullFirst());
1047+
});
1048+
SortInfo sortInfo = new SortInfo(orderingExprs, ascOrders, nullsFirstParams, outputTupleDesc);
1049+
aggregationNode.setSortByGroupKey(sortInfo);
1050+
if (aggregationNode.getLimit() == -1) {
1051+
aggregationNode.setLimit(aggregate.getTopnPushInfo().limit);
1052+
}
1053+
} else {
1054+
aggregationNode.setSortByGroupKey(null);
1055+
}
10381056
setPlanRoot(inputPlanFragment, aggregationNode, aggregate);
10391057
if (aggregate.getStats() != null) {
10401058
aggregationNode.setCardinality((long) aggregate.getStats().getRowCount());

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,9 @@ public List<PlanPostProcessor> getProcessors() {
6464
builder.add(new RecomputeLogicalPropertiesProcessor());
6565
builder.add(new AddOffsetIntoDistribute());
6666
builder.add(new CommonSubExpressionOpt());
67+
if (cascadesContext.getConnectContext().getSessionVariable().pushLimitToLocalAgg) {
68+
builder.add(new PushLimitToLocalAgg());
69+
}
6770
// DO NOT replace PLAN NODE from here
6871
builder.add(new TopNScanOpt());
6972
builder.add(new FragmentProcessor());

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

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
// This file is copied from
18+
// https://github.com/apache/impala/blob/branch-2.9.0/fe/src/main/java/org/apache/impala/AggregationNode.java
19+
// and modified by Doris
20+
21+
package org.apache.doris.nereids.processor.post;
22+
23+
import org.apache.doris.nereids.CascadesContext;
24+
import org.apache.doris.nereids.properties.OrderKey;
25+
import org.apache.doris.nereids.trees.expressions.Expression;
26+
import org.apache.doris.nereids.trees.plans.Plan;
27+
import org.apache.doris.nereids.trees.plans.physical.PhysicalDistribute;
28+
import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate;
29+
import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate.TopnPushInfo;
30+
import org.apache.doris.nereids.trees.plans.physical.PhysicalLimit;
31+
import org.apache.doris.nereids.trees.plans.physical.PhysicalProject;
32+
import org.apache.doris.nereids.trees.plans.physical.PhysicalTopN;
33+
34+
import org.apache.hadoop.util.Lists;
35+
36+
import java.util.List;
37+
import java.util.stream.Collectors;
38+
39+
/**
40+
Pattern1:
41+
limit(n) -> aggGlobal -> distribute -> aggLocal
42+
=>
43+
limit(n) -> aggGlobal(topn=n) -> distribute -> aggLocal(topn=n)
44+
45+
Pattern2: topn orderkeys are the prefix of group keys
46+
topn -> aggGlobal -> distribute -> aggLocal
47+
=>
48+
topn(n) -> aggGlobal(topn=n) -> distribute -> aggLocal(topn=n)
49+
*/
50+
public class PushLimitToLocalAgg extends PlanPostProcessor {
51+
@Override
52+
public Plan visitPhysicalTopN(PhysicalTopN<? extends Plan> topN, CascadesContext ctx) {
53+
Plan topnChild = topN.child();
54+
if (topnChild instanceof PhysicalProject) {
55+
topnChild = topnChild.child(0);
56+
}
57+
if (topnChild instanceof PhysicalHashAggregate) {
58+
PhysicalHashAggregate<? extends Plan> upperAgg = (PhysicalHashAggregate<? extends Plan>) topnChild;
59+
if (upperAgg.getAggPhase().isGlobal()
60+
&& upperAgg.child() instanceof PhysicalDistribute
61+
&& upperAgg.child().child(0) instanceof PhysicalHashAggregate) {
62+
List<OrderKey> orderKeys = tryGenerateOrderKeyByGroupKeyAndTopnKey(topN, upperAgg);
63+
if (!orderKeys.isEmpty()) {
64+
upperAgg.setTopnPushInfo(new TopnPushInfo(
65+
orderKeys,
66+
topN.getLimit() + topN.getOffset()));
67+
PhysicalHashAggregate<? extends Plan> bottomAgg =
68+
(PhysicalHashAggregate<? extends Plan>) upperAgg.child().child(0);
69+
bottomAgg.setTopnPushInfo(new TopnPushInfo(
70+
orderKeys,
71+
topN.getLimit() + topN.getOffset()));
72+
}
73+
}
74+
}
75+
topN.child().accept(this, ctx);
76+
return topN;
77+
}
78+
79+
@Override
80+
public Plan visitPhysicalLimit(PhysicalLimit<? extends Plan> limit, CascadesContext ctx) {
81+
Plan limitChild = limit.child();
82+
if (limitChild instanceof PhysicalProject) {
83+
limitChild = limitChild.child(0);
84+
}
85+
if (limitChild instanceof PhysicalHashAggregate) {
86+
PhysicalHashAggregate<? extends Plan> upperAgg = (PhysicalHashAggregate<? extends Plan>) limitChild;
87+
if (upperAgg.getAggPhase().isGlobal()
88+
&& upperAgg.child() instanceof PhysicalDistribute
89+
&& upperAgg.child().child(0) instanceof PhysicalHashAggregate) {
90+
upperAgg.setTopnPushInfo(new TopnPushInfo(
91+
generateOrderKeyByGroupKey(upperAgg),
92+
limit.getLimit() + limit.getOffset()));
93+
PhysicalHashAggregate<? extends Plan> bottomAgg =
94+
(PhysicalHashAggregate<? extends Plan>) upperAgg.child().child(0);
95+
bottomAgg.setTopnPushInfo(new TopnPushInfo(
96+
generateOrderKeyByGroupKey(bottomAgg),
97+
limit.getLimit() + limit.getOffset()));
98+
}
99+
}
100+
limit.child().accept(this, ctx);
101+
102+
return limit;
103+
}
104+
105+
private List<OrderKey> generateOrderKeyByGroupKey(PhysicalHashAggregate<? extends Plan> agg) {
106+
return agg.getGroupByExpressions().stream()
107+
.map(key -> new OrderKey(key, true, false))
108+
.collect(Collectors.toList());
109+
}
110+
111+
/**
112+
return true, if topn order-key is prefix of agg group-key, ignore asc/desc and null_first
113+
TODO order-key can be subset of group-key. BE does not support now.
114+
*/
115+
private List<OrderKey> tryGenerateOrderKeyByGroupKeyAndTopnKey(PhysicalTopN<? extends Plan> topN,
116+
PhysicalHashAggregate<? extends Plan> agg) {
117+
List<OrderKey> orderKeys = Lists.newArrayListWithCapacity(agg.getGroupByExpressions().size());
118+
if (topN.getOrderKeys().size() > agg.getGroupByExpressions().size()) {
119+
return orderKeys;
120+
}
121+
List<Expression> topnKeys = topN.getOrderKeys().stream()
122+
.map(OrderKey::getExpr).collect(Collectors.toList());
123+
for (int i = 0; i < topN.getOrderKeys().size(); i++) {
124+
// prefix check
125+
if (!topnKeys.get(i).equals(agg.getGroupByExpressions().get(i))) {
126+
return Lists.newArrayList();
127+
}
128+
orderKeys.add(topN.getOrderKeys().get(i));
129+
}
130+
for (int i = topN.getOrderKeys().size(); i < agg.getGroupByExpressions().size(); i++) {
131+
orderKeys.add(new OrderKey(agg.getGroupByExpressions().get(i), true, false));
132+
}
133+
return orderKeys;
134+
}
135+
}

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import org.apache.doris.nereids.memo.GroupExpression;
2121
import org.apache.doris.nereids.properties.LogicalProperties;
22+
import org.apache.doris.nereids.properties.OrderKey;
2223
import org.apache.doris.nereids.properties.PhysicalProperties;
2324
import org.apache.doris.nereids.properties.RequireProperties;
2425
import org.apache.doris.nereids.properties.RequirePropertiesSupplier;
@@ -60,6 +61,9 @@ public class PhysicalHashAggregate<CHILD_TYPE extends Plan> extends PhysicalUnar
6061

6162
private final RequireProperties requireProperties;
6263

64+
// only used in post processor
65+
private TopnPushInfo topnPushInfo = null;
66+
6367
public PhysicalHashAggregate(List<Expression> groupByExpressions, List<NamedExpression> outputExpressions,
6468
AggregateParam aggregateParam, boolean maybeUsingStream, LogicalProperties logicalProperties,
6569
RequireProperties requireProperties, CHILD_TYPE child) {
@@ -196,6 +200,7 @@ public String toString() {
196200
"outputExpr", outputExpressions,
197201
"partitionExpr", partitionExpressions,
198202
"requireProperties", requireProperties,
203+
"topnOpt", topnPushInfo != null,
199204
"stats", statistics
200205
);
201206
}
@@ -231,19 +236,22 @@ public PhysicalHashAggregate<Plan> withChildren(List<Plan> children) {
231236
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions, partitionExpressions,
232237
aggregateParam, maybeUsingStream, groupExpression, getLogicalProperties(),
233238
requireProperties, physicalProperties, statistics,
234-
children.get(0));
239+
children.get(0))
240+
.setTopnPushInfo(topnPushInfo);
235241
}
236242

237243
public PhysicalHashAggregate<CHILD_TYPE> withPartitionExpressions(List<Expression> partitionExpressions) {
238244
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions,
239245
Optional.ofNullable(partitionExpressions), aggregateParam, maybeUsingStream,
240-
Optional.empty(), getLogicalProperties(), requireProperties, child());
246+
Optional.empty(), getLogicalProperties(), requireProperties, child())
247+
.setTopnPushInfo(topnPushInfo);
241248
}
242249

243250
@Override
244251
public PhysicalHashAggregate<CHILD_TYPE> withGroupExpression(Optional<GroupExpression> groupExpression) {
245252
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions, partitionExpressions,
246-
aggregateParam, maybeUsingStream, groupExpression, getLogicalProperties(), requireProperties, child());
253+
aggregateParam, maybeUsingStream, groupExpression, getLogicalProperties(), requireProperties, child())
254+
.setTopnPushInfo(topnPushInfo);
247255
}
248256

249257
@Override
@@ -252,7 +260,7 @@ public Plan withGroupExprLogicalPropChildren(Optional<GroupExpression> groupExpr
252260
Preconditions.checkArgument(children.size() == 1);
253261
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions, partitionExpressions,
254262
aggregateParam, maybeUsingStream, groupExpression, logicalProperties.get(),
255-
requireProperties, children.get(0));
263+
requireProperties, children.get(0)).setTopnPushInfo(topnPushInfo);
256264
}
257265

258266
@Override
@@ -261,21 +269,21 @@ public PhysicalHashAggregate<CHILD_TYPE> withPhysicalPropertiesAndStats(Physical
261269
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions, partitionExpressions,
262270
aggregateParam, maybeUsingStream, groupExpression, getLogicalProperties(),
263271
requireProperties, physicalProperties, statistics,
264-
child());
272+
child()).setTopnPushInfo(topnPushInfo);
265273
}
266274

267275
@Override
268276
public PhysicalHashAggregate<CHILD_TYPE> withAggOutput(List<NamedExpression> newOutput) {
269277
return new PhysicalHashAggregate<>(groupByExpressions, newOutput, partitionExpressions,
270278
aggregateParam, maybeUsingStream, Optional.empty(), getLogicalProperties(),
271-
requireProperties, physicalProperties, statistics, child());
279+
requireProperties, physicalProperties, statistics, child()).setTopnPushInfo(topnPushInfo);
272280
}
273281

274282
public <C extends Plan> PhysicalHashAggregate<C> withRequirePropertiesAndChild(
275283
RequireProperties requireProperties, C newChild) {
276284
return new PhysicalHashAggregate<>(groupByExpressions, outputExpressions, partitionExpressions,
277285
aggregateParam, maybeUsingStream, Optional.empty(), getLogicalProperties(),
278-
requireProperties, physicalProperties, statistics, newChild);
286+
requireProperties, physicalProperties, statistics, newChild).setTopnPushInfo(topnPushInfo);
279287
}
280288

281289
@Override
@@ -299,4 +307,26 @@ public PhysicalHashAggregate<CHILD_TYPE> resetLogicalProperties() {
299307
requireProperties, physicalProperties, statistics,
300308
child());
301309
}
310+
311+
/**
312+
* used to push limit down to localAgg
313+
*/
314+
public static class TopnPushInfo {
315+
public List<OrderKey> orderkeys;
316+
public long limit;
317+
318+
public TopnPushInfo(List<OrderKey> orderkeys, long limit) {
319+
this.orderkeys = ImmutableList.copyOf(orderkeys);
320+
this.limit = limit;
321+
}
322+
}
323+
324+
public TopnPushInfo getTopnPushInfo() {
325+
return topnPushInfo;
326+
}
327+
328+
public PhysicalHashAggregate<CHILD_TYPE> setTopnPushInfo(TopnPushInfo topnPushInfo) {
329+
this.topnPushInfo = topnPushInfo;
330+
return this;
331+
}
302332
}

‎fe/fe-core/src/main/java/org/apache/doris/planner/AggregationNode.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.doris.analysis.FunctionCallExpr;
2727
import org.apache.doris.analysis.SlotDescriptor;
2828
import org.apache.doris.analysis.SlotId;
29+
import org.apache.doris.analysis.SortInfo;
2930
import org.apache.doris.analysis.TupleDescriptor;
3031
import org.apache.doris.common.NotImplementedException;
3132
import org.apache.doris.common.UserException;
@@ -65,6 +66,8 @@ public class AggregationNode extends PlanNode {
6566
// If true, use streaming preaggregation algorithm. Not valid if this is a merge agg.
6667
private boolean useStreamingPreagg;
6768

69+
private SortInfo sortByGroupKey;
70+
6871
/**
6972
* Create an agg node that is not an intermediate node.
7073
* isIntermediate is true if it is a slave node in a 2-part agg plan.
@@ -288,6 +291,9 @@ protected void toThrift(TPlanNode msg) {
288291
msg.agg_node.setUseStreamingPreaggregation(useStreamingPreagg);
289292
msg.agg_node.setIsFirstPhase(aggInfo.isFirstPhase());
290293
msg.agg_node.setIsColocate(isColocate);
294+
if (sortByGroupKey != null) {
295+
msg.agg_node.setAggSortInfoByGroupKey(sortByGroupKey.toThrift());
296+
}
291297
List<Expr> groupingExprs = aggInfo.getGroupingExprs();
292298
if (groupingExprs != null) {
293299
msg.agg_node.setGroupingExprs(Expr.treesToThrift(groupingExprs));
@@ -333,6 +339,7 @@ public String getNodeExplainString(String detailPrefix, TExplainLevel detailLeve
333339
if (!conjuncts.isEmpty()) {
334340
output.append(detailPrefix).append("having: ").append(getExplainString(conjuncts)).append("\n");
335341
}
342+
output.append(detailPrefix).append("sortByGroupKey:").append(sortByGroupKey != null).append("\n");
336343
output.append(detailPrefix).append(String.format(
337344
"cardinality=%,d", cardinality)).append("\n");
338345
return output.toString();
@@ -411,4 +418,13 @@ public void finalize(Analyzer analyzer) throws UserException {
411418
public void setColocate(boolean colocate) {
412419
isColocate = colocate;
413420
}
421+
422+
423+
public boolean isSortByGroupKey() {
424+
return sortByGroupKey != null;
425+
}
426+
427+
public void setSortByGroupKey(SortInfo sortByGroupKey) {
428+
this.sortByGroupKey = sortByGroupKey;
429+
}
414430
}

‎fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1189,6 +1189,9 @@ public void setEnableLeftZigZag(boolean enableLeftZigZag) {
11891189
@VariableMgr.VarAttr(name = REWRITE_OR_TO_IN_PREDICATE_THRESHOLD, fuzzy = true)
11901190
private int rewriteOrToInPredicateThreshold = 2;
11911191

1192+
@VariableMgr.VarAttr(name = "push_limit_to_local_agg", fuzzy = false, needForward = true)
1193+
public boolean pushLimitToLocalAgg = true;
1194+
11921195
@VariableMgr.VarAttr(name = NEREIDS_CBO_PENALTY_FACTOR, needForward = true)
11931196
private double nereidsCboPenaltyFactor = 0.7;
11941197

‎gensrc/thrift/PlanNodes.thrift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -904,7 +904,7 @@ struct TAggregationNode {
904904
7: optional list<TSortInfo> agg_sort_infos
905905
8: optional bool is_first_phase
906906
9: optional bool is_colocate
907-
// 9: optional bool use_fixed_length_serialization_opt
907+
10: optional TSortInfo agg_sort_info_by_group_key
908908
}
909909

910910
struct TRepeatNode {

0 commit comments

Comments
 (0)