Fix CTE materialization state isolation

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"};

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.
+ *
+ * <p>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<NodeRef<Table>, Query> cteQueries = new HashMap<>();
+
+  // Optional.empty() records an attempted materialization that fell back to inline planning.
+  private final Map<NodeRef<Query>, Optional<CteDataStore>> materializationResults =
+      new HashMap<>();
+
+  public void addCteQuery(Table table, Query query) {
+    cteQueries.put(NodeRef.of(table), query);
+  }
+
+  public Map<NodeRef<Table>, 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();
+  }
+}

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<NodeRef<Table>, 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<NodeRef<Table>, Long> getCteMaterializationCosts() {
   }
 
   public void addCteQuery(Table table, Query query) {
-    cteQueries.put(NodeRef.of(table), query);
+    cteMaterializationContext.addCteQuery(table, query);
   }
 
   public Map<NodeRef<Table>, 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<NodeRef<Table>, Query> cteQueries) {
-    this.cteQueries = cteQueries;
+  public void setCteMaterializationContext(CteMaterializationContext cteMaterializationContext) {
+    this.cteMaterializationContext = cteMaterializationContext;
   }
 
   public void addSubQueryTables(Query query, List<Identifier> tables) {

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<NodeRef<Table>, 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);

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<Identifier> tables = context.getTables(query);
         List<WithQuery> 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(),

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

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java

@@ -123,11 +123,7 @@ public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
         List<Identifier> tables = context.getTables(query);
         List<WithQuery> 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<Literal> 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(),