feat: add AstMemoryEstimator for SQL AST memory estimation (#16882)

Author: Young-Leo

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/session/PrepareTask.java

@@ -26,7 +26,6 @@
 import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
 import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
 import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
-import org.apache.iotdb.db.queryengine.plan.relational.sql.AstMemoryEstimator;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement;
 import org.apache.iotdb.rpc.TSStatusCode;
 
@@ -69,7 +68,7 @@ public ListenableFuture<ConfigTaskResult> execute(IConfigTaskExecutor configTask
     }
 
     // Estimate memory size of the AST
-    long memorySizeInBytes = AstMemoryEstimator.estimateMemorySize(sql);
+    long memorySizeInBytes = sql == null ? 0L : sql.ramBytesUsed();
 
     // Allocate memory from CoordinatorMemoryManager
     // This memory is shared across all sessions using a single MemoryBlock

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

@@ -31,6 +31,7 @@
 import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache;
 
 import org.apache.tsfile.read.common.block.TsBlock;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -42,6 +43,9 @@
 
 public abstract class AbstractQueryDeviceWithCache extends AbstractTraverseDevice {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AbstractQueryDeviceWithCache.class);
+
   // For query devices fully in cache
   protected List<ShowDevicesResult> results = new ArrayList<>();
 
@@ -111,4 +115,21 @@ public static List<ColumnHeader> getDeviceColumnHeaderList(
   public abstract DatasetHeader getDataSetHeader();
 
   public abstract TsBlock getTsBlock(final Analysis analysis);
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE + ramBytesUsedForCommonFields();
+  }
+
+  @Override
+  protected long ramBytesUsedForCommonFields() {
+    long size = super.ramBytesUsedForCommonFields();
+    if (results != null) {
+      size += RamUsageEstimator.shallowSizeOf(results);
+      for (ShowDevicesResult result : results) {
+        size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(result);
+      }
+    }
+    return size;
+  }
 }

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

@@ -32,6 +32,7 @@
 
 import com.google.common.collect.ImmutableList;
 import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import java.util.Collections;
 import java.util.List;
@@ -46,6 +47,9 @@
 // Show, Count, Update, Delete Devices
 public abstract class AbstractTraverseDevice extends Statement {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AbstractTraverseDevice.class);
+
   protected String database;
 
   protected String tableName;
@@ -187,6 +191,10 @@ public void setAttributeColumns(final List<String> attributeColumns) {
     this.attributeColumns = attributeColumns;
   }
 
+  public List<String> getAttributeColumns() {
+    return attributeColumns;
+  }
+
   public List<ColumnHeader> getColumnHeaderList() {
     return columnHeaderList;
   }
@@ -241,4 +249,44 @@ protected String toStringContent() {
         + tagFuzzyPredicate
         + '}';
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE + ramBytesUsedForCommonFields();
+  }
+
+  protected long ramBytesUsedForCommonFields() {
+    long size = 0;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(table);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(where);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(tagFuzzyPredicate);
+    size += RamUsageEstimator.sizeOf(database);
+    size += RamUsageEstimator.sizeOf(tableName);
+    if (tagDeterminedFilterList != null) {
+      size += RamUsageEstimator.shallowSizeOf(tagDeterminedFilterList);
+      for (List<SchemaFilter> filters : tagDeterminedFilterList) {
+        if (filters != null) {
+          size += RamUsageEstimator.shallowSizeOf(filters);
+          for (SchemaFilter filter : filters) {
+            size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(filter);
+          }
+        }
+      }
+    }
+    if (columnHeaderList != null) {
+      size += RamUsageEstimator.shallowSizeOf(columnHeaderList);
+      for (ColumnHeader header : columnHeaderList) {
+        size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(header);
+      }
+    }
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfStringList(attributeColumns);
+    if (partitionKeyList != null) {
+      size += RamUsageEstimator.shallowSizeOf(partitionKeyList);
+      for (IDeviceID deviceID : partitionKeyList) {
+        size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(deviceID);
+      }
+    }
+    return size;
+  }
 }

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

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
+import org.apache.tsfile.utils.RamUsageEstimator;
+
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -27,6 +29,8 @@
 import static java.util.Objects.requireNonNull;
 
 public class AddColumn extends Statement {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AddColumn.class);
 
   private final QualifiedName tableName;
   private final ColumnDefinition column;
@@ -112,4 +116,13 @@ public String toString() {
         .add("view", view)
         .toString();
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += tableName == null ? 0L : tableName.ramBytesUsed();
+    size += column == null ? 0L : column.ramBytesUsed();
+    return size;
+  }
 }

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

@@ -20,6 +20,7 @@
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import java.util.List;
 import java.util.Objects;
@@ -28,6 +29,9 @@
 import static java.util.Objects.requireNonNull;
 
 public class AliasedRelation extends Relation {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AliasedRelation.class);
+
   private final Relation relation;
   private final Identifier alias;
   private final List<Identifier> columnNames;
@@ -109,4 +113,14 @@ public boolean shallowEquals(Node other) {
     return alias.equals(otherRelation.alias)
         && Objects.equals(columnNames, otherRelation.columnNames);
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(relation);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(alias);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeList(columnNames);
+    return size;
+  }
 }

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

@@ -21,6 +21,7 @@
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import javax.annotation.Nullable;
 
@@ -32,6 +33,9 @@
 
 public class AllColumns extends SelectItem {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AllColumns.class);
+
   private final List<Identifier> aliases;
   @Nullable private final Expression target;
 
@@ -127,4 +131,13 @@ public boolean shallowEquals(Node other) {
 
     return aliases.equals(((AllColumns) other).aliases);
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeList(aliases);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(target);
+    return size;
+  }
 }

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

@@ -20,6 +20,7 @@
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.tsfile.utils.RamUsageEstimator;
 
 import javax.annotation.Nonnull;
 
@@ -28,6 +29,7 @@
 import static java.util.Objects.requireNonNull;
 
 public final class AllRows extends Expression {
+  private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(AllRows.class);
 
   public AllRows() {
     super(null);
@@ -64,4 +66,11 @@ public int hashCode() {
   public boolean shallowEquals(Node other) {
     return sameClass(this, other);
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    return size;
+  }
 }

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

@@ -21,12 +21,16 @@
 
 import org.apache.iotdb.db.queryengine.plan.statement.metadata.DatabaseSchemaStatement;
 
+import org.apache.tsfile.utils.RamUsageEstimator;
+
 import java.util.List;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
 
 public class AlterDB extends DatabaseStatement {
+  private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(AlterDB.class);
+
   public AlterDB(
       final NodeLocation location,
       final boolean exists,
@@ -53,4 +57,9 @@ public String toString() {
         .add("properties", properties)
         .toString();
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE + ramBytesUsedForCommonFields();
+  }
 }

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

@@ -19,13 +19,17 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
+import org.apache.tsfile.utils.RamUsageEstimator;
+
 import java.util.Map;
 import java.util.Objects;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
 
 public class AlterPipe extends PipeStatement {
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AlterPipe.class);
 
   private final String pipeName;
   private final boolean ifExistsCondition;
@@ -140,4 +144,15 @@ public String toString() {
         .add("isReplaceAllConnectorAttributes", isReplaceAllConnectorAttributes)
         .toString();
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += RamUsageEstimator.sizeOf(pipeName);
+    size += RamUsageEstimator.sizeOfMap(extractorAttributes);
+    size += RamUsageEstimator.sizeOfMap(processorAttributes);
+    size += RamUsageEstimator.sizeOfMap(connectorAttributes);
+    return size;
+  }
 }