diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java index e02c299e71..cbfe3d28ac 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java @@ -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(), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 61efe9e32f..e57934c920 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -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), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/MergeTopNs.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/MergeTopNs.java new file mode 100644 index 0000000000..39ea7a384d --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/MergeTopNs.java @@ -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 + *

+ * 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 childTopN = topN.child(); + List orderKeys = topN.getOrderKeys(); + List 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); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/MergeTopNsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/MergeTopNsTest.java new file mode 100644 index 0000000000..e16a38be6b --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/MergeTopNsTest.java @@ -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) + ); + } + +}