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);