[fix](Nereids) column pruning should prune map in cte consumer

fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java

@@ -50,7 +50,6 @@
 import org.apache.doris.nereids.rules.exploration.mv.MaterializationContext;
 import org.apache.doris.nereids.trees.expressions.CTEId;
 import org.apache.doris.nereids.trees.expressions.Expression;
-import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
 import org.apache.doris.nereids.trees.plans.Plan;
@@ -102,7 +101,7 @@ public class CascadesContext implements ScheduleContext {
     private Optional<RootRewriteJobContext> currentRootRewriteJobContext;
     // in optimize stage, the plan will storage in the memo
     private Memo memo;
-    private StatementContext statementContext;
+    private final StatementContext statementContext;
 
     private final CTEContext cteContext;
     private final RuleSet ruleSet;
@@ -616,16 +615,6 @@ public void putCTEIdToConsumer(LogicalCTEConsumer cteConsumer) {
         consumers.add(cteConsumer);
     }
 
-    public void putCTEIdToProject(CTEId cteId, NamedExpression p) {
-        Set<NamedExpression> projects = this.statementContext.getCteIdToProjects()
-                .computeIfAbsent(cteId, k -> new HashSet<>());
-        projects.add(p);
-    }
-
-    public Set<NamedExpression> getProjectForProducer(CTEId cteId) {
-        return this.statementContext.getCteIdToProjects().get(cteId);
-    }
-
     public Map<CTEId, Set<LogicalCTEConsumer>> getCteIdToConsumers() {
         return this.statementContext.getCteIdToConsumers();
     }
@@ -639,17 +628,6 @@ public Map<RelationId, Set<Expression>> getConsumerIdToFilters() {
         return this.statementContext.getConsumerIdToFilters();
     }
 
-    public void markConsumerUnderProject(LogicalCTEConsumer cteConsumer) {
-        Set<RelationId> consumerIds = this.statementContext.getCteIdToConsumerUnderProjects()
-                .computeIfAbsent(cteConsumer.getCteId(), k -> new HashSet<>());
-        consumerIds.add(cteConsumer.getRelationId());
-    }
-
-    public boolean couldPruneColumnOnProducer(CTEId cteId) {
-        Set<RelationId> consumerIds = this.statementContext.getCteIdToConsumerUnderProjects().get(cteId);
-        return consumerIds.size() == this.statementContext.getCteIdToConsumers().get(cteId).size();
-    }
-
     public void addCTEConsumerGroup(CTEId cteId, Group g, Map<Slot, Slot> producerSlotToConsumerSlot) {
         List<Pair<Map<Slot, Slot>, Group>> consumerGroups =
                 this.statementContext.getCteIdToConsumerGroup().computeIfAbsent(cteId, k -> new ArrayList<>());
@@ -746,7 +724,7 @@ public void printPlanProcess() {
 
     public static void printPlanProcess(List<PlanProcess> planProcesses) {
         for (PlanProcess row : planProcesses) {
-            LOG.info("RULE: " + row.ruleName + "\nBEFORE:\n" + row.beforeShape + "\nafter:\n" + row.afterShape);
+            LOG.info("RULE: {}\nBEFORE:\n{}\nafter:\n{}", row.ruleName, row.beforeShape, row.afterShape);
         }
     }
 

fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java

@@ -27,7 +27,6 @@
 import org.apache.doris.nereids.trees.expressions.CTEId;
 import org.apache.doris.nereids.trees.expressions.ExprId;
 import org.apache.doris.nereids.trees.expressions.Expression;
-import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.plans.ObjectId;
@@ -54,7 +53,6 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -103,9 +101,8 @@ public class StatementContext implements Closeable {
     private final IdGenerator<CTEId> cteIdGenerator = CTEId.createGenerator();
 
     private final Map<CTEId, Set<LogicalCTEConsumer>> cteIdToConsumers = new HashMap<>();
-    private final Map<CTEId, Set<NamedExpression>> cteIdToProjects = new HashMap<>();
+    private final Map<CTEId, Set<Slot>> cteIdToOutputIds = new HashMap<>();
     private final Map<RelationId, Set<Expression>> consumerIdToFilters = new HashMap<>();
-    private final Map<CTEId, Set<RelationId>> cteIdToConsumerUnderProjects = new HashMap<>();
     // Used to update consumer's stats
     private final Map<CTEId, List<Pair<Map<Slot, Slot>, Group>>> cteIdToConsumerGroup = new HashMap<>();
     private final Map<CTEId, LogicalPlan> rewrittenCteProducer = new HashMap<>();
@@ -134,12 +131,13 @@ public class StatementContext implements Closeable {
     private BitSet disableRules;
 
     // table locks
-    private Stack<CloseableResource> plannerResources = new Stack<>();
+    private final Stack<CloseableResource> plannerResources = new Stack<>();
 
     // for create view support in nereids
     // key is the start and end position of the sql substring that needs to be replaced,
     // and value is the new string used for replacement.
-    private TreeMap<Pair<Integer, Integer>, String> indexInSqlToString = new TreeMap<>(new Pair.PairComparator<>());
+    private final TreeMap<Pair<Integer, Integer>, String> indexInSqlToString
+            = new TreeMap<>(new Pair.PairComparator<>());
 
     public StatementContext() {
         this(ConnectContext.get(), null, 0);
@@ -216,10 +214,6 @@ public Optional<SqlCacheContext> getSqlCacheContext() {
         return Optional.ofNullable(sqlCacheContext);
     }
 
-    public int getMaxContinuousJoin() {
-        return joinCount;
-    }
-
     public Set<SlotReference> getAllPathsSlots() {
         Set<SlotReference> allSlotReferences = Sets.newHashSet();
         for (Map<List<String>, SlotReference> slotReferenceMap : subColumnSlotRefMap.values()) {
@@ -240,19 +234,16 @@ public Slot getRewrittenSlotRefByOriginalExpr(Expression originalExpr) {
      * Add a slot ref attached with paths in context to avoid duplicated slot
      */
     public void addPathSlotRef(Slot root, List<String> paths, SlotReference slotRef, Expression originalExpr) {
-        subColumnSlotRefMap.computeIfAbsent(root, k -> Maps.newTreeMap(new Comparator<List<String>>() {
-            @Override
-            public int compare(List<String> lst1, List<String> lst2) {
-                Iterator<String> it1 = lst1.iterator();
-                Iterator<String> it2 = lst2.iterator();
-                while (it1.hasNext() && it2.hasNext()) {
-                    int result = it1.next().compareTo(it2.next());
-                    if (result != 0) {
-                        return result;
-                    }
+        subColumnSlotRefMap.computeIfAbsent(root, k -> Maps.newTreeMap((lst1, lst2) -> {
+            Iterator<String> it1 = lst1.iterator();
+            Iterator<String> it2 = lst2.iterator();
+            while (it1.hasNext() && it2.hasNext()) {
+                int result = it1.next().compareTo(it2.next());
+                if (result != 0) {
+                    return result;
                 }
-                return Integer.compare(lst1.size(), lst2.size());
             }
+            return Integer.compare(lst1.size(), lst2.size());
         }));
         subColumnSlotRefMap.get(root).put(paths, slotRef);
         subColumnOriginalExprMap.put(slotRef, originalExpr);
@@ -349,18 +340,14 @@ public Map<CTEId, Set<LogicalCTEConsumer>> getCteIdToConsumers() {
         return cteIdToConsumers;
     }
 
-    public Map<CTEId, Set<NamedExpression>> getCteIdToProjects() {
-        return cteIdToProjects;
+    public Map<CTEId, Set<Slot>> getCteIdToOutputIds() {
+        return cteIdToOutputIds;
     }
 
     public Map<RelationId, Set<Expression>> getConsumerIdToFilters() {
         return consumerIdToFilters;
     }
 
-    public Map<CTEId, Set<RelationId>> getCteIdToConsumerUnderProjects() {
-        return cteIdToConsumerUnderProjects;
-    }
-
     public Map<CTEId, List<Pair<Map<Slot, Slot>, Group>>> getCteIdToConsumerGroup() {
         return cteIdToConsumerGroup;
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java

@@ -1123,9 +1123,13 @@ public PlanFragment visitPhysicalCTEConsumer(PhysicalCTEConsumer cteConsumer,
         // update expr to slot mapping
         TupleDescriptor tupleDescriptor = null;
         for (Slot producerSlot : cteProducer.getOutput()) {
-            Slot consumerSlot = cteConsumer.getProducerToConsumerSlotMap().get(producerSlot);
             SlotRef slotRef = context.findSlotRef(producerSlot.getExprId());
             tupleDescriptor = slotRef.getDesc().getParent();
+            Slot consumerSlot = cteConsumer.getProducerToConsumerSlotMap().get(producerSlot);
+            // consumerSlot could be null if we prune partial consumers' columns
+            if (consumerSlot == null) {
+                continue;
+            }
             context.addExprIdSlotRefPair(consumerSlot.getExprId(), slotRef);
         }
         CTEScanNode cteScanNode = new CTEScanNode(tupleDescriptor);

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -43,8 +43,8 @@
 import org.apache.doris.nereids.rules.rewrite.CheckMatchExpression;
 import org.apache.doris.nereids.rules.rewrite.CheckMultiDistinct;
 import org.apache.doris.nereids.rules.rewrite.CheckPrivileges;
+import org.apache.doris.nereids.rules.rewrite.CollectCteConsumerOutput;
 import org.apache.doris.nereids.rules.rewrite.CollectFilterAboveConsumer;
-import org.apache.doris.nereids.rules.rewrite.CollectProjectAboveConsumer;
 import org.apache.doris.nereids.rules.rewrite.ColumnPruning;
 import org.apache.doris.nereids.rules.rewrite.ConvertInnerOrCrossJoin;
 import org.apache.doris.nereids.rules.rewrite.CountDistinctRewrite;
@@ -415,7 +415,7 @@ public class Rewriter extends AbstractBatchJobExecutor {
             topic("Push project and filter on cte consumer to cte producer",
                     topDown(
                             new CollectFilterAboveConsumer(),
-                            new CollectProjectAboveConsumer()
+                            new CollectCteConsumerOutput()
                     )
             )
     );

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

@@ -306,8 +306,7 @@ public enum RuleType {
 
     COLLECT_FILTER(RuleTypeClass.REWRITE),
     COLLECT_JOIN_CONSTRAINT(RuleTypeClass.REWRITE),
-    COLLECT_PROJECT_ABOVE_CTE_CONSUMER(RuleTypeClass.REWRITE),
-    COLLECT_PROJECT_ABOVE_FILTER_CTE_CONSUMER(RuleTypeClass.REWRITE),
+    COLLECT_CTE_CONSUMER_OUTPUT(RuleTypeClass.REWRITE),
 
     LEADING_JOIN(RuleTypeClass.REWRITE),
     REWRITE_SENTINEL(RuleTypeClass.REWRITE),

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

@@ -0,0 +1,41 @@
+// 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.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.rules.Rule;
+import org.apache.doris.nereids.rules.RuleType;
+import org.apache.doris.nereids.trees.expressions.Slot;
+
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Collect outputs of CTE Consumer.
+ */
+public class CollectCteConsumerOutput extends OneRewriteRuleFactory {
+
+    @Override
+    public Rule build() {
+        return logicalCTEConsumer().thenApply(ctx -> {
+            Set<Slot> producerOutputs = ctx.statementContext
+                    .getCteIdToOutputIds().computeIfAbsent(ctx.root.getCteId(), k -> new HashSet<>());
+            producerOutputs.addAll(ctx.root.getProducerToConsumerOutputMap().keySet());
+            return null;
+        }).toRule(RuleType.COLLECT_CTE_CONSUMER_OUTPUT);
+    }
+}