column prune

englefly · Jan 9, 2025 · 1ecf72f · 1ecf72f
1 parent 8fc5aab
commit 1ecf72f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/.../main/java/org/apache/doris/nereids/processor/post/ProjectAggregateExpressionsForCse.java b/.../main/java/org/apache/doris/nereids/processor/post/ProjectAggregateExpressionsForCse.java
@@ -46,6 +46,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * create project under aggregate to enable CSE
@@ -102,10 +103,13 @@ public Plan visitPhysicalHashAggregate(PhysicalHashAggregate<? extends Plan> agg
         }
 
         if (aggregate.child() instanceof PhysicalProject) {
-            PhysicalProject<? extends Plan> project = (PhysicalProject<? extends Plan>) aggregate.child();
             List<NamedExpression> newProjections = Lists.newArrayList();
+            PhysicalProject<? extends Plan> project = (PhysicalProject<? extends Plan>) aggregate.child();
+            Set<Slot> newInputSlots = aggOutputReplaced.stream()
+                    .flatMap(expr -> expr.getInputSlots().stream())
+                    .collect(Collectors.toSet());
             for (NamedExpression expr : project.getProjects()) {
-                if (!(expr instanceof SlotReference) || aggOutputReplaced.contains(expr)) {
+                if (!(expr instanceof SlotReference) || newInputSlots.contains(expr)) {
                     newProjections.add(expr);
                 }
             }