Optimize pipe request serialization buffer sizing

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tablet/PipeInsertNodeTabletInsertionEvent.java

@@ -94,6 +94,8 @@ public class PipeInsertNodeTabletInsertionEvent extends PipeInsertionEvent
 
   private final AtomicReference<PipeTabletMemoryBlock> allocatedMemoryBlock;
   private volatile List<Tablet> tablets;
+  // Calculated together with tablets so downstream batching does not rescan Tablet internals.
+  private volatile long tabletsMemoryUsageInBytes;
 
   private List<TabletInsertionEventParser> eventParsers;
 
@@ -481,22 +483,31 @@ public boolean isAligned(final int i) {
   // TODO: for table model insertion, we need to get the database name
   public synchronized List<Tablet> convertToTablets() {
     if (Objects.isNull(tablets)) {
-      tablets =
-          initEventParsers().stream()
-              .map(TabletInsertionEventParser::convertToTablet)
-              .collect(Collectors.toList());
+      final List<TabletInsertionEventParser> parsers = initEventParsers();
+      final List<Tablet> convertedTablets = new ArrayList<>(parsers.size());
+      long tabletMemoryUsageInBytes = 0;
+      for (final TabletInsertionEventParser parser : parsers) {
+        final Tablet tablet = parser.convertToTablet();
+        convertedTablets.add(tablet);
+        // Tablet.ramBytesUsed() is required for the memory block to account for the actual
+        // retained tablet size. Calculate it while converting to avoid a second stream traversal.
+        tabletMemoryUsageInBytes += PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet);
+      }
+      tablets = convertedTablets;
+      tabletsMemoryUsageInBytes = tabletMemoryUsageInBytes;
       allocatedMemoryBlock.compareAndSet(
           null,
           PipeDataNodeResourceManager.memory()
-              .forceAllocateForTabletWithRetry(
-                  tablets.stream()
-                      .map(PipeMemoryWeightUtil::calculateTabletSizeInBytes)
-                      .reduce(Long::sum)
-                      .orElse(0L)));
+              .forceAllocateForTabletWithRetry(tabletMemoryUsageInBytes));
     }
     return tablets;
   }
 
+  public long getTabletsMemoryUsageInBytes() {
+    convertToTablets();
+    return tabletsMemoryUsageInBytes;
+  }
+
   /////////////////////////// event parser ///////////////////////////
 
   private List<TabletInsertionEventParser> initEventParsers() {
@@ -505,35 +516,27 @@ private List<TabletInsertionEventParser> initEventParsers() {
         return eventParsers;
       }
 
-      eventParsers = new ArrayList<>();
       final InsertNode node = getInsertNode();
       if (Objects.isNull(node)) {
         throw new PipeException(DataNodePipeMessages.INSERTNODE_HAS_BEEN_RELEASED);
       }
+      eventParsers = new ArrayList<>(getEventParserCount(node));
+      final UserEntity userEntity =
+          shouldParse4Privilege
+              ? new UserEntity(Long.parseLong(userId), userName, cliHostname)
+              : null;
       switch (node.getType()) {
         case INSERT_ROW:
         case INSERT_TABLET:
           eventParsers.add(
               new TabletInsertionEventTreePatternParser(
-                  pipeTaskMeta,
-                  this,
-                  node,
-                  treePattern,
-                  shouldParse4Privilege
-                      ? new UserEntity(Long.parseLong(userId), userName, cliHostname)
-                      : null));
+                  pipeTaskMeta, this, node, treePattern, userEntity));
           break;
         case INSERT_ROWS:
           for (final InsertRowNode insertRowNode : ((InsertRowsNode) node).getInsertRowNodeList()) {
             eventParsers.add(
                 new TabletInsertionEventTreePatternParser(
-                    pipeTaskMeta,
-                    this,
-                    insertRowNode,
-                    treePattern,
-                    shouldParse4Privilege
-                        ? new UserEntity(Long.parseLong(userId), userName, cliHostname)
-                        : null));
+                    pipeTaskMeta, this, insertRowNode, treePattern, userEntity));
           }
           break;
         case RELATIONAL_INSERT_ROW:
