diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementIT.java
index c7c5902e8e2dc..09da6f3ac9e07 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBTablePreparedStatementIT.java
@@ -190,6 +190,37 @@ public void testPrepareAndExecuteMultipleTimes() {
}
}
+ @Test
+ public void testMaterializedCteIsRefetchedForEachExecute() {
+ String[] expectedHeader = new String[] {"_col0"};
+ try (Connection connection = EnvFactory.getEnv().getTableConnection();
+ Statement statement = connection.createStatement()) {
+ statement.execute("USE " + DATABASE_NAME);
+ statement.execute(
+ "PREPARE materialized_cte_stmt FROM "
+ + "WITH cte AS MATERIALIZED (SELECT * FROM test_table) SELECT COUNT(*) FROM cte");
+
+ executePreparedStatementAndVerify(
+ connection,
+ statement,
+ "EXECUTE materialized_cte_stmt",
+ expectedHeader,
+ new String[] {"5,"});
+
+ statement.execute("INSERT INTO test_table VALUES (2025-01-01T00:05:00, 6, 'Frank', 600.1)");
+
+ executePreparedStatementAndVerify(
+ connection,
+ statement,
+ "EXECUTE materialized_cte_stmt",
+ expectedHeader,
+ new String[] {"6,"});
+ statement.execute("DEALLOCATE PREPARE materialized_cte_stmt");
+ } catch (SQLException e) {
+ fail("testMaterializedCteIsRefetchedForEachExecute failed: " + e.getMessage());
+ }
+ }
+
@Test
public void testPrepareWithMultipleParameters() {
String[] expectedHeader = new String[] {"time", "id", "name", "value"};
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/CteMaterializationContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/CteMaterializationContext.java
new file mode 100644
index 0000000000000..88dffc4d6aacb
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/CteMaterializationContext.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.common;
+
+import org.apache.iotdb.commons.queryengine.plan.relational.analyzer.NodeRef;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Table;
+import org.apache.iotdb.commons.queryengine.utils.cte.CteDataStore;
+
+import javax.annotation.Nullable;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * Holds CTE materialization state for one top-level query execution.
+ *
+ *
A relational {@link Query} AST can be cached by a prepared statement and reused by independent
+ * executions. Materialized data therefore cannot be stored on that AST. This context is owned by an
+ * {@link MPPQueryContext} instead and is shared only with inner queries spawned while executing the
+ * same top-level query. It is designed to be accessed only by the planner thread.
+ */
+public final class CteMaterializationContext {
+
+ private final Map, Query> cteQueries = new HashMap<>();
+
+ // Optional.empty() records an attempted materialization that fell back to inline planning.
+ private final Map, Optional> materializationResults =
+ new HashMap<>();
+
+ public void addCteQuery(Table table, Query query) {
+ cteQueries.put(NodeRef.of(table), query);
+ }
+
+ public Map, Query> getCteQueries() {
+ return Collections.unmodifiableMap(cteQueries);
+ }
+
+ public boolean isMaterializationAttempted(Query query) {
+ return materializationResults.containsKey(NodeRef.of(query));
+ }
+
+ public void recordMaterializationResult(Query query, @Nullable CteDataStore dataStore) {
+ materializationResults.put(NodeRef.of(query), Optional.ofNullable(dataStore));
+ }
+
+ @Nullable
+ public CteDataStore getCteDataStore(Query query) {
+ return materializationResults.getOrDefault(NodeRef.of(query), Optional.empty()).orElse(null);
+ }
+
+ @Nullable
+ public CteDataStore getCteDataStore(Table table) {
+ Query query = cteQueries.get(NodeRef.of(table));
+ return query == null ? null : getCteDataStore(query);
+ }
+
+ public void clear() {
+ cteQueries.clear();
+ materializationResults.clear();
+ }
+}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
index d3926bcac79ce..750e92dcb8628 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
@@ -57,6 +57,8 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import javax.annotation.Nullable;
+
import java.nio.file.Paths;
import java.time.ZoneId;
import java.util.Collections;
@@ -171,7 +173,9 @@ public enum ExplainType {
private boolean debug = false;
- private Map, Query> cteQueries = new HashMap<>();
+ // This state belongs to one top-level execution. Inner CTE and scalar-subquery executions share
+ // the same instance so they can consume materialized data without putting runtime state on ASTs.
+ private CteMaterializationContext cteMaterializationContext = new CteMaterializationContext();
// Stores the EXPLAIN/EXPLAIN ANALYZE results for Common Table Expressions (CTEs)
// Key: CTE table reference
@@ -276,7 +280,7 @@ public void prepareForRetry() {
}
private void cleanUpCte() {
- cteQueries.clear();
+ cteMaterializationContext.clear();
cteExplainResults.clear();
cteMaterializationCosts.clear();
subQueryTables.clear();
@@ -907,23 +911,37 @@ public Map, Long> getCteMaterializationCosts() {
}
public void addCteQuery(Table table, Query query) {
- cteQueries.put(NodeRef.of(table), query);
+ cteMaterializationContext.addCteQuery(table, query);
}
public Map, Query> getCteQueries() {
- return cteQueries;
+ return cteMaterializationContext.getCteQueries();
+ }
+
+ public boolean isCteMaterializationAttempted(Query query) {
+ return cteMaterializationContext.isMaterializationAttempted(query);
+ }
+
+ public void recordCteMaterializationResult(Query query, @Nullable CteDataStore dataStore) {
+ cteMaterializationContext.recordMaterializationResult(query, dataStore);
}
+ @Nullable
+ public CteDataStore getCteDataStore(Query query) {
+ return cteMaterializationContext.getCteDataStore(query);
+ }
+
+ @Nullable
public CteDataStore getCteDataStore(Table table) {
- Query query = cteQueries.get(NodeRef.of(table));
- if (query == null) {
- return null;
- }
- return query.getCteDataStore();
+ return cteMaterializationContext.getCteDataStore(table);
+ }
+
+ public CteMaterializationContext getCteMaterializationContext() {
+ return cteMaterializationContext;
}
- public void setCteQueries(Map, Query> cteQueries) {
- this.cteQueries = cteQueries;
+ public void setCteMaterializationContext(CteMaterializationContext cteMaterializationContext) {
+ this.cteMaterializationContext = cteMaterializationContext;
}
public void addSubQueryTables(Query query, List tables) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
index a7edb91c5cfbe..b98755377c41b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
@@ -38,8 +38,6 @@
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Literal;
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Parameter;
-import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query;
-import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Table;
import org.apache.iotdb.commons.queryengine.plan.relational.type.InternalTypeManager;
import org.apache.iotdb.commons.queryengine.plan.relational.type.TypeManager;
import org.apache.iotdb.db.auth.AuthorityChecker;
@@ -49,6 +47,7 @@
import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
import org.apache.iotdb.db.protocol.session.IClientSession;
import org.apache.iotdb.db.protocol.session.PreparedStatementInfo;
+import org.apache.iotdb.db.queryengine.common.CteMaterializationContext;
import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints;
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
import org.apache.iotdb.db.queryengine.common.MPPQueryContext.ExplainType;
@@ -464,7 +463,7 @@ public ExecutionResult executeForTableModel(
SessionInfo session,
String sql,
Metadata metadata,
- Map, Query> cteQueries,
+ CteMaterializationContext cteMaterializationContext,
ExplainType explainType,
ExplainOutputFormat explainOutputFormat,
long timeOut,
@@ -479,7 +478,10 @@ public ExecutionResult executeForTableModel(
debug,
((queryContext, startTime) -> {
queryContext.setInnerTriggeredQuery(true);
- queryContext.setCteQueries(cteQueries);
+ // Inner queries are part of the same top-level execution, so they must see its
+ // materialized CTE data while remaining isolated from other EXECUTE commands that may
+ // reuse the same prepared AST.
+ queryContext.setCteMaterializationContext(cteMaterializationContext);
queryContext.setExplainType(explainType);
queryContext.setExplainOutputFormat(explainOutputFormat);
queryContext.setVerbose(isVerbose);
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializer.java
index 322e0777b22b1..615beb267d062 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializer.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializer.java
@@ -83,17 +83,13 @@ public void materializeCTE(Analysis analysis, MPPQueryContext context) {
(tableRef, query) -> {
Table table = tableRef.getNode();
if (query.isMaterialized()) {
- if (!query.isExecuted()) {
+ if (!context.isCteMaterializationAttempted(query)) {
CteDataStore dataStore =
fetchCteQueryResult(context, table, query, analysis.getWith());
- query.setExecuted(true);
- if (dataStore == null) {
- // CTE query execution failed. Use inline instead of materialization
- // in the outer query
- query.setCteDataStore(null);
- return;
- }
- query.setCteDataStore(dataStore);
+ // Record a null result as well. A failed materialization falls back to inline
+ // planning and must not be attempted again for another reference in this
+ // execution.
+ context.recordCteMaterializationResult(query, dataStore);
}
context.addCteQuery(table, query);
}
@@ -112,11 +108,7 @@ public CteDataStore fetchCteQueryResult(
List tables = context.getTables(query);
List withQueries =
with.getQueries().stream()
- .filter(
- x ->
- tables.contains(x.getName())
- && !x.getQuery().isMaterialized()
- && !x.getQuery().isDone())
+ .filter(x -> tables.contains(x.getName()) && !x.getQuery().isMaterialized())
.collect(Collectors.toList());
if (!withQueries.isEmpty()) {
@@ -140,7 +132,7 @@ public CteDataStore fetchCteQueryResult(
sessionManager.getSessionInfoOfTableModel(sessionManager.getCurrSession()),
String.format("Materialize query for CTE '%s'", table.getName()),
LocalExecutionPlanner.getInstance().metadata,
- context.getCteQueries(),
+ context.getCteMaterializationContext(),
context.getExplainType(),
context.getExplainOutputFormat(),
context.getTimeOut(),
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
index 117d5db4e5c1d..33cfb7f034d95 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
@@ -287,7 +287,7 @@ private RelationPlan processNamedQuery(Table table, Query namedQuery, Scope scop
throw new SemanticException(DataNodeQueryMessages.UNEXPECTED_RECURSIVE_CTE);
}
- if (namedQuery.isMaterialized() && namedQuery.isDone()) {
+ if (namedQuery.isMaterialized()) {
RelationPlan materializedCtePlan = processMaterializedCte(table, namedQuery, scope);
if (materializedCtePlan != null) {
return materializedCtePlan;
@@ -298,7 +298,7 @@ private RelationPlan processNamedQuery(Table table, Query namedQuery, Scope scop
}
private RelationPlan processMaterializedCte(Table table, Query query, Scope scope) {
- CteDataStore dataStore = query.getCteDataStore();
+ CteDataStore dataStore = queryContext.getCteDataStore(query);
if (dataStore == null) {
return null;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java
index 7789de35f3a4e..e144b10151dff 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java
@@ -123,11 +123,7 @@ public Optional fetchUncorrelatedSubqueryResultForPredicate(
List tables = context.getTables(query);
List withQueries =
with.getQueries().stream()
- .filter(
- x ->
- tables.contains(x.getName())
- && !x.getQuery().isMaterialized()
- && !x.getQuery().isDone())
+ .filter(x -> tables.contains(x.getName()) && !x.getQuery().isMaterialized())
.collect(Collectors.toList());
if (!withQueries.isEmpty()) {
@@ -152,7 +148,7 @@ public Optional fetchUncorrelatedSubqueryResultForPredicate(
.getSessionInfoOfTableModel(SessionManager.getInstance().getCurrSession()),
"Try to Fetch Uncorrelated Scalar Subquery Result for Predicate",
LocalExecutionPlanner.getInstance().metadata,
- context.getCteQueries(),
+ context.getCteMaterializationContext(),
ExplainType.NONE,
ExplainOutputFormat.GRAPHVIZ,
context.getTimeOut(),
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerStateTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerStateTest.java
new file mode 100644
index 0000000000000..b68566e8dcad7
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerStateTest.java
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.relational.planner;
+
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QualifiedName;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QueryBody;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Table;
+import org.apache.iotdb.commons.queryengine.utils.cte.CteDataStore;
+import org.apache.iotdb.db.protocol.session.SessionManager;
+import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
+import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.core.classloader.annotations.SuppressStaticInitializationFor;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.Collections;
+import java.util.Optional;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.powermock.api.mockito.PowerMockito.mockStatic;
+import static org.powermock.api.mockito.PowerMockito.when;
+
+@PowerMockIgnore({
+ "com.sun.org.apache.xerces.*",
+ "javax.xml.*",
+ "org.xml.*",
+ "javax.management.*",
+ "javax.crypto.*",
+ "sun.security.*",
+ "java.time.*"
+})
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({Coordinator.class, SessionManager.class})
+@SuppressStaticInitializationFor("org.apache.iotdb.db.queryengine.plan.Coordinator")
+public class CteMaterializerStateTest {
+
+ @BeforeClass
+ public static void prepareEnvironment() {
+ mockStatic(Coordinator.class);
+ when(Coordinator.getInstance()).thenReturn(Mockito.mock(Coordinator.class));
+ mockStatic(SessionManager.class);
+ when(SessionManager.getInstance()).thenReturn(Mockito.mock(SessionManager.class));
+ }
+
+ @Test
+ public void testMaterializedCteIsRefetchedWhenSameStatementIsExecutedAgain() {
+ // PreparedStatementInfo caches the Query AST. The two contexts model two independent EXECUTE
+ // commands that must not share a materialized result through that cached AST.
+ final Query cteQuery = createMaterializedQuery();
+ final Table cteReference = new Table(QualifiedName.of("cte1"));
+
+ final Analysis firstAnalysis = new Analysis(cteQuery, Collections.emptyMap());
+ firstAnalysis.registerNamedQuery(cteReference, cteQuery);
+ final Analysis secondAnalysis = new Analysis(cteQuery, Collections.emptyMap());
+ secondAnalysis.registerNamedQuery(cteReference, cteQuery);
+
+ final CteDataStore firstResult = Mockito.mock(CteDataStore.class);
+ final CteDataStore secondResult = Mockito.mock(CteDataStore.class);
+ final CteMaterializer materializer = Mockito.spy(new CteMaterializer());
+ Mockito.doReturn(firstResult, secondResult)
+ .when(materializer)
+ .fetchCteQueryResult(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
+
+ final MPPQueryContext firstContext = new MPPQueryContext(new QueryId("first"));
+ materializer.materializeCTE(firstAnalysis, firstContext);
+ assertSame(firstResult, firstContext.getCteDataStore(cteQuery));
+
+ final MPPQueryContext secondContext = new MPPQueryContext(new QueryId("second"));
+ materializer.materializeCTE(secondAnalysis, secondContext);
+ assertSame(firstResult, firstContext.getCteDataStore(cteQuery));
+ assertSame(secondResult, secondContext.getCteDataStore(cteQuery));
+ Mockito.verify(materializer, Mockito.times(2))
+ .fetchCteQueryResult(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
+ }
+
+ @Test
+ public void testMaterializedCteIsFetchedOnceForMultipleReferencesInOneExecution() {
+ final Query cteQuery = createMaterializedQuery();
+ final Table firstReference = new Table(QualifiedName.of("cte1"));
+ final Table secondReference = new Table(QualifiedName.of("cte1"));
+ final Analysis analysis = new Analysis(cteQuery, Collections.emptyMap());
+ analysis.registerNamedQuery(firstReference, cteQuery);
+ analysis.registerNamedQuery(secondReference, cteQuery);
+
+ final CteDataStore result = Mockito.mock(CteDataStore.class);
+ final CteMaterializer materializer = Mockito.spy(new CteMaterializer());
+ Mockito.doReturn(result)
+ .when(materializer)
+ .fetchCteQueryResult(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
+
+ final MPPQueryContext context = new MPPQueryContext(new QueryId("same_execution"));
+ materializer.materializeCTE(analysis, context);
+ // Calling the planner hook again in the same execution must also reuse the recorded result.
+ materializer.materializeCTE(analysis, context);
+
+ assertSame(result, context.getCteDataStore(firstReference));
+ assertSame(result, context.getCteDataStore(secondReference));
+ Mockito.verify(materializer, Mockito.times(1))
+ .fetchCteQueryResult(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any());
+ }
+
+ @Test
+ public void testInnerQuerySharesOnlyItsParentExecutionState() {
+ final Query cteQuery = createMaterializedQuery();
+ final Table cteReference = new Table(QualifiedName.of("cte1"));
+ final CteDataStore result = Mockito.mock(CteDataStore.class);
+
+ final MPPQueryContext parentContext = new MPPQueryContext(new QueryId("parent"));
+ parentContext.recordCteMaterializationResult(cteQuery, result);
+ parentContext.addCteQuery(cteReference, cteQuery);
+
+ final MPPQueryContext innerContext = new MPPQueryContext(new QueryId("inner"));
+ innerContext.setCteMaterializationContext(parentContext.getCteMaterializationContext());
+ assertSame(result, innerContext.getCteDataStore(cteQuery));
+ assertSame(result, innerContext.getCteDataStore(cteReference));
+
+ final MPPQueryContext independentContext = new MPPQueryContext(new QueryId("independent"));
+ assertNull(independentContext.getCteDataStore(cteQuery));
+ assertNull(independentContext.getCteDataStore(cteReference));
+ }
+
+ @Test
+ public void testRetryClearsTopLevelExecutionState() {
+ final Query cteQuery = createMaterializedQuery();
+ final Table cteReference = new Table(QualifiedName.of("cte1"));
+ final MPPQueryContext context = new MPPQueryContext(new QueryId("retry"));
+ context.recordCteMaterializationResult(cteQuery, Mockito.mock(CteDataStore.class));
+ context.addCteQuery(cteReference, cteQuery);
+
+ context.prepareForRetry();
+
+ assertFalse(context.isCteMaterializationAttempted(cteQuery));
+ assertNull(context.getCteDataStore(cteQuery));
+ assertNull(context.getCteDataStore(cteReference));
+ }
+
+ private static Query createMaterializedQuery() {
+ final Query query =
+ new Query(
+ Optional.empty(),
+ Mockito.mock(QueryBody.class),
+ Optional.empty(),
+ Optional.empty(),
+ Optional.empty(),
+ Optional.empty());
+ query.setMaterialized(true);
+ return query;
+ }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java
index 21225e38e590f..6135ad98864b5 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java
@@ -113,7 +113,7 @@ public static void prepareEnv() {
Mockito.any(), // SessionInfo
Mockito.anyString(), // String
Mockito.any(), // Metadata
- Mockito.anyMap(), // Map, CteDataStore>
+ Mockito.any(), // CteMaterializationContext
Mockito.any(), // ExplainType
Mockito.any(), // ExplainOutputFormat
Mockito.anyLong(), // timeOut
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteSubqueryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteSubqueryTest.java
index ef1b289a892b8..c1ca0cdebcd6e 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteSubqueryTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteSubqueryTest.java
@@ -136,7 +136,7 @@ private void mockExecuteForTableModel() throws IoTDBException {
Mockito.any(), // SessionInfo
Mockito.anyString(), // String
Mockito.any(), // Metadata
- Mockito.anyMap(), // Map, CteDataStore>
+ Mockito.any(), // CteMaterializationContext
Mockito.any(), // ExplainType
Mockito.any(), // ExplainOutputFormat
Mockito.anyLong(), // timeOut
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/relational/sql/ast/Query.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/relational/sql/ast/Query.java
index 6c1d0395ad9db..5712b73ab2456 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/relational/sql/ast/Query.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/relational/sql/ast/Query.java
@@ -20,7 +20,6 @@
package org.apache.iotdb.commons.queryengine.plan.relational.sql.ast;
import org.apache.iotdb.commons.i18n.QueryMessages;
-import org.apache.iotdb.commons.queryengine.utils.cte.CteDataStore;
import com.google.common.collect.ImmutableList;
import org.apache.tsfile.utils.RamUsageEstimator;
@@ -45,10 +44,6 @@ public class Query extends Statement {
private final Optional limit;
// whether this query needs materialization
private boolean materialized = false;
- // whether this query has ever been executed
- private boolean isExecuted = false;
- // materialization has been executed successfully if cteDataStore is not null
- private CteDataStore cteDataStore = null;
public Query(
Optional with,
@@ -119,26 +114,6 @@ public void setMaterialized(boolean materialized) {
this.materialized = materialized;
}
- public boolean isExecuted() {
- return isExecuted;
- }
-
- public void setExecuted(boolean executed) {
- isExecuted = executed;
- }
-
- public boolean isDone() {
- return cteDataStore != null;
- }
-
- public void setCteDataStore(CteDataStore cteDataStore) {
- this.cteDataStore = cteDataStore;
- }
-
- public CteDataStore getCteDataStore() {
- return this.cteDataStore;
- }
-
@Override
public R accept(IAstVisitor visitor, C context) {
return ((CommonQueryAstVisitor) visitor).visitQuery(this, context);