[feat](Nereids) after partition prune, output rows of scan node only contains rows from selected partitions

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

@@ -18,7 +18,6 @@
 package org.apache.doris.nereids;
 
 import org.apache.doris.analysis.StatementBase;
-import org.apache.doris.catalog.Column;
 import org.apache.doris.catalog.TableIf;
 import org.apache.doris.catalog.constraint.TableIdentifier;
 import org.apache.doris.common.Id;
@@ -32,6 +31,7 @@
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.Placeholder;
 import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
 import org.apache.doris.nereids.trees.plans.ObjectId;
 import org.apache.doris.nereids.trees.plans.PlaceholderId;
@@ -135,7 +135,7 @@ public class StatementContext implements Closeable {
     private final Map<Slot, Relation> slotToRelation = Maps.newHashMap();
 
     // the columns in Plan.getExpressions(), such as columns in join condition or filter condition, group by expression
-    private final Set<Column> keyColumns = Sets.newHashSet();
+    private final Set<SlotReference> keySlots = Sets.newHashSet();
     private BitSet disableRules;
 
     // table locks
@@ -516,12 +516,12 @@ public String toString() {
         }
     }
 
-    public void addKeyColumn(Column column) {
-        keyColumns.add(column);
+    public void addKeySlot(SlotReference slot) {
+        keySlots.add(slot);
     }
 
-    public boolean isKeyColumn(Column column) {
-        return keyColumns.contains(column);
+    public boolean isKeySlot(SlotReference slot) {
+        return keySlots.contains(slot);
     }
 
     /** Get table id with lazy */

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

@@ -138,7 +138,7 @@ public Plan rewriteRoot(Plan plan, JobContext jobContext) {
             if (stmtContext != null) {
                 for (Slot key : keys) {
                     if (key instanceof SlotReference) {
-                        ((SlotReference) key).getColumn().ifPresent(stmtContext::addKeyColumn);
+                        stmtContext.addKeySlot((SlotReference) key);
                     }
                 }
             }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/OlapScan.java

@@ -28,6 +28,12 @@ public interface OlapScan {
 
     long getSelectedIndexId();
 
+    /**
+     * if this is mv, return selectedIndexId, o.w -1
+     * @return -1 or selectedIndexId
+     */
+    long getSelectedIndexIdForMV();
+
     List<Long> getSelectedPartitionIds();
 
     List<Long> getSelectedTabletIds();

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

@@ -93,6 +93,11 @@ public long getSelectedIndexId() {
         return logicalOlapScan.getSelectedIndexId();
     }
 
+    @Override
+    public long getSelectedIndexIdForMV() {
+        return logicalOlapScan.getSelectedIndexIdForMV();
+    }
+
     @Override
     public List<Long> getSelectedPartitionIds() {
         return logicalOlapScan.getSelectedPartitionIds();

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

@@ -333,6 +333,14 @@ public long getSelectedIndexId() {
         return selectedIndexId;
     }
 
+    @Override
+    public long getSelectedIndexIdForMV() {
+        if (getTable().getBaseIndexId() != selectedIndexId) {
+            return selectedIndexId;
+        }
+        return -1;
+    }
+
     public boolean isIndexSelected() {
         return indexSelected;
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalDeferMaterializeOlapScan.java

@@ -90,6 +90,11 @@ public long getSelectedIndexId() {
         return physicalOlapScan.getSelectedIndexId();
     }
 
+    @Override
+    public long getSelectedIndexIdForMV() {
+        return physicalOlapScan.getSelectedIndexIdForMV();
+    }
+
     @Override
     public List<Long> getSelectedPartitionIds() {
         return physicalOlapScan.getSelectedPartitionIds();

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java

@@ -90,6 +90,14 @@ public long getSelectedIndexId() {
         return selectedIndexId;
     }
 
+    @Override
+    public long getSelectedIndexIdForMV() {
+        if (getTable().getBaseIndexId() != selectedIndexId) {
+            return selectedIndexId;
+        }
+        return -1;
+    }
+
     @Override
     public List<Long> getSelectedTabletIds() {
         return selectedTabletIds;

fe/fe-core/src/main/java/org/apache/doris/statistics/StatisticsBuilder.java

@@ -21,11 +21,12 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 public class StatisticsBuilder {
 
     private double rowCount;
-    private int widthInJoinCluster;
+    private int widthInJoinCluster = 1;
     private final Map<Expression, ColumnStatistic> expressionToColumnStats;
 
     public StatisticsBuilder() {
@@ -60,6 +61,10 @@ public StatisticsBuilder putColumnStatistics(Expression expression, ColumnStatis
         return this;
     }
 
+    public Set<Map.Entry<Expression, ColumnStatistic>> getExpressionColumnStatsEntries() {
+        return expressionToColumnStats.entrySet();
+    }
+
     public Statistics build() {
         return new Statistics(rowCount, widthInJoinCluster, expressionToColumnStats);
     }

regression-test/suites/nereids_p0/stats/partition_col_stats.groovy

@@ -36,7 +36,7 @@ suite("partition_col_stats") {
     """
     //run this sql to make stats be cached
     sql "select * from pt where k1<3;"
-    sleep(10)
+    sleep(10000)
     explain{
         sql "physical plan select * from pt where k1<3;"
         contains("stats=4")