From e737a7ca41953836920ecdd998db48e5cc2b7494 Mon Sep 17 00:00:00 2001 From: seawinde Date: Mon, 3 Aug 2026 11:40:10 +0800 Subject: [PATCH 1/3] [fix](fe) Clear stale partition state for prepared statements ### What problem does this PR solve? Issue Number: N/A Related PR: N/A Problem Summary: Server prepared statements reuse a StatementContext across executions. Query partition and relation mappings from prior executions therefore accumulate, while partition compensation allocates a temporary BitSet for each stored RelationId. Clear the per-execution mappings before planning each EXECUTE and query the existing BitSet directly. ### Release note Fixed excessive FE CPU and memory usage when repeatedly executing prepared statements with materialized view rewrite. ### Check List (For Author) - Test: Unit Test - ExecuteCommandTest - PartitionCompensatorTest - Behavior changed: No - Does this need documentation: No --- .../rules/exploration/mv/PartitionCompensator.java | 4 +--- .../trees/plans/commands/ExecuteCommand.java | 4 ++++ .../trees/plans/commands/ExecuteCommandTest.java | 14 +++++++++++++- 3 files changed, 18 insertions(+), 4 deletions(-) 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..d020f34d1fd4d6 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,10 @@ 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 + // information collected by the previous execution before planning the current one. + statementContext.getTableUsedPartitionNameMap().clear(); + statementContext.getCommonTableIdToRelationIdMap().clear(); 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..3859c8de0cd1a7 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 @@ -20,6 +20,7 @@ import org.apache.doris.analysis.TableScanParams; import org.apache.doris.catalog.DatabaseIf; 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; @@ -27,6 +28,7 @@ import org.apache.doris.nereids.analyzer.UnboundRelation; import org.apache.doris.nereids.parser.NereidsParser; 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.qe.ConnectContext; @@ -47,7 +49,7 @@ public class ExecuteCommandTest { @Test - public void testResolvedScanOptionsAreResetForEveryExecute() throws Exception { + public void testExecutionStateIsResetForEveryExecute() throws Exception { String sql = "select * from p@options('scan.mode'='latest')"; LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); UnboundRelation relation = logicalPlan.collectToList( @@ -69,10 +71,20 @@ public void testResolvedScanOptionsAreResetForEveryExecute() throws Exception { Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); Mockito.when(executor.getContext()).thenReturn(connectContext); + statementContext.getTableUsedPartitionNameMap().put( + Collections.singletonList("table"), Pair.of(new RelationId(0), Collections.singleton("p1"))); + statementContext.getCommonTableIdToRelationIdMap().put(0, 0); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + Assertions.assertTrue(statementContext.getTableUsedPartitionNameMap().isEmpty()); + Assertions.assertTrue(statementContext.getCommonTableIdToRelationIdMap().isEmpty()); Assertions.assertEquals("2", resolveNextSnapshot(scanParams, snapshotId)); + statementContext.getTableUsedPartitionNameMap().put( + Collections.singletonList("table"), Pair.of(new RelationId(1), Collections.singleton("p2"))); + statementContext.getCommonTableIdToRelationIdMap().put(0, 1); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); + Assertions.assertTrue(statementContext.getTableUsedPartitionNameMap().isEmpty()); + Assertions.assertTrue(statementContext.getCommonTableIdToRelationIdMap().isEmpty()); Assertions.assertEquals("3", resolveNextSnapshot(scanParams, snapshotId)); Mockito.verify(executor, Mockito.times(2)).execute(); } From 5b4221b695c8ede1682b83dafff68599cc4da966 Mon Sep 17 00:00:00 2001 From: seawinde Date: Mon, 3 Aug 2026 17:04:33 +0800 Subject: [PATCH 2/3] [test](fe) Cover prepared statement partition state reset ### What problem does this PR solve? Issue Number: N/A Related PR: #49514, #58643 Problem Summary: The existing unit test only checked stale partition mappings after the mocked executor returned, so it could not distinguish cleanup before planning from cleanup after execution. Add a focused test that simulates current planning state being written during each EXECUTE and verifies that only the current RelationId remains. ### Release note None ### Check List (For Author) - Test: Unit Test - ./run-fe-ut.sh --run org.apache.doris.nereids.trees.plans.commands.ExecuteCommandTest - Behavior changed: No - Does this need documentation: No --- .../plans/commands/ExecuteCommandTest.java | 59 +++++++++++++++---- 1 file changed, 48 insertions(+), 11 deletions(-) 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 3859c8de0cd1a7..e31da7aefd5c15 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 @@ -43,13 +43,14 @@ import org.mockito.Mockito; import java.util.Collections; +import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; public class ExecuteCommandTest { @Test - public void testExecutionStateIsResetForEveryExecute() throws Exception { + public void testResolvedScanOptionsAreResetForEveryExecute() throws Exception { String sql = "select * from p@options('scan.mode'='latest')"; LogicalPlan logicalPlan = new NereidsParser().parseSingle(sql); UnboundRelation relation = logicalPlan.collectToList( @@ -71,24 +72,60 @@ public void testExecutionStateIsResetForEveryExecute() throws Exception { Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); Mockito.when(executor.getContext()).thenReturn(connectContext); - statementContext.getTableUsedPartitionNameMap().put( - Collections.singletonList("table"), Pair.of(new RelationId(0), Collections.singleton("p1"))); - statementContext.getCommonTableIdToRelationIdMap().put(0, 0); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - Assertions.assertTrue(statementContext.getTableUsedPartitionNameMap().isEmpty()); - Assertions.assertTrue(statementContext.getCommonTableIdToRelationIdMap().isEmpty()); Assertions.assertEquals("2", resolveNextSnapshot(scanParams, snapshotId)); - statementContext.getTableUsedPartitionNameMap().put( - Collections.singletonList("table"), Pair.of(new RelationId(1), Collections.singleton("p2"))); - statementContext.getCommonTableIdToRelationIdMap().put(0, 1); new ExecuteCommand("stmt", prepareCommand, statementContext).run(connectContext, executor); - Assertions.assertTrue(statementContext.getTableUsedPartitionNameMap().isEmpty()); - Assertions.assertTrue(statementContext.getCommonTableIdToRelationIdMap().isEmpty()); Assertions.assertEquals("3", resolveNextSnapshot(scanParams, snapshotId)); 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 testResolvedScanOptionsAreResetForPreparedDeleteUsing() throws Exception { String sql = "delete from target using source@options('scan.mode'='latest') " From b959429e41cfff2cea5f642b4deb8d9cd99e7ee9 Mon Sep 17 00:00:00 2001 From: seawinde Date: Tue, 4 Aug 2026 14:47:38 +0800 Subject: [PATCH 3/3] [fix](fe) Reset MV state for prepared executions ### What problem does this PR solve? Issue Number: N/A Related PR: #66370 Problem Summary: Server prepared statements reuse one StatementContext across external EXECUTEs. Materialized-view partition caches, pre-rewrite plans, phase decisions, hints, timing, and MV relation statistics from an earlier execution could therefore affect later planning or grow without bound. Reset only the partition and materialized-view planning state at the external EXECUTE boundary while preserving unrelated prepared-statement state. ### Release note Fixed stale materialized-view planning state across repeated prepared statement executions. ### Check List (For Author) - Test: Unit Test - ./run-fe-ut.sh --run org.apache.doris.nereids.trees.plans.commands.ExecuteCommandTest - Behavior changed: Yes (each prepared EXECUTE independently computes partition and materialized-view rewrite state) - Does this need documentation: No --- .../doris/nereids/StatementContext.java | 16 ++++ .../trees/plans/commands/ExecuteCommand.java | 7 +- .../plans/commands/ExecuteCommandTest.java | 75 +++++++++++++++++++ 3 files changed, 94 insertions(+), 4 deletions(-) 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/trees/plans/commands/ExecuteCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExecuteCommand.java index d020f34d1fd4d6..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,10 +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 - // information collected by the previous execution before planning the current one. - statementContext.getTableUsedPartitionNameMap().clear(); - statementContext.getCommonTableIdToRelationIdMap().clear(); + // 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 e31da7aefd5c15..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,23 +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; @@ -126,6 +135,55 @@ public void testPartitionStateIsResetForEveryExecute() throws Exception { 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') " @@ -229,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(