CTE scan operator and CTE materialization memory control

Author: shizy818

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

@@ -426,6 +426,9 @@ public class IoTDBConfig {
   /** The buffer for sort operation */
   private long sortBufferSize = 32 * 1024 * 1024L;
 
+  /** The buffer for cte scan operation */
+  private long cteBufferSize = 32 * 1024 * 1024L;
+
   /** Mods cache size limit per fi */
   private long modsCacheSizeLimitPerFI = 32 * 1024 * 1024;
 
@@ -1370,6 +1373,7 @@ private void formulateFolders() {
     extPipeDir = addDataHomeDir(extPipeDir);
     queryDir = addDataHomeDir(queryDir);
     sortTmpDir = addDataHomeDir(sortTmpDir);
+    cteTmpDir = addDataHomeDir(cteTmpDir);
     formulateDataDirs(tierDataDirs);
   }
 
@@ -4076,6 +4080,14 @@ public long getSortBufferSize() {
     return sortBufferSize;
   }
 
+  public void setCteBufferSize(long cteBufferSize) {
+    this.cteBufferSize = cteBufferSize;
+  }
+
+  public long getCteBufferSize() {
+    return cteBufferSize;
+  }
+
   public void setModsCacheSizeLimitPerFI(long modsCacheSizeLimitPerFI) {
     this.modsCacheSizeLimitPerFI = modsCacheSizeLimitPerFI;
   }
@@ -4092,6 +4104,10 @@ public String getSortTmpDir() {
     return sortTmpDir;
   }
 
+  public void setCteTmpDir(String cteTmpDir) {
+    this.cteTmpDir = cteTmpDir;
+  }
+
   public String getCteTmpDir() {
     return cteTmpDir;
   }

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

@@ -1030,12 +1030,18 @@ public void loadProperties(TrimProperties properties) throws BadNodeUrlException
     // The buffer for sort operator to calculate
     loadFixedSizeLimitForQuery(properties, "sort_buffer_size_in_bytes", conf::setSortBufferSize);
 
+    // The buffer for cte scan operator
+    loadFixedSizeLimitForQuery(properties, "cte_buffer_size_in_bytes", conf::setCteBufferSize);
+
     loadFixedSizeLimitForQuery(
         properties, "mods_cache_size_limit_per_fi_in_bytes", conf::setModsCacheSizeLimitPerFI);
 
     // tmp filePath for sort operator
     conf.setSortTmpDir(properties.getProperty("sort_tmp_dir", conf.getSortTmpDir()));
 
+    // tmp filePath for cte
+    conf.setCteTmpDir(properties.getProperty("cte_tmp_dir", conf.getCteTmpDir()));
+
     conf.setRateLimiterType(properties.getProperty("rate_limiter_type", conf.getRateLimiterType()));
 
     conf.setDataNodeSchemaCacheEvictionPolicy(

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

@@ -96,6 +96,9 @@ public class MPPQueryContext {
 
   private boolean userQuery = false;
 
+  /** check if there is tmp file to be deleted. */
+  private boolean mayHaveTmpFile = false;
+
   public MPPQueryContext(QueryId queryId) {
     this.queryId = queryId;
     this.endPointBlackList = ConcurrentHashMap.newKeySet();
@@ -426,4 +429,12 @@ public boolean isQuery() {
   public void setUserQuery(boolean userQuery) {
     this.userQuery = userQuery;
   }
+
+  public void setMayHaveTmpFile(boolean mayHaveTmpFile) {
+    this.mayHaveTmpFile = mayHaveTmpFile;
+  }
+
+  public boolean mayHaveTmpFile() {
+    return mayHaveTmpFile;
+  }
 }

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

@@ -22,6 +22,7 @@
 package org.apache.iotdb.db.queryengine.execution.operator.source.relational;
 
 import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.db.queryengine.execution.MemoryEstimationHelper;
 import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
 import org.apache.iotdb.db.queryengine.execution.operator.source.AbstractSourceOperator;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
@@ -30,14 +31,17 @@
 import org.apache.iotdb.db.utils.cte.MemoryReader;
 
 import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.utils.RamUsageEstimator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.List;
 
 public class CteScanOperator extends AbstractSourceOperator {
-  private static final Logger logger = LoggerFactory.getLogger(CteScanOperator.class);
+  private static final Logger LOGGER = LoggerFactory.getLogger(CteScanOperator.class);
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(CteScanOperator.class);
 
   private final CteDataStore dataStore;
   private List<CteDataReader> dataReaders;
@@ -82,7 +86,7 @@ public void close() throws Exception {
         }
       }
     } catch (Exception e) {
-      logger.error("Fail to close fileChannel", e);
+      LOGGER.error("Fail to close fileChannel", e);
     }
   }
 
@@ -93,22 +97,27 @@ public boolean isFinished() throws Exception {
 
   @Override
   public long calculateMaxPeekMemory() {
-    return 0;
+    if (dataReaders == null || readerIndex >= dataReaders.size()) {
+      return 0;
+    }
+    return dataReaders.get(readerIndex).bytesUsed();
   }
 
   @Override
   public long calculateMaxReturnSize() {
-    return 0;
+    // The returned object is a reference to TsBlock in CteDataReader
+    return RamUsageEstimator.NUM_BYTES_OBJECT_REF;
   }
 
   @Override
   public long calculateRetainedSizeAfterCallingNext() {
-    return 0;
+    return calculateMaxPeekMemoryWithCounter();
   }
 
   @Override
   public long ramBytesUsed() {
-    return 0;
+    return INSTANCE_SIZE
+        + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext);
   }
 
   private void prepareReaders() throws IoTDBException {

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

@@ -24,6 +24,7 @@
 import org.apache.iotdb.commons.exception.IoTDBException;
 import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
 import org.apache.iotdb.commons.service.metric.PerformanceOverviewMetrics;
+import org.apache.iotdb.commons.utils.FileUtils;
 import org.apache.iotdb.db.conf.IoTDBConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.exception.query.KilledByOthersException;
@@ -47,6 +48,8 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil;
+import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.ir.CteMaterializer;
 import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler;
 import org.apache.iotdb.db.utils.SetThreadName;
 import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
@@ -59,6 +62,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.File;
 import java.nio.ByteBuffer;
 import java.time.format.DateTimeParseException;
 import java.util.List;
@@ -349,6 +353,9 @@ private void releaseResource() {
       resultHandle.close();
       cleanUpResultHandle();
     }
+    if (getSQLDialect().equals(IClientSession.SqlDialect.TABLE)) {
+      CteMaterializer.cleanUpCTE((Analysis) analysis, context);
+    }
   }
 
   private void cleanUpResultHandle() {
@@ -373,6 +380,15 @@ private void cleanUpResultHandle() {
   public void stopAndCleanup(Throwable t) {
     stop(t);
     releaseResource(t);
+
+    try {
+      // delete cte tmp file if exists
+      deleteTmpFile();
+    } catch (Throwable err) {
+      LOGGER.error(
+          "Errors occurred while attempting to delete tmp files, potentially leading to resource leakage.",
+          err);
+    }
   }
 
   /** Release the resources that current QueryExecution hold with a specified exception */
@@ -392,6 +408,9 @@ private void releaseResource(Throwable t) {
       }
       cleanUpResultHandle();
     }
+    if (getSQLDialect().equals(IClientSession.SqlDialect.TABLE)) {
+      CteMaterializer.cleanUpCTE((Analysis) analysis, context);
+    }
   }
 
   /**
@@ -712,4 +731,18 @@ public String toString() {
   public ScheduledExecutorService getScheduledExecutor() {
     return planner.getScheduledExecutorService();
   }
+
+  private void deleteTmpFile() {
+    if (context.mayHaveTmpFile()) {
+      String tmpFilePath =
+          IoTDBDescriptor.getInstance().getConfig().getCteTmpDir()
+              + File.separator
+              + context.getQueryId()
+              + File.separator;
+      File tmpFile = new File(tmpFilePath);
+      if (tmpFile.exists()) {
+        FileUtils.deleteFileOrDirectory(tmpFile, true);
+      }
+    }
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanNodeType.java

@@ -117,6 +117,7 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTreeDeviceViewScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AssignUniqueId;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.node.CteScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.EnforceSingleRowNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.GapFillNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.GroupNode;
@@ -310,6 +311,7 @@ public enum PlanNodeType {
   TABLE_WINDOW_FUNCTION((short) 1032),
   TABLE_INTO_NODE((short) 1033),
   TABLE_UNION_NODE((short) 1034),
+  TABLE_CTE_SCAN_NODE((short) 1035),
 
   RELATIONAL_INSERT_TABLET((short) 2000),
   RELATIONAL_INSERT_ROW((short) 2001),
@@ -697,6 +699,8 @@ public static PlanNode deserialize(ByteBuffer buffer, short nodeType) {
             buffer);
       case 1034:
         return UnionNode.deserialize(buffer);
+      case 1035:
+        return CteScanNode.deserialize(buffer);
       case 2000:
         return RelationalInsertTabletNode.deserialize(buffer);
       case 2001:

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

@@ -122,7 +122,7 @@ public class Analysis implements IAnalysis {
 
   private String updateType;
 
-  private final Map<Query, CteDataStore> namedQueriesDataStore = new LinkedHashMap<>();
+  private final Map<String, CteDataStore> cteDataStores = new LinkedHashMap<>();
 
   private final Map<NodeRef<Table>, Query> namedQueries = new LinkedHashMap<>();
 
@@ -285,12 +285,16 @@ public Map<NodeRef<Table>, Query> getNamedQueries() {
     return namedQueries;
   }
 
-  public void addCteDataStore(Query query, CteDataStore dataStore) {
-    namedQueriesDataStore.put(query, dataStore);
+  public void addCteDataStore(String cteName, CteDataStore dataStore) {
+    cteDataStores.put(cteName, dataStore);
   }
 
-  public CteDataStore getCTEDataStore(Query query) {
-    return namedQueriesDataStore.get(query);
+  public Map<String, CteDataStore> getCteDataStores() {
+    return cteDataStores;
+  }
+
+  public CteDataStore getCteDataStore(String cteName) {
+    return cteDataStores.get(cteName);
   }
 
   public boolean isAnalyzed(Expression expression) {

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

@@ -248,7 +248,7 @@ protected RelationPlan visitTable(final Table table, final Void context) {
       }
 
       if (namedQuery.isMaterialized()) {
-        CteDataStore dataStore = analysis.getCTEDataStore(namedQuery);
+        CteDataStore dataStore = analysis.getCteDataStore(table.getName().toString());
         if (dataStore != null) {
           List<Symbol> outputSymbols =
               analysis.getOutputDescriptor(table).getAllFields().stream()

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

@@ -40,11 +40,11 @@
 
 import java.io.File;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
 public class CteMaterializer {
-  private static final SqlParser relationSqlParser = new SqlParser();
 
   private static final Coordinator coordinator = Coordinator.getInstance();
 
@@ -57,16 +57,33 @@ public static void materializeCTE(Analysis analysis, MPPQueryContext context) {
         .forEach(
             (tableRef, query) -> {
               if (query.isMaterialized() && !materializedQueries.contains(query)) {
-                CteDataStore dataStore =
-                    fetchCteQueryResult(tableRef.getNode().getName().toString(), query, context);
+                if (!context.mayHaveTmpFile()) {
+                  context.setMayHaveTmpFile(true);
+                }
+
+                String cteName = tableRef.getNode().getName().toString();
+                CteDataStore dataStore = fetchCteQueryResult(cteName, query, context);
                 if (dataStore != null) {
-                  analysis.addCteDataStore(query, dataStore);
+                  analysis.addCteDataStore(cteName, dataStore);
+                  context.reserveMemoryForFrontEnd(dataStore.getCachedBytes());
                   materializedQueries.add(query);
                 }
               }
             });
   }
 
+  public static void cleanUpCTE(Analysis analysis, MPPQueryContext context) {
+    Map<String, CteDataStore> cteDataStores = analysis.getCteDataStores();
+    cteDataStores
+        .values()
+        .forEach(
+            dataStore -> {
+              context.releaseMemoryReservedForFrontEnd(dataStore.getCachedBytes());
+              dataStore.clear();
+            });
+    cteDataStores.clear();
+  }
+
   public static CteDataStore fetchCteQueryResult(
       String cteName, Query query, MPPQueryContext context) {
     final long queryId = SessionManager.getInstance().requestQueryId();
@@ -91,7 +108,7 @@ public static CteDataStore fetchCteQueryResult(
       String folderPath =
           IoTDBDescriptor.getInstance().getConfig().getCteTmpDir()
               + File.separator
-              + queryId
+              + context.getQueryId()
               + File.separator;
       CteDataStore cteDataStore =
           new CteDataStore(new DiskSpiller(folderPath, folderPath + cteName));

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

@@ -119,4 +119,8 @@ protected void serializeAttributes(ByteBuffer byteBuffer) {}
 
   @Override
   protected void serializeAttributes(DataOutputStream stream) throws IOException {}
+
+  public static CteScanNode deserialize(ByteBuffer byteBuffer) {
+    return null;
+  }
 }