feat: add AstMemoryEstimator for SQL AST memory estimation

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/ast/AbstractTraverseDevice.java

@@ -187,6 +187,10 @@ public void setAttributeColumns(final List<String> attributeColumns) {
     this.attributeColumns = attributeColumns;
   }
 
+  public List<String> getAttributeColumns() {
+    return attributeColumns;
+  }
+
   public List<ColumnHeader> getColumnHeaderList() {
     return columnHeaderList;
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AddColumn.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/AliasedRelation.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,16 @@ 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);
+    if (columnNames != null) {
+      size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeList(columnNames);
+    }
+    return size;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AllColumns.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/AllRows.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/AlterDB.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,13 @@ public String toString() {
         .add("properties", properties)
         .toString();
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfString(dbName);
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeList(properties);
+    return size;
+  }
 }

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

@@ -19,13 +19,18 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.sql.ast;
 
+import org.apache.tsfile.utils.RamUsageEstimator;
+
+import java.util.ArrayList;
 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 +145,30 @@ public String toString() {
         .add("isReplaceAllConnectorAttributes", isReplaceAllConnectorAttributes)
         .toString();
   }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = INSTANCE_SIZE;
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
+    size += AstMemoryEstimationHelper.getEstimatedSizeOfString(pipeName);
+    size +=
+        AstMemoryEstimationHelper.getShallowSizeOfList(
+            new ArrayList<>(extractorAttributes.entrySet()));
+    size +=
+        AstMemoryEstimationHelper.getEstimatedSizeOfStringList(
+            new ArrayList<>(extractorAttributes.values()));
+    size +=
+        AstMemoryEstimationHelper.getShallowSizeOfList(
+            new ArrayList<>(processorAttributes.entrySet()));
+    size +=
+        AstMemoryEstimationHelper.getEstimatedSizeOfStringList(
+            new ArrayList<>(processorAttributes.values()));
+    size +=
+        AstMemoryEstimationHelper.getShallowSizeOfList(
+            new ArrayList<>(connectorAttributes.entrySet()));
+    size +=
+        AstMemoryEstimationHelper.getEstimatedSizeOfStringList(
+            new ArrayList<>(connectorAttributes.values()));
+    return size;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AnchorPattern.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;
@@ -33,6 +34,9 @@ public enum Type {
     PARTITION_END
   }
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(AnchorPattern.class);
+
   private final Type type;
 
   public AnchorPattern(NodeLocation location, Type type) {
@@ -80,4 +84,11 @@ public String toString() {
   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/ArithmeticBinaryExpression.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 org.apache.tsfile.utils.ReadWriteIOUtils;
 
 import javax.annotation.Nonnull;
@@ -34,6 +35,9 @@
 
 public class ArithmeticBinaryExpression extends Expression {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(ArithmeticBinaryExpression.class);
+
   public enum Operator {
     ADD("+"),
     SUBTRACT("-"),
@@ -139,4 +143,12 @@ public boolean shallowEquals(Node other) {
 
     return operator == ((ArithmeticBinaryExpression) other).operator;
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE
+        + AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal())
+        + AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(left)
+        + AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(right);
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/ArithmeticUnaryExpression.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 org.apache.tsfile.utils.ReadWriteIOUtils;
 
 import javax.annotation.Nonnull;
@@ -34,6 +35,9 @@
 
 public class ArithmeticUnaryExpression extends Expression {
 
+  private static final long INSTANCE_SIZE =
+      RamUsageEstimator.shallowSizeOfInstance(ArithmeticUnaryExpression.class);
+
   public enum Sign {
     PLUS,
     MINUS
@@ -141,4 +145,11 @@ public boolean shallowEquals(Node other) {
 
     return sign == ((ArithmeticUnaryExpression) other).sign;
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return INSTANCE_SIZE
+        + AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal())
+        + AstMemoryEstimationHelper.getEstimatedSizeOfAccountableObject(value);
+  }
 }