Pipe: improved the memory calculation for tablets (apache#16739)

* fix

* fix

* opti

* fix

* fix

* fix function name

* fix

* fix

* fix ut

---------

Co-authored-by: Tian Jiang <[email protected]>

Author: Caideyipi

integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java

@@ -1630,6 +1630,7 @@ public void insertRelationalTabletWithAutoCastTest()
     Collections.addAll(dataTypes, TSDataType.values());
     dataTypes.remove(TSDataType.VECTOR);
     dataTypes.remove(TSDataType.UNKNOWN);
+    dataTypes.remove(TSDataType.OBJECT);
 
     try {
       for (TSDataType from : dataTypes) {
@@ -1720,6 +1721,7 @@ public void insertRelationalRowWithAutoCastTest()
     Collections.addAll(dataTypes, TSDataType.values());
     dataTypes.remove(TSDataType.VECTOR);
     dataTypes.remove(TSDataType.UNKNOWN);
+    dataTypes.remove(TSDataType.OBJECT);
 
     for (TSDataType from : dataTypes) {
       for (TSDataType to : dataTypes) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/InsertNodeMemoryEstimator.java

@@ -189,7 +189,7 @@ private static long calculateFullInsertNodeSize(final InsertNode node) {
     // MeasurementSchemas
     size += sizeOfMeasurementSchemas(node.getMeasurementSchemas());
     // Measurement
-    size += sizeOfStringArray(node.getMeasurements());
+    size += RamUsageEstimator.sizeOf(node.getMeasurements());
     // dataTypes
     size += RamUsageEstimator.shallowSizeOf(node.getDataTypes());
     // columnCategories
@@ -231,7 +231,7 @@ private static long sizeOfInsertTabletNode(final InsertTabletNode node) {
     long size = INSERT_TABLET_NODE_SIZE;
     size += calculateFullInsertNodeSize(node);
     size += RamUsageEstimator.sizeOf(node.getTimes());
-    size += sizeOfBitMapArray(node.getBitMaps());
+    size += RamUsageEstimator.sizeOf(node.getBitMaps());
     size += sizeOfColumns(node.getColumns(), node.getMeasurementSchemas());
     final List<Integer> range = node.getRange();
     if (range != null) {
@@ -247,7 +247,7 @@ private static long calculateInsertTabletNodeSizeExcludingSchemas(final InsertTa
 
     size += RamUsageEstimator.sizeOf(node.getTimes());
 
-    size += sizeOfBitMapArray(node.getBitMaps());
+    size += RamUsageEstimator.sizeOf(node.getBitMaps());
 
     size += sizeOfColumns(node.getColumns(), node.getMeasurementSchemas());
 
@@ -390,7 +390,7 @@ private static long sizeOfRelationalInsertTabletNode(final RelationalInsertTable
 
     size += RamUsageEstimator.sizeOf(node.getTimes());
 
-    size += sizeOfBitMapArray(node.getBitMaps());
+    size += RamUsageEstimator.sizeOf(node.getBitMaps());
 
     size += sizeOfColumns(node.getColumns(), node.getMeasurementSchemas());
 
@@ -413,7 +413,7 @@ public static long sizeOfPartialPath(final PartialPath partialPath) {
     if (nodes != null) {
       // Since fullPath may be lazy loaded, lazy loading will not be triggered here, so it is
       // assumed that the memory size of fullPath is the same as that of nodes.
-      size += sizeOfStringArray(nodes) * 2;
+      size += RamUsageEstimator.sizeOf(nodes) * 2;
       size += TsFileConstant.PATH_SEPARATOR.length() * (nodes.length - 1) + NUM_BYTES_OBJECT_HEADER;
     }
     return size;
@@ -555,34 +555,6 @@ private static long sizeOfBinary(final Binary binary) {
     return Objects.nonNull(binary) ? binary.ramBytesUsed() : 0L;
   }
 
-  public static long sizeOfStringArray(final String[] values) {
-    return Objects.nonNull(values) ? RamUsageEstimator.sizeOf(values) : 0L;
-  }
-
-  public static long sizeOfBitMapArray(BitMap[] bitMaps) {
-    if (bitMaps == null) {
-      return 0L;
-    }
-    long size =
-        RamUsageEstimator.alignObjectSize(
-            NUM_BYTES_ARRAY_HEADER + NUM_BYTES_OBJECT_REF * bitMaps.length);
-    for (BitMap bitMap : bitMaps) {
-      size += sizeOfBitMap(bitMap);
-    }
-    return size;
-  }
-
-  private static long sizeOfBitMap(final BitMap bitMaps) {
-    if (bitMaps == null) {
-      return 0L;
-    }
-    long size = BIT_MAP_SIZE;
-
-    size +=
-        RamUsageEstimator.alignObjectSize(NUM_BYTES_ARRAY_HEADER + bitMaps.getByteArray().length);
-    return size;
-  }
-
   public static long sizeOfColumns(
       final Object[] columns, final MeasurementSchema[] measurementSchemas) {
     // Directly calculate if measurementSchemas are absent
@@ -631,7 +603,7 @@ public static long sizeOfColumns(
         case TEXT:
         case BLOB:
           {
-            size += getBinarySize((Binary[]) columns[i]);
+            size += RamUsageEstimator.sizeOf((Binary[]) columns[i]);
             break;
           }
       }
@@ -641,23 +613,10 @@ public static long sizeOfColumns(
 
   private static long getNumBytesUnknownObject(final Object obj) {
     return obj instanceof Binary[]
-        ? getBinarySize((Binary[]) obj)
+        ? RamUsageEstimator.sizeOf((Binary[]) obj)
         : RamUsageEstimator.sizeOfObject(obj);
   }
 
-  private static long getBinarySize(final Binary[] binaries) {
-    if (binaries == null) {
-      return 0L;
-    }
-
-    long size = 0L;
-    for (Binary binary : binaries) {
-      size += InsertNodeMemoryEstimator.sizeOfBinary(binary);
-    }
-
-    return size + RamUsageEstimator.shallowSizeOf(binaries);
-  }
-
   public static long sizeOfValues(
       final Object[] values, final MeasurementSchema[] measurementSchemas) {
     // Directly calculate if measurementSchemas are absent

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryWeightUtil.java

@@ -35,7 +35,6 @@
 import org.apache.tsfile.read.common.Field;
 import org.apache.tsfile.read.common.RowRecord;
 import org.apache.tsfile.utils.Binary;
-import org.apache.tsfile.utils.BitMap;
 import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.utils.TsPrimitiveType;
 import org.apache.tsfile.write.record.Tablet;
@@ -44,6 +43,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import static org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
 import static org.apache.tsfile.utils.RamUsageEstimator.NUM_BYTES_OBJECT_REF;
@@ -203,64 +203,8 @@ private static Pair<Integer, Integer> calculateTabletRowCountAndMemoryBySize(
     }
   }
 
-  public static long calculateTabletSizeInBytes(Tablet tablet) {
-    long totalSizeInBytes = 0;
-
-    if (tablet == null) {
-      return totalSizeInBytes;
-    }
-
-    long[] timestamps = tablet.getTimestamps();
-    Object[] tabletValues = tablet.getValues();
-
-    // timestamps
-    if (timestamps != null) {
-      totalSizeInBytes += timestamps.length * 8L;
-    }
-
-    // values
-    final List<IMeasurementSchema> timeseries = tablet.getSchemas();
-    if (timeseries != null) {
-      for (int column = 0; column < timeseries.size(); column++) {
-        final IMeasurementSchema measurementSchema = timeseries.get(column);
-        if (measurementSchema == null) {
-          continue;
-        }
-
-        final TSDataType tsDataType = measurementSchema.getType();
-        if (tsDataType == null) {
-          continue;
-        }
-
-        if (tsDataType.isBinary()) {
-          if (tabletValues == null || tabletValues.length <= column) {
-            continue;
-          }
-          final Binary[] values = ((Binary[]) tabletValues[column]);
-          if (values == null) {
-            continue;
-          }
-          for (Binary value : values) {
-            totalSizeInBytes += value == null ? 8 : value.ramBytesUsed();
-          }
-        } else {
-          totalSizeInBytes += (long) tablet.getMaxRowNumber() * tsDataType.getDataTypeSize();
-        }
-      }
-    }
-
-    // bitMaps
-    BitMap[] bitMaps = tablet.getBitMaps();
-    if (bitMaps != null) {
-      for (int i = 0; i < bitMaps.length; i++) {
-        totalSizeInBytes += bitMaps[i] == null ? 0 : bitMaps[i].getSize();
-      }
-    }
-
-    // estimate other dataStructures size
-    totalSizeInBytes += 100;
-
-    return totalSizeInBytes;
+  public static long calculateTabletSizeInBytes(final Tablet tablet) {
+    return Objects.nonNull(tablet) ? tablet.ramBytesUsed() : 0L;
   }
 
   public static long calculateTableSchemaBytesUsed(TableSchema tableSchema) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertBaseStatement.java

@@ -686,7 +686,7 @@ public long ramBytesUsed() {
     ramBytesUsed =
         InsertNodeMemoryEstimator.sizeOfPartialPath(devicePath)
             + InsertNodeMemoryEstimator.sizeOfMeasurementSchemas(measurementSchemas)
-            + InsertNodeMemoryEstimator.sizeOfStringArray(measurements)
+            + RamUsageEstimator.sizeOf(measurements)
             + RamUsageEstimator.shallowSizeOf(dataTypes)
             + RamUsageEstimator.shallowSizeOf(columnCategories)
             // We assume that the integers are all cached by JVM

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/crud/InsertTabletStatement.java

@@ -615,7 +615,7 @@ public void swapColumn(int src, int target) {
   protected long calculateBytesUsed() {
     return INSTANCE_SIZE
         + RamUsageEstimator.sizeOf(times)
-        + InsertNodeMemoryEstimator.sizeOfBitMapArray(nullBitMaps)
+        + RamUsageEstimator.sizeOf(nullBitMaps)
         + InsertNodeMemoryEstimator.sizeOfColumns(columns, measurementSchemas)
         + (Objects.nonNull(deviceIDs)
             ? Arrays.stream(deviceIDs)

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/PrimitiveArrayManager.java

@@ -249,6 +249,7 @@ private static Object createPrimitiveArray(TSDataType dataType) {
       case TEXT:
       case STRING:
       case BLOB:
+      case OBJECT:
         dataArray = new Binary[ARRAY_SIZE];
         break;
       default:

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/datastructure/TVList.java

@@ -970,7 +970,7 @@ public TsBlock nextBatch() {
       TSDataType dataType = getDataType();
       int maxRowCountOfCurrentBatch =
           Math.min(
-              paginationController.hasLimit()
+              paginationController.hasSetLimit()
                   ? (int) paginationController.getCurLimit()
                   : Integer.MAX_VALUE,
               Math.min(maxNumberOfPointsInPage, rows - index));

pom.xml

@@ -173,7 +173,7 @@
         <thrift.version>0.14.1</thrift.version>
         <xz.version>1.9</xz.version>
         <zstd-jni.version>1.5.6-3</zstd-jni.version>
-        <tsfile.version>2.2.0-251027-SNAPSHOT</tsfile.version>
+        <tsfile.version>2.2.0-251113-SNAPSHOT</tsfile.version>
     </properties>
     <!--
     if we claim dependencies in dependencyManagement, then we do not claim