[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.

Author: morrySnow

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AdjustNullable.java

@@ -246,7 +246,7 @@ public Plan visitLogicalWindow(LogicalWindow<? extends Plan> window, Map<ExprId,
         List<NamedExpression> 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

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.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 Plan visitLogicalCTEConsumer(LogicalCTEConsumer cteConsumer, PruneContext
         return super.visitLogicalCTEConsumer(cteConsumer, context);
     }
 
+    @Override
+    public Plan visitLogicalWindow(LogicalWindow<? extends Plan> window, PruneContext context) {
+        boolean pruned = false;
+        boolean reserved = false;
+        ImmutableList.Builder<NamedExpression> 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<? extends Plan> 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<? extends Plan> prunedOutputAgg = pruneOutput(agg, agg.getOutputs(), agg::pruneOutputs, context);

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SimplifyWindowExpression.java

@@ -120,7 +120,7 @@ private Plan simplify(MatchingContext<LogicalWindow<Plan>> ctx) {
             }
             List<NamedExpression> finalProjections = Lists.newArrayList(projections);
             finalProjections.addAll(windowOutputs);
-            return new LogicalProject(finalProjections, window.withExpression(remainWindows,
+            return new LogicalProject(finalProjections, window.withExpressionsAndChild(remainWindows,
                     window.child(0)));
         }
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java

@@ -86,7 +86,7 @@ public List<? extends Expression> getExpressions() {
         return windowExpressions;
     }
 
-    public LogicalWindow<Plan> withExpression(List<NamedExpression> windowExpressions, Plan child) {
+    public LogicalWindow<Plan> withExpressionsAndChild(List<NamedExpression> windowExpressions, Plan child) {
         return new LogicalWindow<>(windowExpressions, isChecked, child);
     }
 

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