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" + } +}