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

[fix](Nereids) prune not required window expressions on window operator (#35593)

Browse Source 
pick from master #35504

if window expression is not required by its parent, we should prune this
column. If all window expressions of window operator are pruned, we
remove this window operator directly.
This commit is contained in:
morrySnow
2024-05-29 15:27:27 +08:00
committed by GitHub
parent 8c0c05b9c6
commit 1f5edae090
5 changed files with 88 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java
View File

@ -246,7 +246,7 @@ public class AdjustNullable extends DefaultPlanRewriter<Map<ExprId, Slot>> imple
List<NamedExpression> windowExpressions =
updateExpressions(window.getWindowExpressions(), replaceMap);
windowExpressions.forEach(w -> replaceMap.put(w.getExprId(), w.toSlot()));
return window.withExpression(windowExpressions, window.child());
return window.withExpressionsAndChild(windowExpressions, window.child());
}
@Override

25
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java
View File

@ -37,6 +37,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat;
import org.apache.doris.nereids.trees.plans.logical.LogicalSink;
import org.apache.doris.nereids.trees.plans.logical.LogicalUnion;
import org.apache.doris.nereids.trees.plans.logical.LogicalWindow;
import org.apache.doris.nereids.trees.plans.logical.OutputPrunable;
import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter;
import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
@ -211,6 +212,30 @@ public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements
return super.visitLogicalCTEConsumer(cteConsumer, context);
}
@Override
public Plan visitLogicalWindow(LogicalWindow<? extends Plan> window, PruneContext context) {
boolean pruned = false;
boolean reserved = false;
ImmutableList.Builder<NamedExpression> reservedWindowExpressions = ImmutableList.builder();
for (NamedExpression windowExpression : window.getWindowExpressions()) {
if (context.requiredSlots.contains(windowExpression.toSlot())) {
reservedWindowExpressions.add(windowExpression);
reserved = true;
} else {
pruned = true;
}
}
if (!pruned) {
return pruneChildren(window, context.requiredSlots);
}
if (!reserved) {
return window.child().accept(this, context);
}
LogicalWindow<? extends Plan> prunedWindow
= window.withExpressionsAndChild(reservedWindowExpressions.build(), window.child());
return pruneChildren(prunedWindow, context.requiredSlots);
}
private Plan pruneAggregate(Aggregate<?> agg, PruneContext context) {
// first try to prune group by and aggregate functions
Aggregate<? extends Plan> prunedOutputAgg = pruneOutput(agg, agg.getOutputs(), agg::pruneOutputs, context);

2
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java
View File

@ -120,7 +120,7 @@ public class SimplifyWindowExpression extends OneRewriteRuleFactory {
}
List<NamedExpression> finalProjections = Lists.newArrayList(projections);
finalProjections.addAll(windowOutputs);
return new LogicalProject(finalProjections, window.withExpression(remainWindows,
return new LogicalProject(finalProjections, window.withExpressionsAndChild(remainWindows,
window.child(0)));
}
}

2
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java
View File

@ -86,7 +86,7 @@ public class LogicalWindow<CHILD_TYPE extends Plan> extends LogicalUnary<CHILD_T
return windowExpressions;
}
public LogicalWindow<Plan> withExpression(List<NamedExpression> windowExpressions, Plan child) {
public LogicalWindow<Plan> withExpressionsAndChild(List<NamedExpression> windowExpressions, Plan child) {
return new LogicalWindow<>(windowExpressions, isChecked, child);
}

60
regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy Normal file
View File

@ -0,0 +1,60 @@
// 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.
suite("window_column_pruning") {
sql "set disable_nereids_rules=PRUNE_EMPTY_PARTITION"
sql """
DROP TABLE IF EXISTS window_column_pruning
"""
sql """
CREATE TABLE IF NOT EXISTS window_column_pruning(
`id` int NULL,
`c1` int NULL
) ENGINE = OLAP
DISTRIBUTED BY HASH(id) BUCKETS 4
PROPERTIES (
"replication_allocation" = "tag.location.default: 1"
);
"""
// should prune
explain {
sql "select id from (select id, row_number() over (partition by id) as rn from window_column_pruning) tmp where id > 1;"
notContains "row_number"
}
// should not prune
explain {
sql "select id, rn from (select id, row_number() over (partition by id) as rn from window_column_pruning) tmp where id > 1;"
contains "row_number"
}
// should prune rank, but not prune row_number
explain {
sql "select id, rn1 from (select id, row_number() over (partition by id) as rn1, rank() over (partition by id) as rk from window_column_pruning) tmp where id > 1;"
contains "row_number"
notContains "rank"
}
// prune through union all
explain {
sql "select id from (select id, rank() over() px from window_column_pruning union all select id, rank() over() px from window_column_pruning) a"
notContains "rank"
}
}