code review

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -2199,6 +2199,24 @@ public synchronized void loadHotModifiedProps(TrimProperties properties)
           Long.parseLong(
               properties.getProperty(
                   "max_object_file_size_in_byte", String.valueOf(conf.getMaxObjectSizeInByte()))));
+
+      // The buffer for cte materialization.
+      long cteBufferSizeInBytes =
+          Long.parseLong(
+              properties.getProperty(
+                  "cte_buffer_size_in_bytes", Long.toString(conf.getCteBufferSize())));
+      if (cteBufferSizeInBytes > 0) {
+        conf.setCteBufferSize(cteBufferSizeInBytes);
+      }
+      // max number of rows for cte materialization
+      int maxRowsInCteBuffer =
+          Integer.parseInt(
+              properties.getProperty(
+                  "max_rows_in_cte_buffer", Integer.toString(conf.getMaxRowsInCteBuffer())));
+      if (maxRowsInCteBuffer > 0) {
+        conf.setMaxRowsInCteBuffer(maxRowsInCteBuffer);
+      }
+
     } catch (Exception e) {
       if (e instanceof InterruptedException) {
         Thread.currentThread().interrupt();

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

@@ -137,8 +137,10 @@ public enum ExplainType {
   // planning
   private final Map<NodeRef<Table>, Long> cteMaterializationCosts = new HashMap<>();
 
-  // Never materialize CTE in a subquery.
-  private boolean subquery = false;
+  // Indicates whether this query context is for a sub-query triggered by the main query.
+  // Sub-queries are independent queries spawned from the main query (e.g., CTE sub-queries).
+  // When true, CTE materialization is skipped as it's handled by the main query context.
+  private boolean innerTriggeredQuery = false;
 
   // Tables in the subquery
   private final Map<NodeRef<Query>, List<Identifier>> subQueryTables = new HashMap<>();
@@ -210,7 +212,7 @@ public void releaseMemoryForSchemaTree() {
   }
 
   public void prepareForRetry() {
-    if (!isSubquery()) {
+    if (!isInnerTriggeredQuery()) {
       cleanUpCte();
     }
     this.initResultNodeContext();
@@ -223,6 +225,7 @@ private void cleanUpCte() {
     cteMaterializationCosts.clear();
     subQueryTables.clear();
   }
+
   private void initResultNodeContext() {
     this.resultNodeContext = new ResultNodeContext(queryId);
   }
@@ -500,12 +503,12 @@ public void setUserQuery(boolean userQuery) {
     this.userQuery = userQuery;
   }
 
-  public boolean isSubquery() {
-    return subquery;
+  public boolean isInnerTriggeredQuery() {
+    return innerTriggeredQuery;
   }
 
-  public void setSubquery(boolean subquery) {
-    this.subquery = subquery;
+  public void setInnerTriggeredQuery(boolean innerTriggeredQuery) {
+    this.innerTriggeredQuery = innerTriggeredQuery;
   }
 
   public void addCteMaterializationCost(Table table, long cost) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java

@@ -59,6 +59,11 @@ public class FragmentInstanceExecution {
   // It will be set to null while this FI is FINISHED
   private List<IDriver> drivers;
 
+  // Indicates whether this fragment instance should be ignored for statistics collection.
+  // This is true when the fragment instance contains ExplainAnalyzeOperator, which is
+  // a virtual fragment used for EXPLAIN ANALYZE and should not be included in query statistics.
+  boolean shouldIgnoreForStatistics;
+
   // It will be set to null while this FI is FINISHED
   private ISink sink;
 
@@ -110,6 +115,7 @@ private FragmentInstanceExecution(
     this.stateMachine = stateMachine;
     this.timeoutInMs = timeoutInMs;
     this.exchangeManager = exchangeManager;
+    this.shouldIgnoreForStatistics = shouldIgnoreForStatistics();
   }
 
   public FragmentInstanceState getInstanceState() {
@@ -141,7 +147,7 @@ public FragmentInstanceStateMachine getStateMachine() {
   }
 
   // Check if this fragment instance should be ignored for statistics
-  // (i.e., it contains ExplainAnalyzeOperator only)
+  // (i.e., it contains ExplainAnalyzeOperator)
   private boolean shouldIgnoreForStatistics() {
     if (drivers == null || drivers.isEmpty()) {
       return false;
@@ -166,7 +172,7 @@ private boolean fillFragmentInstanceStatistics(
     statistics.setState(getInstanceState().toString());
     // Previously we ignore statistics when current data region is instance of
     // VirtualDataRegion. Now data region of a CteScanNode is also virtual.
-    if (shouldIgnoreForStatistics()) {
+    if (shouldIgnoreForStatistics) {
       // We don't need to output the region having ExplainAnalyzeOperator only.
       return false;
     }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/CteScanOperator.java

@@ -30,6 +30,7 @@
 import org.apache.iotdb.db.utils.cte.CteDataStore;
 import org.apache.iotdb.db.utils.cte.MemoryReader;
 
+import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.read.common.block.TsBlock;
 import org.apache.tsfile.utils.RamUsageEstimator;
 import org.slf4j.Logger;
@@ -40,6 +41,9 @@ public class CteScanOperator implements SourceOperator {
   private static final long INSTANCE_SIZE =
       RamUsageEstimator.shallowSizeOfInstance(CteScanOperator.class);
 
+  private final long maxReturnSize =
+      TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+
   private final OperatorContext operatorContext;
   private final PlanNodeId sourceId;
 
@@ -79,7 +83,7 @@ public void close() throws Exception {
         dataReader.close();
       }
     } catch (Exception e) {
-      LOGGER.error("Fail to close fileChannel", e);
+      LOGGER.error("Fail to close CteDataReader", e);
     }
   }
 
@@ -95,8 +99,7 @@ public long calculateMaxPeekMemory() {
 
   @Override
   public long calculateMaxReturnSize() {
-    // The returned object is a reference to TsBlock in CteDataReader
-    return RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+    return maxReturnSize;
   }
 
   @Override

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

@@ -427,7 +427,7 @@ public ExecutionResult executeForTableModel(
         sql,
         userQuery,
         ((queryContext, startTime) -> {
-          queryContext.setSubquery(true);
+          queryContext.setInnerTriggeredQuery(true);
           queryContext.setCteQueries(cteQueries);
           queryContext.setExplainType(explainType);
           return createQueryExecutionForTableModel(

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

@@ -70,6 +70,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.TableFunctionInvocation;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.WindowFrame;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.With;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
 import org.apache.iotdb.db.queryengine.plan.statement.component.FillPolicy;
 
 import com.google.common.collect.ArrayListMultimap;
@@ -124,6 +125,8 @@ public class Analysis implements IAnalysis {
 
   private final Map<NodeRef<Table>, Query> namedQueries = new LinkedHashMap<>();
 
+  // WITH clause stored during analyze phase. Required for constant folding and CTE materialization
+  // subqueries, which cannot directly access the WITH clause
   private With with;
 
   // map expandable query to the node being the inner recursive reference
@@ -255,6 +258,11 @@ public class Analysis implements IAnalysis {
 
   private boolean isQuery = false;
 
+  // SqlParser is needed during query planning phase for executing uncorrelated scalar subqueries
+  // in advance (predicate folding). The planner needs to parse and execute these subqueries
+  // independently to utilize predicate pushdown optimization.
+  private SqlParser sqlParser;
+
   public Analysis(@Nullable Statement root, Map<NodeRef<Parameter>, Expression> parameters) {
     this.root = root;
     this.parameters = ImmutableMap.copyOf(requireNonNull(parameters, "parameters is null"));
@@ -277,6 +285,14 @@ public void setUpdateType(String updateType) {
     this.updateType = updateType;
   }
 
+  public SqlParser getSqlParser() {
+    return sqlParser;
+  }
+
+  public void setSqlParser(SqlParser sqlParser) {
+    this.sqlParser = sqlParser;
+  }
+
   public Query getNamedQuery(Table table) {
     return namedQueries.get(NodeRef.of(table));
   }

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

@@ -80,6 +80,7 @@ public Analysis analyze(Statement statement) {
             warningCollector);
 
     Analysis analysis = new Analysis(rewrittenStatement, parameterLookup);
+    analysis.setSqlParser(statementAnalyzerFactory.getSqlParser());
     // Register CTE passed by parent query.
     context
         .getCteQueries()

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

@@ -54,7 +54,11 @@ public class Scope {
   private final RelationId relationId;
   private final RelationType relation;
   private final Map<String, WithQuery> namedQueries;
-  private final List<Identifier> tables;
+
+  // Tables to access for the current relation. For CTE materialization and constant folding
+  // subqueries, non-materialized CTEs in tables must be identified, and their definitions
+  // attached to the subquery context.
+  private List<Identifier> tables;
 
   public static Scope create() {
     return builder().build();
@@ -83,12 +87,8 @@ public void addTable(Table table) {
     tables.add(new Identifier(table.getName().getSuffix()));
   }
 
-  public void addTables(List<Identifier> tables) {
-    this.tables.addAll(tables);
-  }
-
-  public Scope copy() {
-    return builder().like(this).build();
+  public void setTables(List<Identifier> tables) {
+    this.tables = tables;
   }
 
   public List<Identifier> getTables() {

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

@@ -3630,15 +3630,13 @@ protected Scope visitJoin(Join node, Optional<Scope> scope) {
 
       joinConditionCheck(criteria);
 
-      Optional<Scope> leftScope = scope.map(Scope::copy);
-      Scope left = process(node.getLeft(), leftScope);
-      Optional<Scope> rightScope = scope.map(Scope::copy);
-      Scope right = process(node.getRight(), rightScope);
+      // remember current tables in the scope
+      List<Identifier> tables = new ArrayList<>();
+      scope.ifPresent(s -> tables.addAll(s.getTables()));
 
-      if (scope.isPresent()) {
-        leftScope.ifPresent(l -> scope.get().addTables(l.getTables()));
-        rightScope.ifPresent(l -> scope.get().addTables(l.getTables()));
-      }
+      Scope left = process(node.getLeft(), scope);
+      scope.ifPresent(s -> s.setTables(tables));
+      Scope right = process(node.getRight(), scope);
 
       if (criteria instanceof JoinUsing) {
         return analyzeJoinUsing(node, ((JoinUsing) criteria).getColumns(), scope, left, right);

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

@@ -77,4 +77,8 @@ public static StatementAnalyzerFactory createTestingStatementAnalyzerFactory(
   public AccessControl getAccessControl() {
     return accessControl;
   }
+
+  public SqlParser getSqlParser() {
+    return sqlParser;
+  }
 }