database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature](Nereids): merge topNs (#28246)

Browse Source 
merge topNs like 
```
TopN
|
TopN

merge ->

TopN
```
This commit is contained in:
jakevin
2024-01-10 15:33:45 +08:00
committed by yiguolei
parent 3e9cd3a8b9
commit ddf2e8d3dd
4 changed files with 136 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -82,6 +82,7 @@ import org.apache.doris.nereids.rules.rewrite.MergeFilters;
import org.apache.doris.nereids.rules.rewrite.MergeOneRowRelationIntoUnion;
import org.apache.doris.nereids.rules.rewrite.MergeProjects;
import org.apache.doris.nereids.rules.rewrite.MergeSetOperations;
import org.apache.doris.nereids.rules.rewrite.MergeTopNs;
import org.apache.doris.nereids.rules.rewrite.NormalizeSort;
import org.apache.doris.nereids.rules.rewrite.OrExpansion;
import org.apache.doris.nereids.rules.rewrite.PruneEmptyPartition;
@ -323,6 +324,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 MergeTopNs()),
topDown(new SplitLimit()),
topDown(
new PushDownLimit(),

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

@ -245,6 +245,7 @@ public enum RuleType {
PUSH_PROJECT_INTO_ONE_ROW_RELATION(RuleTypeClass.REWRITE),
PUSH_PROJECT_INTO_UNION(RuleTypeClass.REWRITE),
MERGE_SET_OPERATION(RuleTypeClass.REWRITE),
MERGE_TOP_N(RuleTypeClass.REWRITE),
BUILD_AGG_FOR_UNION(RuleTypeClass.REWRITE),
COUNT_DISTINCT_REWRITE(RuleTypeClass.REWRITE),
INNER_TO_CROSS_JOIN(RuleTypeClass.REWRITE),

56
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/MergeTopNs.java Normal file
View File

@ -0,0 +1,56 @@
// 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.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
import java.util.List;
/**
* this rule aims to merge consecutive topN
* <p>
* topN - child topN
* If child topN orderby list is prefix subList of topN =>
* topN with limit = min(topN.limit, childTopN.limit), offset = topN.offset + childTopN.offset
*/
public class MergeTopNs extends OneRewriteRuleFactory {
@Override
public Rule build() {
return logicalTopN(logicalTopN())
.then(topN -> {
LogicalTopN<Plan> childTopN = topN.child();
List<OrderKey> orderKeys = topN.getOrderKeys();
List<OrderKey> childOrderKeys = childTopN.getOrderKeys();
if (!orderKeys.subList(0, childOrderKeys.size()).equals(childOrderKeys)) {
return null;
}
long offset = topN.getOffset();
long limit = topN.getLimit();
long childOffset = childTopN.getOffset();
long childLimit = childTopN.getLimit();
long newOffset = offset + childOffset;
// choose min limit
long newLimit = Math.min(limit, childLimit);
return topN.withLimitChild(newLimit, newOffset, childTopN.child());
}).toRule(RuleType.MERGE_TOP_N);
}
}

77
fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/MergeTopNsTest.java Normal file
View File

@ -0,0 +1,77 @@
// 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.trees.expressions.StatementScopeIdGenerator;
import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.util.LogicalPlanBuilder;
import org.apache.doris.nereids.util.MemoPatternMatchSupported;
import org.apache.doris.nereids.util.MemoTestUtils;
import org.apache.doris.nereids.util.PlanChecker;
import org.apache.doris.nereids.util.PlanConstructor;
import com.google.common.collect.ImmutableList;
import org.junit.jupiter.api.Test;
/**
* MergeConsecutiveProjects ut
*/
class MergeTopNsTest implements MemoPatternMatchSupported {
LogicalOlapScan score = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score);
@Test
void testMergeSameOrderBy() {
LogicalPlan plan = new LogicalPlanBuilder(score)
.topN(10, 0, ImmutableList.of(0, 1))
.topN(20, 0, ImmutableList.of(0, 1))
.build();
PlanChecker.from(MemoTestUtils.createConnectContext(), plan)
.applyTopDown(new MergeTopNs())
.matches(
logicalTopN(logicalOlapScan()).when(topN -> topN.getLimit() == 10)
);
}
@Test
void testMergeSubOrderBy() {
LogicalPlan plan = new LogicalPlanBuilder(score)
.topN(10, 0, ImmutableList.of(0))
.topN(20, 0, ImmutableList.of(0, 1))
.build();
PlanChecker.from(MemoTestUtils.createConnectContext(), plan)
.applyTopDown(new MergeTopNs())
.matches(
logicalTopN(logicalOlapScan()).when(topN -> topN.getLimit() == 10)
);
}
@Test
void testOffset() {
LogicalPlan plan = new LogicalPlanBuilder(score)
.topN(10, 3, ImmutableList.of(0))
.topN(20, 1, ImmutableList.of(0, 1))
.build();
PlanChecker.from(MemoTestUtils.createConnectContext(), plan)
.applyTopDown(new MergeTopNs())
.matches(
logicalTopN(logicalOlapScan()).when(topN -> topN.getLimit() == 10 && topN.getOffset() == 4)
);
}
}