From e7982dc52f9dfb4562026f3fccdd85b974864efa Mon Sep 17 00:00:00 2001 From: feiniaofeiafei Date: Tue, 25 Aug 2026 12:11:44 +0800 Subject: [PATCH] [fix](fe) Preserve partition pruning state in plan copies ### What problem does this PR solve? Issue Number: N/A Related PR: N/A Problem Summary: LogicalOlapScan.withRelationId reset partitionPruned and hasPartitionPredicate as a hidden side effect, which broke default deep-copy semantics. Preserve both fields while replacing the relation ID. Add an explicit DeepCopierContext option that StructInfo enables for MTMV union compensation, so copied OLAP scans invalidate only partitionPruned during the existing copy traversal and the following whole-tree rewrite reapplies partition pruning. ### Release note None ### Check List (For Author) - Test: Unit Test - ./run-fe-ut.sh --run org.apache.doris.nereids.trees.copier.LogicalPlanDeepCopierTest,org.apache.doris.nereids.mv.StructInfoTest - mvn checkstyle:check -pl fe-core - Behavior changed: Yes. Default plan copies preserve partition pruning metadata, while MTMV union compensation explicitly invalidates partitionPruned during its copy. - Does this need documentation: No --- .../rules/exploration/mv/StructInfo.java | 6 +++- .../trees/copier/DeepCopierContext.java | 10 +++++++ .../trees/copier/LogicalPlanDeepCopier.java | 4 +++ .../trees/plans/logical/LogicalOlapScan.java | 18 +++++++++-- .../logical/LogicalOlapTableStreamScan.java | 13 ++++++++ .../copier/LogicalPlanDeepCopierTest.java | 30 +++++++++++++++++++ 6 files changed, 78 insertions(+), 3 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java index d4e0e3f50031d0..e5109b5516cc89 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java @@ -975,8 +975,12 @@ public static Pair 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(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/DeepCopierContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/DeepCopierContext.java index c4772fb7e819ed..b898a08f7f1bbe 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/DeepCopierContext.java +++ b/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 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); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java index 29fd9e9e4c5c34..01472f1c920671 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java +++ b/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 virtualColumns = catalogRelation.getVirtualColumns().stream() .map(e -> { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java index 94bed913bcb8a6..b8faa2473c49b8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java @@ -656,11 +656,10 @@ public LogicalOlapScan withManuallySpecifiedTabletIds(List 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 getSelectedTabletIds() { return selectedTabletIds; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapTableStreamScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapTableStreamScan.java index 725bf059c55455..156433924eff8b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapTableStreamScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapTableStreamScan.java @@ -300,6 +300,19 @@ public LogicalOlapTableStreamScan withSelectedPartitionIds(List 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 diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java index 98f07698a9ef29..b139d05910ad7c 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java +++ b/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);