From 1f5edae090f72afede1fd74ee2fe0e846ec59245 Mon Sep 17 00:00:00 2001 From: morrySnow <101034200+morrySnow@users.noreply.github.com> Date: Wed, 29 May 2024 15:27:27 +0800 Subject: [PATCH] [fix](Nereids) prune not required window expressions on window operator (#35593) 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. --- .../nereids/rules/rewrite/AdjustNullable.java | 2 +- .../nereids/rules/rewrite/ColumnPruning.java | 25 ++++++++ .../rewrite/SimplifyWindowExpression.java | 2 +- .../trees/plans/logical/LogicalWindow.java | 2 +- .../window_column_pruning.groovy | 60 +++++++++++++++++++ 5 files changed, 88 insertions(+), 3 deletions(-) create mode 100644 regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java index 0404e7b71c..44b88c54e8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java @@ -246,7 +246,7 @@ public class AdjustNullable extends DefaultPlanRewriter> imple List 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 diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java index 4cb18e8a38..32478ceb95 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java @@ -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 implements return super.visitLogicalCTEConsumer(cteConsumer, context); } + @Override + public Plan visitLogicalWindow(LogicalWindow window, PruneContext context) { + boolean pruned = false; + boolean reserved = false; + ImmutableList.Builder 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 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 prunedOutputAgg = pruneOutput(agg, agg.getOutputs(), agg::pruneOutputs, context); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java index c0548a4257..3c59657919 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java @@ -120,7 +120,7 @@ public class SimplifyWindowExpression extends OneRewriteRuleFactory { } List finalProjections = Lists.newArrayList(projections); finalProjections.addAll(windowOutputs); - return new LogicalProject(finalProjections, window.withExpression(remainWindows, + return new LogicalProject(finalProjections, window.withExpressionsAndChild(remainWindows, window.child(0))); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java index 56a306b600..ed99c26516 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java @@ -86,7 +86,7 @@ public class LogicalWindow extends LogicalUnary withExpression(List windowExpressions, Plan child) { + public LogicalWindow withExpressionsAndChild(List windowExpressions, Plan child) { return new LogicalWindow<>(windowExpressions, isChecked, child); } diff --git a/regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy b/regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy new file mode 100644 index 0000000000..a83f8ed752 --- /dev/null +++ b/regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy @@ -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" + } +} +