[fix](partition pruning) Preserve partition pruning state in plan copies

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java

@@ -975,8 +975,12 @@ public static Pair<Plan, Boolean> addFilterOnTableScan(Plan queryPlan,
         }
         // Deep copy the plan to avoid the plan output is the same with the later union output, this may cause
         // exec by mistake
+        DeepCopierContext deepCopierContext = new DeepCopierContext();
+        // The compensation filter changes the partition range. Invalidate pruning during the copy to avoid
+        // another traversal before the following whole-tree rewrite applies partition pruning again.
+        deepCopierContext.setInvalidatePartitionPruning(true);
         queryPlanWithUnionFilter = new LogicalPlanDeepCopier().deepCopy(
-                (LogicalPlan) queryPlanWithUnionFilter, new DeepCopierContext());
+                (LogicalPlan) queryPlanWithUnionFilter, deepCopierContext);
         // rbo rewrite after adding filter on origin plan
         return Pair.of(MaterializedViewUtils.rewriteByRules(parentCascadesContext, context -> {
             Rewriter.getWholeTreeRewriter(context).execute();

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/DeepCopierContext.java

@@ -39,6 +39,16 @@ public class DeepCopierContext {
      * to ensure they keep same after deep copy
      */
     private final Map<RelationId, LogicalRelation> relationReplaceMap = Maps.newHashMap();
+    /** Whether copied OLAP scans should be marked for partition pruning again. */
+    private boolean invalidatePartitionPruning = false;
+
+    public void setInvalidatePartitionPruning(boolean invalidatePartitionPruning) {
+        this.invalidatePartitionPruning = invalidatePartitionPruning;
+    }
+
+    public boolean shouldInvalidatePartitionPruning() {
+        return invalidatePartitionPruning;
+    }
 
     public void putRelation(RelationId relationId, LogicalRelation newRelation) {
         relationReplaceMap.put(relationId, newRelation);

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java

@@ -46,6 +46,7 @@
 import org.apache.doris.nereids.trees.plans.logical.LogicalIntersect;
 import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
 import org.apache.doris.nereids.trees.plans.logical.LogicalLimit;
+import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPartitionTopN;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
@@ -107,6 +108,9 @@ public Plan visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation,
         }
         LogicalCatalogRelation newRelation =
                 catalogRelation.withRelationId(StatementScopeIdGenerator.newRelationId());
+        if (context.shouldInvalidatePartitionPruning() && newRelation instanceof LogicalOlapScan) {
+            newRelation = ((LogicalOlapScan) newRelation).withPartitionPruned(false);
+        }
         updateReplaceMapWithOutput(catalogRelation, newRelation, context.exprIdReplaceMap);
         List<NamedExpression> virtualColumns = catalogRelation.getVirtualColumns().stream()
                 .map(e -> {

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

@@ -656,11 +656,10 @@ public LogicalOlapScan withManuallySpecifiedTabletIds(List<Long> manuallySpecifi
 
     @Override
     public LogicalOlapScan withRelationId(RelationId relationId) {
-        // we have to set partitionPruned to false, so that mtmv rewrite can prevent deadlock when rewriting union
         return AbstractPlan.copyWithSameId(this, () ->
                 new LogicalOlapScan(relationId, (Table) table, qualifier,
                 Optional.empty(), Optional.empty(),
-                selectedPartitionIds, false, false, selectedTabletIds,
+                selectedPartitionIds, partitionPruned, hasPartitionPredicate, selectedTabletIds,
                 selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions,
                 hints, Maps.newHashMap(), Optional.empty(), tableSample, directMvScan,
                 colToSubPathsMap, selectedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys,
@@ -763,6 +762,21 @@ public boolean isPartitionPruned() {
         return partitionPruned;
     }
 
+    /**
+     * Return a new scan with the specified partition pruning state.
+     */
+    public LogicalOlapScan withPartitionPruned(boolean partitionPruned) {
+        return AbstractPlan.copyWithSameId(this, () ->
+                new LogicalOlapScan(relationId, (Table) table, qualifier,
+                Optional.empty(), Optional.of(getLogicalProperties()),
+                selectedPartitionIds, partitionPruned, hasPartitionPredicate, selectedTabletIds,
+                selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions,
+                hints, cacheSlotWithSlotName, cachedOutput, tableSample, directMvScan,
+                colToSubPathsMap, manuallySpecifiedTabletIds, operativeSlots, virtualColumns,
+                scoreOrderKeys, scoreLimit, scoreRangeInfo, annOrderKeys, annLimit, tableAlias,
+                partitionPrunablePredicates, scanParams));
+    }
+
     public List<Long> getSelectedTabletIds() {
         return selectedTabletIds;
     }

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

@@ -300,6 +300,19 @@ public LogicalOlapTableStreamScan withSelectedPartitionIds(List<Long> selectedPa
                         partitionPrunablePredicates, scanParams, readMode));
     }
 
+    @Override
+    public LogicalOlapTableStreamScan withPartitionPruned(boolean partitionPruned) {
+        return AbstractPlan.copyWithSameId(this, () ->
+                new LogicalOlapTableStreamScan(relationId, (Table) table, qualifier,
+                        groupExpression, Optional.of(getLogicalProperties()),
+                        selectedPartitionIds, partitionPruned, hasPartitionPredicate, selectedTabletIds,
+                        selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions,
+                        hints, cacheSlotWithSlotName, cachedOutput, tableSample, directMvScan,
+                        colToSubPathsMap, manuallySpecifiedTabletIds, operativeSlots, virtualColumns,
+                        scoreOrderKeys, scoreLimit, scoreRangeInfo, annOrderKeys, annLimit, tableAlias,
+                        partitionPrunablePredicates, scanParams, readMode));
+    }
+
     /**
      * Returns a new {@code LogicalOlapScan} carrying the supplied
      * {@link PartitionPrunablePredicate}. It is preserved across all other

fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java

@@ -50,6 +50,36 @@ public void testDeepCopyOlapScan() {
         }
     }
 
+    @Test
+    public void testDeepCopyOlapScanPreservesPartitionPruningState() {
+        LogicalOlapScan relationPlan = PlanConstructor.newLogicalOlapScan(0, "a", 0);
+        relationPlan = relationPlan.withSelectedPartitionIds(relationPlan.getSelectedPartitionIds(), true);
+
+        LogicalOlapScan copiedPlan =
+                (LogicalOlapScan) relationPlan.accept(LogicalPlanDeepCopier.INSTANCE, new DeepCopierContext());
+
+        Assertions.assertTrue(relationPlan.isPartitionPruned());
+        Assertions.assertTrue(relationPlan.hasPartitionPredicate());
+        Assertions.assertTrue(copiedPlan.isPartitionPruned());
+        Assertions.assertTrue(copiedPlan.hasPartitionPredicate());
+    }
+
+    @Test
+    public void testDeepCopyOlapScanInvalidatesPartitionPruning() {
+        LogicalOlapScan relationPlan = PlanConstructor.newLogicalOlapScan(0, "a", 0);
+        relationPlan = relationPlan.withSelectedPartitionIds(relationPlan.getSelectedPartitionIds(), true);
+        DeepCopierContext context = new DeepCopierContext();
+        context.setInvalidatePartitionPruning(true);
+
+        LogicalOlapScan copiedPlan =
+                (LogicalOlapScan) relationPlan.accept(LogicalPlanDeepCopier.INSTANCE, context);
+
+        Assertions.assertTrue(relationPlan.isPartitionPruned());
+        Assertions.assertTrue(relationPlan.hasPartitionPredicate());
+        Assertions.assertFalse(copiedPlan.isPartitionPruned());
+        Assertions.assertTrue(copiedPlan.hasPartitionPredicate());
+    }
+
     @Test
     public void testDeepCopyOlapScanWithNonFirstOperativeSlot() {
         LogicalOlapScan relationPlan = PlanConstructor.newLogicalOlapScan(0, "a", 0);