@@ -565,6 +568,16 @@ private List<TabletInsertionEventParser> initEventParsers() {
     }
   }
 
+  private static int getEventParserCount(final InsertNode node) {
+    if (node instanceof InsertRowsNode) {
+      return ((InsertRowsNode) node).getInsertRowNodeList().size();
+    }
+    if (node instanceof RelationalInsertRowsNode) {
+      return ((RelationalInsertRowsNode) node).getInsertRowNodeList().size();
+    }
+    return 1;
+  }
+
   public long count() {
     long count = 0;
     for (final Tablet covertedTablet : convertToTablets()) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParserTabletIterator.java

@@ -58,12 +58,12 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Predicate;
-import java.util.stream.Collectors;
 
 public class TsFileInsertionEventTableParserTabletIterator implements Iterator<Tablet> {
 
@@ -137,9 +137,12 @@ public TsFileInsertionEventTableParserTabletIterator(
     this.metadataQuerier = new MetadataQuerierByFileImpl(reader);
     fileMetadata = this.metadataQuerier.getWholeFileMetadata();
     final List<Map.Entry<String, TableSchema>> tableSchemaList =
-        fileMetadata.getTableSchemaMap().entrySet().stream()
-            .filter(predicate)
-            .collect(Collectors.toList());
+        new ArrayList<>(fileMetadata.getTableSchemaMap().size());
+    for (final Map.Entry<String, TableSchema> entry : fileMetadata.getTableSchemaMap().entrySet()) {
+      if (predicate.test(entry)) {
+        tableSchemaList.add(entry);
+      }
+    }
 
     this.allocatedMemoryBlockForTablet = allocatedMemoryBlockForTablet;
     this.allocatedMemoryBlockForBatchData = allocatedMemoryBlockForBatchData;
@@ -250,10 +253,10 @@ public boolean hasNext() {
               deviceMetaIterator = metadataQuerier.deviceIterator(tableRoot, null);
 
               final int columnSchemaSize = tableSchema.getColumnSchemas().size();
-              dataTypeList = new ArrayList<>();
-              columnTypes = new ArrayList<>();
-              measurementList = new ArrayList<>();
-              fieldSchemaList = new ArrayList<>();
+              dataTypeList = new ArrayList<>(columnSchemaSize);
+              columnTypes = new ArrayList<>(columnSchemaSize);
+              measurementList = new ArrayList<>(columnSchemaSize);
+              fieldSchemaList = new ArrayList<>(columnSchemaSize);
 
               for (int i = 0; i < columnSchemaSize; i++) {
                 final IMeasurementSchema schema = tableSchema.getColumnSchemas().get(i);
@@ -364,28 +367,27 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta
     timeChunk.getData().rewind();
     long size = timeChunkSize;
 
-    final List<Chunk> valueChunkList = new ArrayList<>();
+    final int fieldSchemaSize = fieldSchemaList.size();
+    final List<Chunk> valueChunkList = new ArrayList<>(fieldSchemaSize);
     final Map<String, IChunkMetadata> valueChunkMetadataMap =
-        alignedChunkMetadata.getValueChunkMetadataList().stream()
-            .filter(Objects::nonNull)
-            .filter(
-                metadata ->
-                    !isFieldDeletedByMods(
-                        metadata.getMeasurementUid(),
-                        alignedChunkMetadata.getStartTime(),
-                        alignedChunkMetadata.getEndTime()))
-            .collect(
-                Collectors.toMap(
-                    IChunkMetadata::getMeasurementUid,
-                    metadata -> metadata,
-                    (left, right) -> left));
+        new HashMap<>((int) (fieldSchemaSize / 0.75f) + 1);
+    for (final IChunkMetadata metadata : alignedChunkMetadata.getValueChunkMetadataList()) {
+      if (metadata != null
+          && !isFieldDeletedByMods(
+              metadata.getMeasurementUid(),
+              alignedChunkMetadata.getStartTime(),
+              alignedChunkMetadata.getEndTime())) {
+        // Keep the first metadata entry to preserve the former merge-function behavior.
+        valueChunkMetadataMap.putIfAbsent(metadata.getMeasurementUid(), metadata);
+      }
+    }
 
     // To ensure that the Tablet has the same alignedChunk column as the current one,
     // you need to create a new Tablet to fill in the data.
     isSameDeviceID = false;
 
     // Need to ensure that columnTypes recreates an array
-    final List<ColumnCategory> categories = new ArrayList<>(deviceIdSize);
+    final List<ColumnCategory> categories = new ArrayList<>(deviceIdSize + fieldSchemaSize);
     for (int i = 0; i < deviceIdSize; i++) {
       categories.add(ColumnCategory.TAG);
     }

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

@@ -733,24 +733,16 @@ private static long sizeOfObjectList(final List<?> list) {
     if (list == null) {
       return 0L;
     }
-    long size = RamUsageEstimator.shallowSizeOf(list);
-    if (list instanceof ArrayList) {
-      size +=
-          RamUsageEstimator.alignObjectSize(
-              NUM_BYTES_ARRAY_HEADER + NUM_BYTES_OBJECT_REF * list.size());
-    }
-    return size;
+    return SIZE_OF_ARRAYLIST
+        + RamUsageEstimator.alignObjectSize(
+            NUM_BYTES_ARRAY_HEADER + NUM_BYTES_OBJECT_REF * list.size());
   }
 
   private static long sizeOfIntegerList(final List<Integer> integers) {
     if (integers == null) {
       return 0L;
     }
-    long size = sizeOfObjectList(integers);
-    for (Integer ignored : integers) {
-      size += SIZE_OF_INT;
-    }
-    return size;
+    return sizeOfObjectList(integers) + (long) SIZE_OF_INT * integers.size();
   }
 
   private static long sizeOfResults(final Map<Integer, TSStatus> results) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java

@@ -131,7 +131,8 @@ public PipeTransferTabletBatchReqV2 toTPipeTransferReq() throws IOException {
           }
         }
         for (final Pair<Boolean, Tablet> tabletPair : batchTablets) {
-          try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+          try (final PublicBAOS byteArrayOutputStream =
+                  new PublicBAOS(calculateTabletSerializedSize(tabletPair.getRight()));
               final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
             tabletPair.getRight().serialize(outputStream);
             ReadWriteIOUtils.write(true, outputStream);
@@ -176,7 +177,12 @@ private long buildTabletInsertionBuffer(final TabletInsertionEvent event) throws
         insertNodeDataBases.add(databaseName);
       } else {
         final List<Tablet> tablets = pipeInsertNodeTabletInsertionEvent.convertToTablets();
-        estimateSize = calculateTabletsSizeInBytes(tablets);
+        // convertToTablets() has already measured every tablet for the event memory block. Reuse
+        // that exact measurement instead of calling Tablet.ramBytesUsed() (which walks the schema
+        // map) once more while building this batch.
+        estimateSize =
+            pipeInsertNodeTabletInsertionEvent.getTabletsMemoryUsageInBytes()
+                + (long) Integer.BYTES * tablets.size();
         increaseTotalBufferSizeAndUpdateMemoryBlock(estimateSize);
         for (final Tablet tablet : tablets) {
           constructTabletBatchWithoutMemoryReservation(
@@ -192,9 +198,11 @@ private long buildTabletInsertionBuffer(final TabletInsertionEvent event) throws
                 pipeRawTabletInsertionEvent.convertToTablet(),
                 pipeRawTabletInsertionEvent.getTableModelDatabaseName());
       } else {
-        try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+        final Tablet tablet = pipeRawTabletInsertionEvent.convertToTablet();
+        try (final PublicBAOS byteArrayOutputStream =
+                new PublicBAOS(calculateTabletSerializedSize(tablet));
             final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
-          pipeRawTabletInsertionEvent.convertToTablet().serialize(outputStream);
+          tablet.serialize(outputStream);
           ReadWriteIOUtils.write(pipeRawTabletInsertionEvent.isAligned(), outputStream);
           buffer = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size());
         }
@@ -226,14 +234,14 @@ private void constructTabletBatchWithoutMemoryReservation(
     currentBatch.getRight().add(tablet);
   }
 
-  private long calculateTabletsSizeInBytes(final List<Tablet> tablets) {
-    return tablets.stream().mapToLong(PipeTabletEventPlainBatch::calculateTabletSizeInBytes).sum();
-  }
-
   private static long calculateTabletSizeInBytes(final Tablet tablet) {
     return PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet) + 4;
   }
 
+  private static int calculateTabletSerializedSize(final Tablet tablet) {
+    return tablet.serializedSize() + Byte.BYTES;
+  }
+
   static boolean mayAppendTablet(final Tablet target, final Tablet source) {
     // Tablet.append already checks schemas and column categories. Avoid repeating those potentially
     // expensive comparisons here because wide-table pipe transfer can have many columns.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReq.java

@@ -46,6 +46,11 @@
 
 public class PipeTransferTabletBatchReq extends TPipeTransferReq {
 
+  private static final int BATCH_REQUEST_COUNT_SERIALIZED_SIZE =
+      Integer.BYTES // legacy binary request count
+          + Integer.BYTES // insert node request count
+          + Integer.BYTES; // raw tablet request count
+
   private final transient List<PipeTransferTabletBinaryReq> binaryReqs = new ArrayList<>();
   private final transient List<PipeTransferTabletInsertNodeReq> insertNodeReqs = new ArrayList<>();
   private final transient List<PipeTransferTabletRawReq> tabletReqs = new ArrayList<>();
@@ -129,19 +134,28 @@ public static PipeTransferTabletBatchReq toTPipeTransferReq(
 
     batchReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
     batchReq.type = PipeRequestType.TRANSFER_TABLET_BATCH.getType();
-    try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+    try (final PublicBAOS byteArrayOutputStream =
+            new PublicBAOS(calculateSerializedSize(insertNodeBuffers, tabletBuffers));
         final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
       // Binary buffer, for rolling upgrade
       ReadWriteIOUtils.write(0, outputStream);
 
+      // Insert-node and raw-tablet serializations are self-delimiting, so their lengths are not
+      // written separately.
       ReadWriteIOUtils.write(insertNodeBuffers.size(), outputStream);
       for (final ByteBuffer insertNodeBuffer : insertNodeBuffers) {
-        outputStream.write(insertNodeBuffer.array(), 0, insertNodeBuffer.limit());
+        outputStream.write(
+            insertNodeBuffer.array(),
+            insertNodeBuffer.arrayOffset() + insertNodeBuffer.position(),
+            insertNodeBuffer.remaining());
       }
 
       ReadWriteIOUtils.write(tabletBuffers.size(), outputStream);
       for (final ByteBuffer tabletBuffer : tabletBuffers) {
-        outputStream.write(tabletBuffer.array(), 0, tabletBuffer.limit());
+        outputStream.write(
+            tabletBuffer.array(),
+            tabletBuffer.arrayOffset() + tabletBuffer.position(),
+            tabletBuffer.remaining());
       }
 
       batchReq.body =
@@ -151,6 +165,13 @@ public static PipeTransferTabletBatchReq toTPipeTransferReq(
     return batchReq;
   }
 
+  static int calculateSerializedSize(
+      final List<ByteBuffer> insertNodeBuffers, final List<ByteBuffer> tabletBuffers) {
+    return BATCH_REQUEST_COUNT_SERIALIZED_SIZE
+        + insertNodeBuffers.stream().mapToInt(ByteBuffer::remaining).sum()
+        + tabletBuffers.stream().mapToInt(ByteBuffer::remaining).sum();
+  }
+
   public static PipeTransferTabletBatchReq fromTPipeTransferReq(
       final TPipeTransferReq transferReq) {
     final PipeTransferTabletBatchReq batchReq = new PipeTransferTabletBatchReq();