[feature](Nereids): Pushdown TopN through Union (#27535)

```
topn
-> Union All 
  -> child plan1
  -> child plan2
  -> child plan3

rewritten to

topn
-> Union All 
 -> topn
  -> child plan1
 -> topn
  -> child plan2
 -> topn
  -> child plan3
```
This commit is contained in:
jakevin
2023-11-27 14:13:18 +08:00
committed by GitHub
parent 1b4cd24b36
commit 234aff3e78
8 changed files with 430 additions and 55 deletions

View File

@ -99,6 +99,7 @@ import org.apache.doris.nereids.rules.rewrite.PushDownLimitDistinctThroughJoin;
import org.apache.doris.nereids.rules.rewrite.PushDownMinMaxThroughJoin;
import org.apache.doris.nereids.rules.rewrite.PushDownSumThroughJoin;
import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughJoin;
import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughUnion;
import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughWindow;
import org.apache.doris.nereids.rules.rewrite.PushFilterInsideJoin;
import org.apache.doris.nereids.rules.rewrite.PushProjectIntoOneRowRelation;
@ -296,7 +297,8 @@ public class Rewriter extends AbstractBatchJobExecutor {
new PushDownLimit(),
new PushDownTopNThroughJoin(),
new PushDownLimitDistinctThroughJoin(),
new PushDownTopNThroughWindow()
new PushDownTopNThroughWindow(),
new PushDownTopNThroughUnion()
),
topDown(new CreatePartitionTopNFromWindow()),
topDown(

View File

@ -256,10 +256,11 @@ public enum RuleType {
PUSH_LIMIT_THROUGH_WINDOW(RuleTypeClass.REWRITE),
LIMIT_SORT_TO_TOP_N(RuleTypeClass.REWRITE),
// topN push down
PUSH_TOP_N_THROUGH_JOIN(RuleTypeClass.REWRITE),
PUSH_TOP_N_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),
PUSH_TOP_N_THROUGH_PROJECT_WINDOW(RuleTypeClass.REWRITE),
PUSH_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),
PUSH_DOWN_TOP_N_THROUGH_JOIN(RuleTypeClass.REWRITE),
PUSH_DOWN_TOP_N_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),
PUSH_DOWN_TOP_N_THROUGH_PROJECT_WINDOW(RuleTypeClass.REWRITE),
PUSH_DOWN_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),
PUSH_DOWN_TOP_N_THROUGH_UNION(RuleTypeClass.REWRITE),
// limit distinct push down
PUSH_LIMIT_DISTINCT_THROUGH_JOIN(RuleTypeClass.REWRITE),
PUSH_LIMIT_DISTINCT_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),

View File

@ -53,7 +53,7 @@ public class PushDownTopNThroughJoin implements RewriteRuleFactory {
}
return topN.withChildren(newJoin);
})
.toRule(RuleType.PUSH_TOP_N_THROUGH_JOIN),
.toRule(RuleType.PUSH_DOWN_TOP_N_THROUGH_JOIN),
// topN -> project -> join
logicalTopN(logicalProject(logicalJoin()))
@ -79,7 +79,7 @@ public class PushDownTopNThroughJoin implements RewriteRuleFactory {
return null;
}
return topN.withChildren(project.withChildren(newJoin));
}).toRule(RuleType.PUSH_TOP_N_THROUGH_PROJECT_JOIN)
}).toRule(RuleType.PUSH_DOWN_TOP_N_THROUGH_PROJECT_JOIN)
);
}

View File

@ -0,0 +1,88 @@
// 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.expressions.NamedExpression;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier;
import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
import org.apache.doris.nereids.trees.plans.logical.LogicalUnion;
import org.apache.doris.nereids.util.ExpressionUtils;
import com.google.common.collect.ImmutableList;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* <pre>
* TopN
* -> Union All
* -> child plan1
* -> child plan2
* -> child plan3
*
* rewritten to
*
* -> Union All
* -> TopN
* -> child plan1
* -> TopN
* -> child plan2
* -> TopN
* -> child plan3
* </pre>
*/
public class PushDownTopNThroughUnion implements RewriteRuleFactory {
@Override
public List<Rule> buildRules() {
return ImmutableList.of(
logicalTopN(logicalUnion().when(union -> union.getQualifier() == Qualifier.ALL))
.then(topN -> {
LogicalUnion union = topN.child();
List<Plan> newChildren = new ArrayList<>();
for (Plan child : union.children()) {
Map<Expression, Expression> replaceMap = new HashMap<>();
for (int i = 0; i < union.getOutputs().size(); ++i) {
NamedExpression output = union.getOutputs().get(i);
replaceMap.put(output, child.getOutput().get(i));
}
List<OrderKey> orderKeys = topN.getOrderKeys().stream()
.map(orderKey -> orderKey.withExpression(
ExpressionUtils.replace(orderKey.getExpr(), replaceMap)))
.collect(ImmutableList.toImmutableList());
newChildren.add(
new LogicalTopN<>(orderKeys, topN.getLimit() + topN.getOffset(), 0, child));
}
if (union.children().equals(newChildren)) {
return null;
}
return topN.withChildren(union.withChildren(newChildren));
})
.toRule(RuleType.PUSH_DOWN_TOP_N_THROUGH_UNION)
);
}
}

View File

@ -59,7 +59,7 @@ public class PushDownTopNThroughWindow implements RewriteRuleFactory {
return topn;
}
return topn.withChildren(newWindow.get());
}).toRule(RuleType.PUSH_TOP_N_THROUGH_WINDOW),
}).toRule(RuleType.PUSH_DOWN_TOP_N_THROUGH_WINDOW),
// topn -> projection -> window
logicalTopN(logicalProject(logicalWindow())).then(topn -> {
@ -79,7 +79,7 @@ public class PushDownTopNThroughWindow implements RewriteRuleFactory {
return topn;
}
return topn.withChildren(project.withChildren(newWindow.get()));
}).toRule(RuleType.PUSH_TOP_N_THROUGH_PROJECT_WINDOW)
}).toRule(RuleType.PUSH_DOWN_TOP_N_THROUGH_PROJECT_WINDOW)
);
}