From 5d117991f0cae68a46f886806c63ec84e886e3b7 Mon Sep 17 00:00:00 2001 From: morrySnow Date: Wed, 29 May 2024 11:44:13 +0800 Subject: [PATCH] [fix](Nereids) prune not required 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 ++++++++ .../trees/plans/logical/LogicalWindow.java | 5 +- .../window_column_pruning.groovy | 59 +++++++++++++++++++ 4 files changed, 87 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 26e9403a00639f..c5c9b8520fe194 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 @@ -223,7 +223,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 ad95f1a552f91f..c759ada4a0c898 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 @@ -34,6 +34,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; @@ -170,6 +171,30 @@ public Plan visitLogicalRepeat(LogicalRepeat repeat, PruneContex return pruneAggregate(repeat, 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/trees/plans/logical/LogicalWindow.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java index 67e18197796bef..13fc55796e9185 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 @@ -88,9 +88,8 @@ public List getExpressions() { return windowExpressions; } - public LogicalWindow withExpression(List windowExpressions, Plan child) { - return new LogicalWindow<>(windowExpressions, isChecked, Optional.empty(), - Optional.empty(), child); + public LogicalWindow withExpressionsAndChild(List windowExpressions, Plan child) { + return new LogicalWindow<>(windowExpressions, isChecked, child); } public LogicalWindow withChecked(List windowExpressions, Plan 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..5cbaae7c9e7225 --- /dev/null +++ b/regression-test/suites/nereids_rules_p0/column_pruning/window_column_pruning.groovy @@ -0,0 +1,59 @@ +// 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" + } +}