cteDataStores -> cteQueries

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java

@@ -117,7 +117,7 @@ public class MPPQueryContext implements IAuditEntity {
   private boolean userQuery = false;
 
   private final Map<NodeRef<Table>, Long> cteMaterializationCosts = new HashMap<>();
-  private Map<NodeRef<Table>, CteDataStore> cteDataStores = new HashMap<>();
+  private Map<NodeRef<Table>, Query> cteQueries = new HashMap<>();
   // table -> (maxLineLength, 'explain' or 'explain analyze' result)
   // Max line length of each CTE should be remembered because we need to standardize
   // the output format of main query and CTE query.
@@ -493,20 +493,24 @@ public Map<NodeRef<Table>, Long> getCteMaterializationCosts() {
     return cteMaterializationCosts;
   }
 
-  public void addCteDataStore(Table table, CteDataStore dataStore) {
-    cteDataStores.put(NodeRef.of(table), dataStore);
+  public void addCteQuery(Table table, Query query) {
+    cteQueries.put(NodeRef.of(table), query);
   }
 
-  public Map<NodeRef<Table>, CteDataStore> getCteDataStores() {
-    return cteDataStores;
+  public Map<NodeRef<Table>, Query> getCteQueries() {
+    return cteQueries;
   }
 
   public CteDataStore getCteDataStore(Table table) {
-    return cteDataStores.get(NodeRef.of(table));
+    Query query = cteQueries.get(NodeRef.of(table));
+    if (query == null) {
+      return null;
+    }
+    return query.getCteDataStore();
   }
 
-  public void setCteDataStores(Map<NodeRef<Table>, CteDataStore> cteDataStores) {
-    this.cteDataStores = cteDataStores;
+  public void setCteQueries(Map<NodeRef<Table>, Query> cteQueries) {
+    this.cteQueries = cteQueries;
   }
 
   public Map<Query, List<Identifier>> getSubQueryTables() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java

@@ -93,6 +93,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.MigrateRegion;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Parameter;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.PipeStatement;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Query;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Prepare;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ReconstructRegion;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.RelationalAuthorStatement;
@@ -143,7 +144,6 @@
 import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement;
 import org.apache.iotdb.db.queryengine.plan.statement.Statement;
 import org.apache.iotdb.db.utils.SetThreadName;
-import org.apache.iotdb.db.utils.cte.CteDataStore;
 
 import org.apache.thrift.TBase;
 import org.apache.tsfile.utils.Accountable;
@@ -418,7 +418,7 @@ public ExecutionResult executeForTableModel(
       SessionInfo session,
       String sql,
       Metadata metadata,
-      Map<NodeRef<Table>, CteDataStore> cteDataStoreMap,
+      Map<NodeRef<Table>, Query> cteQueries,
       ExplainType explainType,
       long timeOut,
       boolean userQuery) {
@@ -429,7 +429,7 @@ public ExecutionResult executeForTableModel(
         userQuery,
         ((queryContext, startTime) -> {
           queryContext.setSubquery(true);
-          queryContext.setCteDataStores(cteDataStoreMap);
+          queryContext.setCteQueries(cteQueries);
           queryContext.setExplainType(explainType);
           return createQueryExecutionForTableModel(
               statement,

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java

@@ -82,10 +82,8 @@ public Analysis analyze(Statement statement) {
     Analysis analysis = new Analysis(rewrittenStatement, parameterLookup);
     // Register CTE passed by parent query.
     context
-        .getCteDataStores()
-        .forEach(
-            (tableRef, dataStore) ->
-                analysis.registerNamedQuery(tableRef.getNode(), dataStore.getQuery()));
+        .getCteQueries()
+        .forEach((tableRef, query) -> analysis.registerNamedQuery(tableRef.getNode(), query));
 
     Statement innerStatement =
         rewrittenStatement instanceof PipeEnriched

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

@@ -82,26 +82,31 @@ public void materializeCTE(Analysis analysis, MPPQueryContext context) {
         .forEach(
             (tableRef, query) -> {
               Table table = tableRef.getNode();
-              if (query.isMaterialized() && !query.isDone()) {
-                CteDataStore dataStore =
-                    fetchCteQueryResult(context, table, query, analysis.getWith());
-                if (dataStore == null) {
-                  // CTE query execution failed. Use inline instead of materialization
-                  // in the outer query
-                  query.setDone(false);
-                  return;
+              if (query.isMaterialized()) {
+                if (!query.isDone()) {
+                  CteDataStore dataStore =
+                      fetchCteQueryResult(context, table, query, analysis.getWith());
+                  if (dataStore == null) {
+                    // CTE query execution failed. Use inline instead of materialization
+                    // in the outer query
+                    query.setCteDataStore(null);
+                    return;
+                  }
+                  query.setCteDataStore(dataStore);
                 }
-
-                context.addCteDataStore(table, dataStore);
-                query.setDone(true);
+                context.addCteQuery(table, query);
               }
             });
   }
 
   public void cleanUpCTE(MPPQueryContext context) {
-    Map<NodeRef<Table>, CteDataStore> cteDataStores = context.getCteDataStores();
-    cteDataStores.values().forEach(CteDataStore::clear);
-    cteDataStores.clear();
+    Map<NodeRef<Table>, Query> cteQueries = context.getCteQueries();
+    cteQueries.values().stream()
+        .map(Query::getCteDataStore)
+        .filter(java.util.Objects::nonNull)
+        .distinct()
+        .forEach(CteDataStore::clear);
+    cteQueries.clear();
   }
 
   public CteDataStore fetchCteQueryResult(
@@ -145,7 +150,7 @@ public CteDataStore fetchCteQueryResult(
               sessionManager.getSessionInfoOfTableModel(sessionManager.getCurrSession()),
               String.format("Materialize query for CTE '%s'", table.getName()),
               LocalExecutionPlanner.getInstance().metadata,
-              context.getCteDataStores(),
+              context.getCteQueries(),
               context.getExplainType(),
               context.getTimeOut(),
               false);
@@ -160,8 +165,7 @@ public CteDataStore fetchCteQueryResult(
       TableSchema tableSchema = getTableSchema(datasetHeader, table.getName().toString());
 
       cteDataStore =
-          new CteDataStore(
-              query, tableSchema, datasetHeader.getColumnIndex2TsBlockColumnIndexList());
+          new CteDataStore(tableSchema, datasetHeader.getColumnIndex2TsBlockColumnIndexList());
       while (execution.hasNextResult()) {
         final Optional<TsBlock> tsBlock;
         try {

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

@@ -146,7 +146,7 @@ public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
                   .getSessionInfoOfTableModel(SessionManager.getInstance().getCurrSession()),
               "Try to Fetch Uncorrelated Scalar Subquery Result for Predicate",
               LocalExecutionPlanner.getInstance().metadata,
-              context.getCteDataStores(),
+              context.getCteQueries(),
               ExplainType.NONE,
               context.getTimeOut(),
               false);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/Query.java

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
+import org.apache.iotdb.db.utils.cte.CteDataStore;
+
 import com.google.common.collect.ImmutableList;
 
 import java.util.List;
@@ -39,8 +41,7 @@ public class Query extends Statement {
   private final Optional<Node> limit;
   // whether this query needs materialization
   private boolean materialized = false;
-  // whether materialization is done
-  private boolean done = false;
+  private CteDataStore cteDataStore = null;
 
   public Query(
       Optional<With> with,
@@ -112,11 +113,15 @@ public void setMaterialized(boolean materialized) {
   }
 
   public boolean isDone() {
-    return done;
+    return cteDataStore != null;
+  }
+
+  public void setCteDataStore(CteDataStore cteDataStore) {
+    this.cteDataStore = cteDataStore;
   }
 
-  public void setDone(boolean done) {
-    this.done = done;
+  public CteDataStore getCteDataStore() {
+    return this.cteDataStore;
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/cte/CteDataStore.java

@@ -24,7 +24,6 @@
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.TableSchema;
-import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Query;
 
 import org.apache.tsfile.read.common.block.TsBlock;
 
@@ -33,7 +32,6 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class CteDataStore {
-  private final Query query;
   private final TableSchema tableSchema;
   private final List<Integer> columnIndex2TsBlockColumnIndexList;
 
@@ -44,9 +42,7 @@ public class CteDataStore {
   // reference count by CteScanOperator
   private final AtomicInteger count;
 
-  public CteDataStore(
-      Query query, TableSchema tableSchema, List<Integer> columnIndex2TsBlockColumnIndexList) {
-    this.query = query;
+  public CteDataStore(TableSchema tableSchema, List<Integer> columnIndex2TsBlockColumnIndexList) {
     this.tableSchema = tableSchema;
     this.columnIndex2TsBlockColumnIndexList = columnIndex2TsBlockColumnIndexList;
     this.cachedData = new ArrayList<>();
@@ -87,10 +83,6 @@ public TableSchema getTableSchema() {
     return tableSchema;
   }
 
-  public Query getQuery() {
-    return query;
-  }
-
   public List<Integer> getColumnIndex2TsBlockColumnIndexList() {
     return columnIndex2TsBlockColumnIndexList;
   }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/CteScanOperatorTest.java

@@ -26,7 +26,6 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.ColumnSchema;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.TableSchema;
-import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Query;
 import org.apache.iotdb.db.utils.cte.CteDataStore;
 
 import com.google.common.collect.ImmutableList;
@@ -72,14 +71,11 @@ public void setUp() {
     // Create a simple table schema for testing
     TableSchema tableSchema = createTestTableSchema();
 
-    // Create mock query
-    Query mockQuery = mock(Query.class);
-
     // Create column index mapping
     List<Integer> columnIndex2TsBlockColumnIndexList = Arrays.asList(0, 1, 2);
 
     // Initialize CteDataStore
-    cteDataStore = new CteDataStore(mockQuery, tableSchema, columnIndex2TsBlockColumnIndexList);
+    cteDataStore = new CteDataStore(tableSchema, columnIndex2TsBlockColumnIndexList);
 
     // Add test data to the data store
     List<TsBlock> testData = createTestTsBlocks();
@@ -105,9 +101,8 @@ public void testConstructor() throws Exception {
   @Test
   public void testEmptyDataStore() throws Exception {
     // Create empty data store
-    Query mockQuery = mock(Query.class);
     TableSchema tableSchema = createTestTableSchema();
-    CteDataStore emptyDataStore = new CteDataStore(mockQuery, tableSchema, Arrays.asList(0, 1, 2));
+    CteDataStore emptyDataStore = new CteDataStore(tableSchema, Arrays.asList(0, 1, 2));
 
     cteScanOperator = new CteScanOperator(operatorContext, planNodeId, emptyDataStore);
     // Should not have data