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