diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index 90a1edb566d61f..ed70f47da1fc10 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -1443,6 +1443,22 @@ public Map> getMvCanRewritePartitionsMap() return mvCanRewritePartitionsMap; } + /** Clear materialized-view planning state retained by a prepared statement between executions. */ + public void resetMaterializedViewStateForPreparedExecution() { + tableUsedPartitionNameMap.clear(); + commonTableIdToRelationIdToMap.clear(); + mvCanRewritePartitionsMap.clear(); + materializedViewRewriteDuration = 0; + hints.removeIf(UseMvHint.class::isInstance); + tmpPlanForMvRewrite.clear(); + rewrittenPlansByMv.clear(); + needPreMvRewriteRuleMasks.clear(); + needPreMvRewrite = false; + preMvRewritten = false; + materializationRewrittenSuccessSet.clear(); + relationIdToStatisticsMap.clear(); + } + public void setPrepareStage(boolean isPrepare) { this.prepareStage = isPrepare; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java index 9064049d3b801b..217535e6aba8ed 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/PartitionCompensator.java @@ -298,10 +298,8 @@ public static Map, Set> getQueryUsedPartitions(StatementCon continue tableLoop; } // If currentUsedRelationIdSet is not empty, need check relation id to get concrete used partitions - BitSet usedPartitionRelation = new BitSet(); - usedPartitionRelation.set(tableUsedPartitionPair.key().asInt()); if (!currentUsedRelationIdSet.isEmpty() - && !currentUsedRelationIdSet.intersects(usedPartitionRelation)) { + && !currentUsedRelationIdSet.get(tableUsedPartitionPair.key().asInt())) { continue; } usedPartitionSet.addAll(tableUsedPartitionPair.value()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java index 21f6c4e5063be9..37b0a20b647b39 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java @@ -91,6 +91,9 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { } PrepareCommand prepareCommand = preparedStmtCtx.command; StatementContext statementContext = preparedStmtCtx.getStatementContext(); + // Prepared statements reuse StatementContext across executions. Discard partition and MV + // planning results collected by the previous execution before planning the current one. + statementContext.resetMaterializedViewStateForPreparedExecution(); statementContext.setPrepareStage(false); statementContext.setIsInsert(false); // A prepared EXECUTE reuses this one StatementContext across executions; drop the connector diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java index 5e1d2506b3bb11..a8a31fb76b8159 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommandTest.java @@ -19,21 +19,32 @@ import org.apache.doris.analysis.TableScanParams; import org.apache.doris.catalog.DatabaseIf; +import org.apache.doris.catalog.Partition; import org.apache.doris.catalog.TableIf; +import org.apache.doris.common.Pair; import org.apache.doris.datasource.CatalogIf; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.datasource.mvcc.MvccTable; +import org.apache.doris.mtmv.BaseTableInfo; +import org.apache.doris.nereids.NereidsPlanner; import org.apache.doris.nereids.StatementContext; import org.apache.doris.nereids.analyzer.UnboundRelation; +import org.apache.doris.nereids.glue.LogicalPlanAdapter; +import org.apache.doris.nereids.hint.Hint; +import org.apache.doris.nereids.hint.UseMvHint; import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.rules.RuleType; import org.apache.doris.nereids.trees.expressions.SubqueryExpr; +import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.commands.merge.MergeIntoCommand; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.OriginStatement; import org.apache.doris.qe.PreparedStatementContext; import org.apache.doris.qe.SessionVariable; import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.statistics.Statistics; import com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Assertions; @@ -41,6 +52,7 @@ import org.mockito.Mockito; import java.util.Collections; +import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; @@ -77,6 +89,101 @@ public void testResolvedScanOptionsAreResetForEveryExecute() throws Exception { Mockito.verify(executor, Mockito.times(2)).execute(); } + @Test + public void testPartitionStateIsResetForEveryExecute() throws Exception { + String sql = "select 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + StatementContext statementContext = new StatementContext(); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + Mockito.when(connectContext.getPreparedStementContext("stmt")).thenReturn(preparedStatement); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + Mockito.when(executor.getContext()).thenReturn(connectContext); + + List tableQualifier = Collections.singletonList("table"); + AtomicInteger relationId = new AtomicInteger(); + Mockito.doAnswer(invocation -> { + int currentRelationId = relationId.getAndIncrement(); + statementContext.getTableUsedPartitionNameMap().put(tableQualifier, + Pair.of(new RelationId(currentRelationId), Collections.singleton("p"))); + statementContext.getCommonTableIdToRelationIdMap().put(0, currentRelationId); + return null; + }).when(executor).execute(); + + statementContext.getTableUsedPartitionNameMap().put( + tableQualifier, Pair.of(new RelationId(100), Collections.singleton("old"))); + statementContext.getCommonTableIdToRelationIdMap().put(0, 100); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + Assertions.assertEquals(1, statementContext.getTableUsedPartitionNameMap().size()); + Assertions.assertEquals(0, statementContext.getTableUsedPartitionNameMap() + .get(tableQualifier).iterator().next().key().asInt()); + Assertions.assertEquals(Collections.singleton(0), + statementContext.getCommonTableIdToRelationIdMap().get(0)); + + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + Assertions.assertEquals(1, statementContext.getTableUsedPartitionNameMap().size()); + Assertions.assertEquals(1, statementContext.getTableUsedPartitionNameMap() + .get(tableQualifier).iterator().next().key().asInt()); + Assertions.assertEquals(Collections.singleton(1), + statementContext.getCommonTableIdToRelationIdMap().get(0)); + Mockito.verify(executor, Mockito.times(2)).execute(); + } + + @Test + public void testMaterializedViewStateIsResetForEveryExecute() throws Exception { + String sql = "select 1"; + LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); + ConnectContext connectContext = MemoTestUtils.createConnectContext(); + StatementContext statementContext = new StatementContext( + connectContext, new OriginStatement(sql, 0)); + connectContext.setStatementContext(statementContext); + PrepareCommand prepareCommand = new PrepareCommand( + "stmt", logicalPlan, Collections.emptyList(), new OriginStatement(sql, 0)); + PreparedStatementContext preparedStatement = new PreparedStatementContext( + prepareCommand, connectContext, statementContext, "stmt"); + StmtExecutor executor = Mockito.mock(StmtExecutor.class); + connectContext.addPreparedStatementContext("stmt", preparedStatement); + Mockito.when(executor.getContext()).thenReturn(connectContext); + + Hint retainedHint = new Hint("Distribute"); + statementContext.addHint(retainedHint); + statementContext.setForceRecordTmpPlan(true); + AtomicInteger executionCount = new AtomicInteger(); + Mockito.doAnswer(invocation -> { + Assertions.assertTrue(statementContext.getTableUsedPartitionNameMap().isEmpty()); + Assertions.assertTrue(statementContext.getCommonTableIdToRelationIdMap().isEmpty()); + Assertions.assertTrue(statementContext.getMvCanRewritePartitionsMap().isEmpty()); + Assertions.assertEquals(0, statementContext.getMaterializedViewRewriteDuration()); + Assertions.assertEquals(Collections.singletonList(retainedHint), statementContext.getHints()); + Assertions.assertTrue(statementContext.getTmpPlanForMvRewrite().isEmpty()); + Assertions.assertTrue(statementContext.getRewrittenPlansByMv().isEmpty()); + Assertions.assertTrue(statementContext.getNeedPreMvRewriteRuleMasks().isEmpty()); + Assertions.assertFalse(statementContext.isNeedPreMvRewrite()); + Assertions.assertFalse(statementContext.isPreMvRewritten()); + Assertions.assertTrue(statementContext.getMaterializationRewrittenSuccessSet().isEmpty()); + Assertions.assertTrue(statementContext.getRelationIdToStatisticsMap().isEmpty()); + Assertions.assertTrue(statementContext.isForceRecordTmpPlan()); + NereidsPlanner planner = new NereidsPlanner(statementContext); + planner.plan(new LogicalPlanAdapter(logicalPlan, statementContext)); + Assertions.assertNotNull(planner.getPhysicalPlan()); + populateMaterializedViewState(statementContext, logicalPlan); + executionCount.incrementAndGet(); + return null; + }).when(executor).execute(); + + populateMaterializedViewState(statementContext, logicalPlan); + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + + Assertions.assertEquals(2, executionCount.get()); + } + @Test public void testResolvedScanOptionsAreResetForPreparedDeleteUsing() throws Exception { String sql = "delete from target using source@options('scan.mode'='latest') " @@ -180,6 +287,23 @@ private String resolveNextSnapshot(TableScanParams scanParams, AtomicInteger sna .get("scan.snapshot-id"); } + private void populateMaterializedViewState(StatementContext statementContext, LogicalPlan logicalPlan) { + statementContext.getTableUsedPartitionNameMap().put(Collections.singletonList("table"), + Pair.of(new RelationId(1), Collections.singleton("partition"))); + statementContext.getCommonTableIdToRelationIdMap().put(1, 1); + statementContext.getMvCanRewritePartitionsMap().put(Mockito.mock(BaseTableInfo.class), + Collections.singleton(Mockito.mock(Partition.class))); + statementContext.addMaterializedViewRewriteDuration(1); + statementContext.addHint(Mockito.mock(UseMvHint.class)); + statementContext.addTmpPlanForMvRewrite(logicalPlan); + statementContext.addRewrittenPlanByMv(logicalPlan); + statementContext.ruleSetApplied(RuleType.REORDER_JOIN); + statementContext.setNeedPreMvRewrite(true); + statementContext.setPreMvRewritten(true); + statementContext.addMaterializationRewrittenSuccess(Collections.singletonList("mv")); + statementContext.addStatistics(new RelationId(1), Mockito.mock(Statistics.class)); + } + private void assertPreparedCommandResetsScanOptions( String sql, LogicalPlan command, LogicalPlan relationRoot) throws Exception { UnboundRelation relation = relationRoot.collectToList(