Pipe: Optimized the table model writing latency by batching & Fixed the NPE caused by tablet event sorting

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

@@ -22,14 +22,18 @@
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeInsertNodeTabletInsertionEvent;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
+import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryWeightUtil;
 import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTabletBatchReqV2;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
-import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
+import org.apache.iotdb.pipe.api.exception.PipeException;
 
 import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.utils.PublicBAOS;
+import org.apache.tsfile.utils.RamUsageEstimator;
 import org.apache.tsfile.utils.ReadWriteIOUtils;
+import org.apache.tsfile.write.record.Tablet;
 
 import java.io.DataOutputStream;
 import java.io.IOException;
@@ -39,6 +43,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
 
 public class PipeTabletEventPlainBatch extends PipeTabletEventBatch {
@@ -52,13 +57,12 @@ public class PipeTabletEventPlainBatch extends PipeTabletEventBatch {
   private final List<String> insertNodeDataBases = new ArrayList<>();
   private final List<String> tabletDataBases = new ArrayList<>();
 
+  private final Map<String, Map<String, Pair<Integer, List<Tablet>>>> tableModelTabletMap =
+      new HashMap<>();
+
   // Used to rate limit when transferring data
   private final Map<Pair<String, Long>, Long> pipe2BytesAccumulated = new HashMap<>();
 
-  PipeTabletEventPlainBatch(final int maxDelayInMs, final long requestMaxBatchSizeInBytes) {
-    super(maxDelayInMs, requestMaxBatchSizeInBytes, null);
-  }
-
   PipeTabletEventPlainBatch(
       final int maxDelayInMs,
       final long requestMaxBatchSizeInBytes,
@@ -67,9 +71,8 @@ public class PipeTabletEventPlainBatch extends PipeTabletEventBatch {
   }
 
   @Override
-  protected boolean constructBatch(final TabletInsertionEvent event)
-      throws WALPipeException, IOException {
-    final int bufferSize = buildTabletInsertionBuffer(event);
+  protected boolean constructBatch(final TabletInsertionEvent event) throws IOException {
+    final long bufferSize = buildTabletInsertionBuffer(event);
     totalBufferSize += bufferSize;
     pipe2BytesAccumulated.compute(
         new Pair<>(
@@ -90,11 +93,43 @@ public synchronized void onSuccess() {
     binaryDataBases.clear();
     insertNodeDataBases.clear();
     tabletDataBases.clear();
+    tableModelTabletMap.clear();
 
     pipe2BytesAccumulated.clear();
   }
 
   public PipeTransferTabletBatchReqV2 toTPipeTransferReq() throws IOException {
+    for (final Map.Entry<String, Map<String, Pair<Integer, List<Tablet>>>> insertTablets :
+        tableModelTabletMap.entrySet()) {
+      final String databaseName = insertTablets.getKey();
+      for (final Map.Entry<String, Pair<Integer, List<Tablet>>> tabletEntry :
+          insertTablets.getValue().entrySet()) {
+        Tablet batchTablet = null;
+        for (final Tablet tablet : tabletEntry.getValue().getRight()) {
+          if (Objects.isNull(batchTablet)) {
+            batchTablet = tablet;
+          } else if (!batchTablet.append(tablet, tabletEntry.getValue().getLeft())) {
+            throw new PipeException(
+                "Failed to merge tablets due to inconsistent schema, database: "
+                    + databaseName
+                    + ", tableName: "
+                    + tablet.getTableName());
+          }
+        }
+        assert batchTablet != null;
+        try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+            final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
+          batchTablet.serialize(outputStream);
+          ReadWriteIOUtils.write(true, outputStream);
+          tabletBuffers.add(
+              ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()));
+        }
+        tabletDataBases.add(databaseName);
+      }
+    }
+
+    tableModelTabletMap.clear();
+
     return PipeTransferTabletBatchReqV2.toTPipeTransferReq(
         binaryBuffers,
         insertNodeBuffers,
@@ -112,57 +147,69 @@ public Map<Pair<String, Long>, Long> getPipe2BytesAccumulated() {
     return pipe2BytesAccumulated;
   }
 
-  private int buildTabletInsertionBuffer(final TabletInsertionEvent event)
-      throws IOException, WALPipeException {
-    int databaseEstimateSize = 0;
+  private long buildTabletInsertionBuffer(final TabletInsertionEvent event) throws IOException {
+    long estimateSize;
     final ByteBuffer buffer;
     if (event instanceof PipeInsertNodeTabletInsertionEvent) {
       final PipeInsertNodeTabletInsertionEvent pipeInsertNodeTabletInsertionEvent =
           (PipeInsertNodeTabletInsertionEvent) event;
-      // Read the bytebuffer from the wal file and transfer it directly without serializing or
-      // deserializing if possible
       final InsertNode insertNode = pipeInsertNodeTabletInsertionEvent.getInsertNode();
-      if (Objects.isNull(insertNode)) {
-        buffer = pipeInsertNodeTabletInsertionEvent.getByteBuffer();
-        binaryBuffers.add(buffer);
-        if (pipeInsertNodeTabletInsertionEvent.isTableModelEvent()) {
-          databaseEstimateSize =
-              pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName().length();
-          binaryDataBases.add(pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName());
-        } else {
-          databaseEstimateSize = 4;
-          binaryDataBases.add(TREE_MODEL_DATABASE_PLACEHOLDER);
-        }
-      } else {
+      if (!(insertNode instanceof RelationalInsertTabletNode)) {
         buffer = insertNode.serializeToByteBuffer();
         insertNodeBuffers.add(buffer);
         if (pipeInsertNodeTabletInsertionEvent.isTableModelEvent()) {
-          databaseEstimateSize =
-              pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName().length();
+          estimateSize =
+              RamUsageEstimator.sizeOf(
+                  pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName());
           insertNodeDataBases.add(pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName());
         } else {
-          databaseEstimateSize = 4;
+          estimateSize = 4;
           insertNodeDataBases.add(TREE_MODEL_DATABASE_PLACEHOLDER);
         }
+        estimateSize += buffer.limit();
+      } else {
+        estimateSize =
+            constructTabletBatch(
+                ((RelationalInsertTabletNode) insertNode).convertToTablet(),
+                pipeInsertNodeTabletInsertionEvent.getTableModelDatabaseName());
       }
     } else {
       final PipeRawTabletInsertionEvent pipeRawTabletInsertionEvent =
           (PipeRawTabletInsertionEvent) event;
-      try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
-          final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
-        pipeRawTabletInsertionEvent.convertToTablet().serialize(outputStream);
-        ReadWriteIOUtils.write(pipeRawTabletInsertionEvent.isAligned(), outputStream);
-        buffer = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size());
-      }
-      tabletBuffers.add(buffer);
       if (pipeRawTabletInsertionEvent.isTableModelEvent()) {
-        databaseEstimateSize = pipeRawTabletInsertionEvent.getTableModelDatabaseName().length();
-        tabletDataBases.add(pipeRawTabletInsertionEvent.getTableModelDatabaseName());
+        estimateSize =
+            constructTabletBatch(
+                pipeRawTabletInsertionEvent.convertToTablet(),
+                pipeRawTabletInsertionEvent.getTableModelDatabaseName());
       } else {
-        databaseEstimateSize = 4;
+        try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+            final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
+          pipeRawTabletInsertionEvent.convertToTablet().serialize(outputStream);
+          ReadWriteIOUtils.write(pipeRawTabletInsertionEvent.isAligned(), outputStream);
+          buffer = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size());
+        }
+        estimateSize = 4 + buffer.limit();
+        tabletBuffers.add(buffer);
         tabletDataBases.add(TREE_MODEL_DATABASE_PLACEHOLDER);
       }
     }
-    return buffer.limit() + databaseEstimateSize;
+
+    return estimateSize;
+  }
+
+  private long constructTabletBatch(final Tablet tablet, final String databaseName) {
+    final AtomicLong size = new AtomicLong(0);
+    final Pair<Integer, List<Tablet>> currentBatch =
+        tableModelTabletMap
+            .computeIfAbsent(
+                databaseName,
+                k -> {
+                  size.addAndGet(RamUsageEstimator.sizeOf(databaseName));
+                  return new HashMap<>();
+                })
+            .computeIfAbsent(tablet.getTableName(), k -> new Pair<>(0, new ArrayList<>()));
+    currentBatch.setLeft(currentBatch.getLeft() + tablet.getRowSize());
+    currentBatch.getRight().add(tablet);
+    return PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet) + 4;
   }
 }

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

@@ -39,7 +39,9 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 public class PipeTransferTabletBatchReqV2 extends TPipeTransferReq {
@@ -62,14 +64,31 @@ public List<InsertBaseStatement> constructStatements() {
 
     final List<InsertRowStatement> insertRowStatementList = new ArrayList<>();
     final List<InsertTabletStatement> insertTabletStatementList = new ArrayList<>();
+    final Map<String, List<InsertRowStatement>> tableModelDatabaseInsertRowStatementMap =
+        new HashMap<>();
 
     for (final PipeTransferTabletBinaryReqV2 binaryReq : binaryReqs) {
       final InsertBaseStatement statement = binaryReq.constructStatement();
       if (statement.isEmpty()) {
         continue;
       }
       if (statement.isWriteToTable()) {
-        statements.add(statement);
+        if (statement instanceof InsertRowStatement) {
+          tableModelDatabaseInsertRowStatementMap
+              .computeIfAbsent(statement.getDatabaseName().get(), k -> new ArrayList<>())
+              .add((InsertRowStatement) statement);
+        } else if (statement instanceof InsertTabletStatement) {
+          statements.add(statement);
+        } else if (statement instanceof InsertRowsStatement) {
+          tableModelDatabaseInsertRowStatementMap
+              .computeIfAbsent(statement.getDatabaseName().get(), k -> new ArrayList<>())
+              .addAll(((InsertRowsStatement) statement).getInsertRowStatementList());
+        } else {
+          throw new UnsupportedOperationException(
+              String.format(
+                  "unknown InsertBaseStatement %s constructed from PipeTransferTabletBinaryReqV2.",
+                  binaryReq));
+        }
         continue;
       }
       if (statement instanceof InsertRowStatement) {
@@ -93,7 +112,22 @@ public List<InsertBaseStatement> constructStatements() {
         continue;
       }
       if (statement.isWriteToTable()) {
-        statements.add(statement);
+        if (statement instanceof InsertRowStatement) {
+          tableModelDatabaseInsertRowStatementMap
+              .computeIfAbsent(statement.getDatabaseName().get(), k -> new ArrayList<>())
+              .add((InsertRowStatement) statement);
+        } else if (statement instanceof InsertTabletStatement) {
+          statements.add(statement);
+        } else if (statement instanceof InsertRowsStatement) {
+          tableModelDatabaseInsertRowStatementMap
+              .computeIfAbsent(statement.getDatabaseName().get(), k -> new ArrayList<>())
+              .addAll(((InsertRowsStatement) statement).getInsertRowStatementList());
+        } else {
+          throw new UnsupportedOperationException(
+              String.format(
+                  "unknown InsertBaseStatement %s constructed from PipeTransferTabletBinaryReqV2.",
+                  insertNodeReq));
+        }
         continue;
       }
       if (statement instanceof InsertRowStatement) {
@@ -131,6 +165,16 @@ public List<InsertBaseStatement> constructStatements() {
     if (!insertMultiTabletsStatement.isEmpty()) {
       statements.add(insertMultiTabletsStatement);
     }
+
+    for (final Map.Entry<String, List<InsertRowStatement>> insertRows :
+        tableModelDatabaseInsertRowStatementMap.entrySet()) {
+      final InsertRowsStatement statement = new InsertRowsStatement();
+      statement.setWriteToTable(true);
+      statement.setDatabaseName(insertRows.getKey());
+      statement.setInsertRowStatementList(insertRows.getValue());
+      statements.add(statement);
+    }
+
     return statements;
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/sorter/PipeTabletEventSorter.java

@@ -152,15 +152,15 @@ protected Object reorderValueListAndBitMap(
 
   private int getLastNonnullIndex(
       final int i, final BitMap originalBitMap, final BitMap deDuplicatedBitMap) {
+    if (originalBitMap == null) {
+      return index[deDuplicatedIndex[i]];
+    }
     if (deDuplicatedIndex == null) {
       if (originalBitMap.isMarked(index[i])) {
         deDuplicatedBitMap.mark(i);
       }
       return index[i];
     }
-    if (originalBitMap == null) {
-      return index[deDuplicatedIndex[i]];
-    }
     int lastNonnullIndex = deDuplicatedIndex[i];
     int lastIndex = i > 0 ? deDuplicatedIndex[i - 1] : -1;
     while (originalBitMap.isMarked(index[lastNonnullIndex])) {

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

@@ -24,6 +24,7 @@
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
 import org.apache.iotdb.commons.utils.CommonDateTimeUtils;
 import org.apache.iotdb.commons.utils.TestOnly;
 import org.apache.iotdb.commons.utils.TimePartitionUtils;
@@ -55,6 +56,7 @@
 import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.apache.tsfile.utils.TsPrimitiveType;
 import org.apache.tsfile.write.UnSupportedDataTypeException;
+import org.apache.tsfile.write.record.Tablet;
 import org.apache.tsfile.write.schema.MeasurementSchema;
 
 import java.io.DataInputStream;
@@ -70,6 +72,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 import static org.apache.iotdb.db.utils.CommonUtils.isAlive;
 
@@ -187,6 +190,22 @@ public void setRange(List<Integer> range) {
     this.range = range;
   }
 
+  public Tablet convertToTablet() {
+    final Tablet tablet =
+        new Tablet(
+            getTableName(),
+            Arrays.asList(measurements),
+            Arrays.asList(dataTypes),
+            Arrays.stream(columnCategories)
+                .map(TsTableColumnCategory::toTsFileColumnType)
+                .collect(Collectors.toList()),
+            rowCount);
+    tablet.setBitMaps(bitMaps);
+    tablet.setTimestamps(times);
+    tablet.setValues(columns);
+    return tablet;
+  }
+
   @Override
   public void addChild(PlanNode child) {}
 

pom.xml

@@ -175,7 +175,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-250722-SNAPSHOT</tsfile.version>
+        <tsfile.version>2.2.0-250730-SNAPSHOT</tsfile.version>
     </properties>
     <!--
     if we claim dependencies in dependencyManagement, then we do not claim