From fbadb21ae45ac6d534282a4788c7f14ee4755cfa Mon Sep 17 00:00:00 2001 From: morrySnow Date: Tue, 28 May 2024 14:54:25 +0800 Subject: [PATCH] [fix](Nereids) prune not require window expressions on window operator 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 +++++++++++++++++++ .../tpch/push_filter_window_eqset.groovy | 2 +- 6 files changed, 89 insertions(+), 4 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 0404e7b71c9bc3..44b88c54e85440 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 Plan visitLogicalWindow(LogicalWindow window, Map 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 1257a3fb3af102..e8d4c6d96abbb2 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 @@ -38,6 +38,7 @@ 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; @@ -242,6 +243,30 @@ public Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, PruneContext 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 c84ac975a9e9d4..0db1d59b9e6fd1 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 @@ private Plan simplify(MatchingContext> ctx) { } 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 99d8f03070684e..163d533db6eef8 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 List getExpressions() { return windowExpressions; } - public LogicalWindow 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 00000000000000..a83f8ed75280e5 --- /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" + } +} + diff --git a/regression-test/suites/nereids_tpch_p0/tpch/push_filter_window_eqset.groovy b/regression-test/suites/nereids_tpch_p0/tpch/push_filter_window_eqset.groovy index e6926b8dd6a6a7..f0dffcf7d4f685 100644 --- a/regression-test/suites/nereids_tpch_p0/tpch/push_filter_window_eqset.groovy +++ b/regression-test/suites/nereids_tpch_p0/tpch/push_filter_window_eqset.groovy @@ -37,7 +37,7 @@ suite("push_filter_window_eqset") { **/ qt_eqset """ explain shape plan - select y + select y, rn from ( select r_regionkey as x, r_regionkey as y, row_number() over(partition by r_regionkey) as rn from region ) T