diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java index bdf6ee1874a4..832287c7148e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java +++ b/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 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); @@ -192,9 +193,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()); } @@ -234,6 +237,10 @@ 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. diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReq.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReq.java index 352ff0bfc63a..917547b278e5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReq.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReq.java @@ -129,7 +129,8 @@ 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); @@ -151,6 +152,13 @@ public static PipeTransferTabletBatchReq toTPipeTransferReq( return batchReq; } + static int calculateSerializedSize( + final List insertNodeBuffers, final List tabletBuffers) { + return Integer.BYTES * 3 + + insertNodeBuffers.stream().mapToInt(ByteBuffer::limit).sum() + + tabletBuffers.stream().mapToInt(ByteBuffer::limit).sum(); + } + public static PipeTransferTabletBatchReq fromTPipeTransferReq( final TPipeTransferReq transferReq) { final PipeTransferTabletBatchReq batchReq = new PipeTransferTabletBatchReq(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReqV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReqV2.java index 6c4607518b4e..2415aaf16c5d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReqV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBatchReqV2.java @@ -33,6 +33,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; +import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; @@ -186,7 +187,10 @@ public static PipeTransferTabletBatchReqV2 toTPipeTransferReq( batchReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); batchReq.type = PipeRequestType.TRANSFER_TABLET_BATCH_V2.getType(); - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS( + calculateSerializedSize( + insertNodeBuffers, tabletBuffers, insertNodeDataBases, tabletDataBases)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { // Binary buffer, for rolling upgrade ReadWriteIOUtils.write(0, outputStream); @@ -212,6 +216,27 @@ public static PipeTransferTabletBatchReqV2 toTPipeTransferReq( return batchReq; } + static int calculateSerializedSize( + final List insertNodeBuffers, + final List tabletBuffers, + final List insertNodeDataBases, + final List tabletDataBases) { + int size = Integer.BYTES * 3; + for (int i = 0; i < insertNodeBuffers.size(); i++) { + size += insertNodeBuffers.get(i).limit(); + size += serializedStringSize(insertNodeDataBases.get(i)); + } + for (int i = 0; i < tabletBuffers.size(); i++) { + size += tabletBuffers.get(i).limit(); + size += serializedStringSize(tabletDataBases.get(i)); + } + return size; + } + + private static int serializedStringSize(final String value) { + return Integer.BYTES + (value == null ? 0 : value.getBytes(TSFileConfig.STRING_CHARSET).length); + } + public static PipeTransferTabletBatchReqV2 fromTPipeTransferReq( final org.apache.iotdb.service.rpc.thrift.TPipeTransferReq transferReq) { final PipeTransferTabletBatchReqV2 batchReq = new PipeTransferTabletBatchReqV2(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReq.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReq.java index b31816c1fcd6..0d790ff1bc1b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReq.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReq.java @@ -102,7 +102,7 @@ public static PipeTransferTabletBinaryReq fromTPipeTransferReq( /////////////////////////////// Air Gap /////////////////////////////// public static byte[] toTPipeTransferBytes(final ByteBuffer byteBuffer) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(Byte.BYTES + Short.BYTES); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_TABLET_BINARY.getType(), outputStream); @@ -110,6 +110,10 @@ public static byte[] toTPipeTransferBytes(final ByteBuffer byteBuffer) throws IO } } + static int calculateSerializedSize(final ByteBuffer byteBuffer) { + return Byte.BYTES + Short.BYTES + byteBuffer.limit(); + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReqV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReqV2.java index 2788033be2d0..f80925da1441 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReqV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletBinaryReqV2.java @@ -32,6 +32,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement; +import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; @@ -119,7 +120,8 @@ public static PipeTransferTabletBinaryReqV2 toTPipeTransferReq( req.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); req.type = PipeRequestType.TRANSFER_TABLET_BINARY_V2.getType(); - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateSerializedSize(byteBuffer, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(byteBuffer.limit(), outputStream); outputStream.write(byteBuffer.array(), 0, byteBuffer.limit()); @@ -151,7 +153,8 @@ public static PipeTransferTabletBinaryReqV2 fromTPipeTransferReq( public static byte[] toTPipeTransferBytes(final ByteBuffer byteBuffer, final String dataBaseName) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateAirGapSerializedSize(byteBuffer, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_TABLET_BINARY_V2.getType(), outputStream); @@ -162,6 +165,18 @@ public static byte[] toTPipeTransferBytes(final ByteBuffer byteBuffer, final Str } } + private static int serializedStringSize(final String value) { + return Integer.BYTES + (value == null ? 0 : value.getBytes(TSFileConfig.STRING_CHARSET).length); + } + + static int calculateSerializedSize(final ByteBuffer byteBuffer, final String dataBaseName) { + return Integer.BYTES + byteBuffer.limit() + serializedStringSize(dataBaseName); + } + + static int calculateAirGapSerializedSize(final ByteBuffer byteBuffer, final String dataBaseName) { + return Byte.BYTES + Short.BYTES + calculateSerializedSize(byteBuffer, dataBaseName); + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReq.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReq.java index bc42630d79b4..14033fb5f135 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReq.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReq.java @@ -31,12 +31,12 @@ import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertBaseStatement; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; -import org.apache.tsfile.utils.BytesUtils; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; import java.io.DataOutputStream; import java.io.IOException; +import java.nio.ByteBuffer; import java.util.Objects; public class PipeTransferTabletInsertNodeReq extends TPipeTransferReq { @@ -86,7 +86,14 @@ public static PipeTransferTabletInsertNodeReq toTPipeTransferReq(final InsertNod req.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); req.type = PipeRequestType.TRANSFER_TABLET_INSERT_NODE.getType(); - req.body = insertNode.serializeToByteBuffer(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateSerializedSize(insertNode)); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + insertNode.serialize(outputStream); + req.body = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } catch (final IOException e) { + throw new RuntimeException(e); + } return req; } @@ -106,15 +113,28 @@ public static PipeTransferTabletInsertNodeReq fromTPipeTransferReq( /////////////////////////////// Air Gap /////////////////////////////// public static byte[] toTPipeTransferBytes(final InsertNode insertNode) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateAirGapSerializedSize(insertNode)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_TABLET_INSERT_NODE.getType(), outputStream); - return BytesUtils.concatByteArray( - byteArrayOutputStream.toByteArray(), insertNode.serializeToByteBuffer().array()); + insertNode.serialize(outputStream); + return byteArrayOutputStream.toByteArray(); } } + static int calculateSerializedSize(final InsertNode insertNode) { + return insertNode.serializeToByteBufferSize(); + } + + static int calculateAirGapSerializedSize(final InsertNode insertNode) { + return calculateAirGapSerializedSize(calculateSerializedSize(insertNode)); + } + + protected static int calculateAirGapSerializedSize(final int bodySize) { + return Byte.BYTES + Short.BYTES + bodySize; + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReqV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReqV2.java index b9d5eb7de85b..4fc18398dda7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReqV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletInsertNodeReqV2.java @@ -120,7 +120,8 @@ public static PipeTransferTabletInsertNodeReqV2 toTPipeTransferReq( req.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); req.type = PipeRequestType.TRANSFER_TABLET_INSERT_NODE_V2.getType(); - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateSerializedSize(insertNode, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { insertNode.serialize(outputStream); ReadWriteIOUtils.write(req.dataBaseName, outputStream); @@ -150,7 +151,8 @@ public static PipeTransferTabletInsertNodeReqV2 fromTPipeTransferReq( public static byte[] toTPipeTransferBytes(final InsertNode insertNode, final String dataBaseName) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateAirGapSerializedSize(insertNode, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write( @@ -161,6 +163,16 @@ public static byte[] toTPipeTransferBytes(final InsertNode insertNode, final Str } } + static int calculateSerializedSize(final InsertNode insertNode, final String dataBaseName) { + return PipeTransferTabletInsertNodeReq.calculateSerializedSize(insertNode) + + ReadWriteIOUtils.sizeToWrite(dataBaseName); + } + + static int calculateAirGapSerializedSize(final InsertNode insertNode, final String dataBaseName) { + return PipeTransferTabletInsertNodeReq.calculateAirGapSerializedSize( + calculateSerializedSize(insertNode, dataBaseName)); + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReq.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReq.java index 01c80758152d..00907ed8302f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReq.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReq.java @@ -135,7 +135,7 @@ public static PipeTransferTabletRawReq toTPipeTransferReq( tabletReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); tabletReq.type = PipeRequestType.TRANSFER_TABLET_RAW.getType(); - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(calculateSerializedSize(tablet)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { tablet.serialize(outputStream); ReadWriteIOUtils.write(isAligned, outputStream); @@ -272,7 +272,8 @@ public byte[] toTPipeTransferBytes() throws IOException { throw new IOException(DataNodePipeMessages.CANNOT_SERIALIZE_BOTH_TABLET_AND_STATEMENT_ARE); } - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateAirGapSerializedSize(tabletToSerialize)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_TABLET_RAW.getType(), outputStream); @@ -298,6 +299,14 @@ public static byte[] toTPipeTransferBytes(final Tablet tablet, final boolean isA return req.toTPipeTransferBytes(); } + static int calculateSerializedSize(final Tablet tablet) { + return tablet.serializedSize() + Byte.BYTES; + } + + static int calculateAirGapSerializedSize(final Tablet tablet) { + return Byte.BYTES + Short.BYTES + calculateSerializedSize(tablet); + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReqV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReqV2.java index d395bf6cf5f2..1b6c24d57166 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReqV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferTabletRawReqV2.java @@ -144,7 +144,8 @@ public static PipeTransferTabletRawReqV2 toTPipeTransferReq( tabletReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); tabletReq.type = PipeRequestType.TRANSFER_TABLET_RAW_V2.getType(); - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS(calculateSerializedSize(tablet, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { tablet.serialize(outputStream); ReadWriteIOUtils.write(isAligned, outputStream); @@ -173,7 +174,9 @@ public static PipeTransferTabletRawReqV2 fromTPipeTransferReq( public static byte[] toTPipeTransferBytes( final Tablet tablet, final boolean isAligned, final String dataBaseName) throws IOException { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + try (final PublicBAOS byteArrayOutputStream = + new PublicBAOS( + Byte.BYTES + Short.BYTES + calculateSerializedSize(tablet, dataBaseName)); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_TABLET_RAW_V2.getType(), outputStream); @@ -184,6 +187,15 @@ public static byte[] toTPipeTransferBytes( } } + static int calculateSerializedSize(final Tablet tablet, final String dataBaseName) { + return tablet.serializedSize() + + Byte.BYTES + + Integer.BYTES + + (dataBaseName == null + ? 0 + : dataBaseName.getBytes(java.nio.charset.StandardCharsets.UTF_8).length); + } + /////////////////////////////// Object /////////////////////////////// @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/builder/IoTConsensusV2TransferBatchReqBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/builder/IoTConsensusV2TransferBatchReqBuilder.java index fc387084a001..8b83c9aba695 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/builder/IoTConsensusV2TransferBatchReqBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/builder/IoTConsensusV2TransferBatchReqBuilder.java @@ -40,7 +40,6 @@ import org.slf4j.LoggerFactory; import java.io.IOException; -import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -209,7 +208,6 @@ public List deepCopyEvents() { } protected int buildTabletInsertionBuffer(TabletInsertionEvent event) throws WALPipeException { - final ByteBuffer buffer; final TCommitId commitId; // event instanceof PipeInsertNodeTabletInsertionEvent) @@ -221,17 +219,16 @@ protected int buildTabletInsertionBuffer(TabletInsertionEvent event) throws WALP pipeInsertNodeTabletInsertionEvent.getCommitterKey().getRestartTimes(), pipeInsertNodeTabletInsertionEvent.getRebootTimes()); - // Read the bytebuffer from the wal file and transfer it directly without serializing or - // deserializing if possible final InsertNode insertNode = pipeInsertNodeTabletInsertionEvent.getInsertNode(); // IoTConsensusV2 will transfer binary data to TIoTConsensusV2TransferReq final ProgressIndex progressIndex = pipeInsertNodeTabletInsertionEvent.getProgressIndex(); - buffer = insertNode.serializeToByteBuffer(); + final int serializedSize = + IoTConsensusV2TabletInsertNodeReq.calculateSerializedSize(insertNode); batchReqs.add( IoTConsensusV2TabletInsertNodeReq.toTIoTConsensusV2TransferReq( insertNode, commitId, consensusGroupId, progressIndex, thisDataNodeId)); - return buffer.limit(); + return serializedSize; } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/request/IoTConsensusV2TabletInsertNodeReq.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/request/IoTConsensusV2TabletInsertNodeReq.java index 5f076b68ec38..af1d0a9fc5f3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/request/IoTConsensusV2TabletInsertNodeReq.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/iotconsensusv2/payload/request/IoTConsensusV2TabletInsertNodeReq.java @@ -95,6 +95,7 @@ public static IoTConsensusV2TabletInsertNodeReq toTIoTConsensusV2TransferReq( req.dataNodeId = thisDataNodeId; req.version = IoTConsensusV2RequestVersion.VERSION_1.getVersion(); req.type = IoTConsensusV2RequestType.TRANSFER_TABLET_INSERT_NODE.getType(); + // InsertNode preallocates this buffer with its manually calculated Pipe serialization size. req.body = insertNode.serializeToByteBuffer(); try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); @@ -109,6 +110,11 @@ public static IoTConsensusV2TabletInsertNodeReq toTIoTConsensusV2TransferReq( return req; } + /** Returns the exact serialized size of an InsertNode request body. */ + public static int calculateSerializedSize(final InsertNode insertNode) { + return insertNode.serializeToByteBufferSize(); + } + public static IoTConsensusV2TabletInsertNodeReq fromTIoTConsensusV2TransferReq( TIoTConsensusV2TransferReq transferReq) { final IoTConsensusV2TabletInsertNodeReq insertNodeReq = new IoTConsensusV2TabletInsertNodeReq(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/pipe/PipeEnrichedInsertNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/pipe/PipeEnrichedInsertNode.java index f6c323a1b575..2abff39a65ce 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/pipe/PipeEnrichedInsertNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/pipe/PipeEnrichedInsertNode.java @@ -39,6 +39,7 @@ import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.MeasurementSchema; import java.io.DataOutputStream; @@ -290,6 +291,16 @@ protected void serializeAttributes(final DataOutputStream stream) throws IOExcep insertNode.serialize(stream); } + @Override + protected int serializedAttributesSize() { + return PlanNodeType.BYTES + insertNode.serializeToByteBufferSize(); + } + + @Override + protected int serializedPlanNodeIdSize() { + return ReadWriteIOUtils.sizeToWrite(super.getPlanNodeId().getId()); + } + public static PipeEnrichedInsertNode deserialize(final ByteBuffer buffer) { return new PipeEnrichedInsertNode((InsertNode) PlanNodeType.deserialize(buffer)); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertMultiTabletsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertMultiTabletsNode.java index 4d1b987b8927..70f5435f6d78 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertMultiTabletsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertMultiTabletsNode.java @@ -280,6 +280,15 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { } } + @Override + protected int serializedAttributesSize() { + int size = PlanNodeType.BYTES + Integer.BYTES; + for (final InsertTabletNode insertTabletNode : insertTabletNodeList) { + size += insertTabletNode.baseSubSerializedSizeForPipe(); + } + return size + parentInsertTabletNodeIndexList.size() * Integer.BYTES; + } + @Override public void markAsGeneratedByPipe() { isGeneratedByPipe = true; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertNode.java index 9c8e36918835..6de154224f43 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertNode.java @@ -22,6 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.commons.consensus.index.ProgressIndex; import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.exception.runtime.SerializationRunTimeException; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; @@ -41,6 +42,7 @@ import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.NotImplementedException; import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.schema.MeasurementSchema; @@ -279,6 +281,31 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { DataNodeQueryMessages.SERIALIZEATTRIBUTES_OF_INSERTNODE_IS_NOT_IMPLEMENTED); } + /** Returns the exact size of the buffer produced by {@link #serializeToByteBuffer()}. */ + public final int serializeToByteBufferSize() { + // InsertNode has no children, so PlanNode.serialize only writes the child count here. + return serializedAttributesSize() + serializedPlanNodeIdSize() + Integer.BYTES; + } + + @Override + public ByteBuffer serializeToByteBuffer() { + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(serializeToByteBufferSize()); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + serialize(outputStream); + return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } catch (final IOException e) { + throw new SerializationRunTimeException(e); + } + } + + /** Returns the exact size of the attributes written by {@link #serializeToByteBuffer()}. */ + protected abstract int serializedAttributesSize(); + + /** Returns the exact size of the plan node id written after the attributes. */ + protected int serializedPlanNodeIdSize() { + return ReadWriteIOUtils.sizeToWrite(getPlanNodeId().getId()); + } + // region Serialization methods for WAL /** Serialized size of measurement schemas, ignoring failed time series */ diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java index 04768c57b502..67671865951b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java @@ -343,6 +343,75 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { subSerialize(stream); } + @Override + protected int serializedAttributesSize() { + return PlanNodeType.BYTES + pipeSubSerializedSize(); + } + + /** Returns the exact size of the row fields written during Pipe serialization. */ + protected int pipeSubSerializedSize() { + return Long.BYTES + + ReadWriteIOUtils.sizeToWrite(targetPath.getFullPath()) + + pipeMeasurementsAndValuesSerializedSize(); + } + + /** Returns the exact size of measurement and value fields written during Pipe serialization. */ + protected int pipeMeasurementsAndValuesSerializedSize() { + int size = Integer.BYTES + Byte.BYTES; + + for (int i = 0; measurements != null && i < measurements.length; i++) { + if (!shouldSerializeMeasurement(i)) { + continue; + } + size += + measurementSchemas == null + ? ReadWriteIOUtils.sizeToWrite(measurements[i]) + : measurementSchemas[i].serializedSize(); + } + + for (int i = 0; values != null && i < values.length; i++) { + if (!shouldSerializeMeasurement(i)) { + continue; + } + size += pipeValueSerializedSize(i); + } + + return size + Byte.BYTES + Byte.BYTES; + } + + private int pipeValueSerializedSize(final int index) { + final TSDataType dataType = getDataTypeIfPresent(index); + if (values[index] == null) { + return Byte.BYTES + (dataType == null ? 0 : Byte.BYTES); + } + + if (isNeedInferType) { + return Byte.BYTES + ReadWriteIOUtils.sizeToWrite(values[index].toString()); + } + + switch (dataType) { + case BOOLEAN: + return Byte.BYTES + Byte.BYTES; + case INT32: + case DATE: + return Byte.BYTES + Integer.BYTES; + case INT64: + case TIMESTAMP: + return Byte.BYTES + Long.BYTES; + case FLOAT: + return Byte.BYTES + Float.BYTES; + case DOUBLE: + return Byte.BYTES + Double.BYTES; + case TEXT: + case STRING: + case BLOB: + case OBJECT: + return Byte.BYTES + ReadWriteIOUtils.sizeToWrite((Binary) values[index]); + default: + throw new UnSupportedDataTypeException(UNSUPPORTED_DATA_TYPE + dataType); + } + } + void subSerialize(ByteBuffer buffer) { ReadWriteIOUtils.write(time, buffer); ReadWriteIOUtils.write(targetPath.getFullPath(), buffer); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java index 4492bf86acde..ec46e4104aaf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java @@ -275,6 +275,15 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { } } + @Override + protected int serializedAttributesSize() { + int size = PlanNodeType.BYTES + Integer.BYTES; + for (InsertRowNode node : insertRowNodeList) { + size += node.pipeSubSerializedSize(); + } + return size + insertRowNodeIndexList.size() * Integer.BYTES; + } + @Override public void markAsGeneratedByPipe() { isGeneratedByPipe = true; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java index ccc4ca810d84..17eb73a3e294 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java @@ -326,6 +326,16 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { } } + @Override + protected int serializedAttributesSize() { + int size = + PlanNodeType.BYTES + ReadWriteIOUtils.sizeToWrite(targetPath.getFullPath()) + Integer.BYTES; + for (InsertRowNode node : insertRowNodeList) { + size += Long.BYTES + node.pipeMeasurementsAndValuesSerializedSize(); + } + return size + insertRowNodeIndexList.size() * Integer.BYTES; + } + @Override public void markAsGeneratedByPipe() { isGeneratedByPipe = true; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java index 976223bf2cf1..c8a979f7e2f1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java @@ -567,6 +567,97 @@ void subSerialize(DataOutputStream stream) throws IOException { ReadWriteIOUtils.write((byte) (isAligned ? 1 : 0), stream); } + @Override + protected int serializedAttributesSize() { + return PlanNodeType.BYTES + baseSubSerializedSizeForPipe(); + } + + /** + * Returns the exact size written by {@link #subSerialize(DataOutputStream)}. + * + *

This deliberately excludes the plan-node type, id, and children. {@link + * InsertMultiTabletsNode} embeds tablet nodes by calling {@code subSerialize}, rather than their + * complete plan-node serialization. + */ + final int baseSubSerializedSizeForPipe() { + int size = ReadWriteIOUtils.sizeToWrite(targetPath.getFullPath()); + + size += Integer.BYTES; + size += Byte.BYTES; + for (int i = 0; measurements != null && i < measurements.length; i++) { + if (!shouldSerializeMeasurement(i)) { + continue; + } + size += + measurementSchemas == null + ? ReadWriteIOUtils.sizeToWrite(measurements[i]) + : measurementSchemas[i].serializedSize(); + } + + for (int i = 0; dataTypes != null && i < dataTypes.length; i++) { + if (shouldSerializeMeasurement(i)) { + size += TSDataType.getSerializedSize(); + } + } + + size += Integer.BYTES; + size += rowCount * Long.BYTES; + + size += Byte.BYTES; + if (bitMaps != null) { + for (int i = 0; measurements != null && i < measurements.length; i++) { + if (!shouldSerializeMeasurement(i)) { + continue; + } + size += Byte.BYTES; + if (getBitMapIfPresent(i) != null) { + size += BitMap.getSizeOfBytes(rowCount); + } + } + } + + for (int i = 0; columns != null && i < columns.length; i++) { + if (shouldSerializeMeasurement(i)) { + size += columnSerializedSizeForPipe(dataTypes[i], columns[i]); + } + } + + return size + Byte.BYTES; + } + + private int columnSerializedSizeForPipe(final TSDataType dataType, final Object column) { + switch (dataType) { + case INT32: + case DATE: + return rowCount * Integer.BYTES; + case INT64: + case TIMESTAMP: + return rowCount * Long.BYTES; + case FLOAT: + return rowCount * Float.BYTES; + case DOUBLE: + return rowCount * Double.BYTES; + case BOOLEAN: + return rowCount * Byte.BYTES; + case TEXT: + case BLOB: + case STRING: + case OBJECT: + int size = 0; + final Binary[] binaryValues = (Binary[]) column; + for (int i = 0; i < rowCount; i++) { + final Binary binary = binaryValues[i]; + size += + binary == null || binary.getValues() == null + ? Integer.BYTES + : Integer.BYTES + binary.getValues().length; + } + return size; + default: + throw new UnSupportedDataTypeException(String.format(DATATYPE_UNSUPPORTED, dataType)); + } + } + /** Serialize measurements or measurement schemas, ignoring failed time series */ private void writeMeasurementsOrSchemas(ByteBuffer buffer) { ReadWriteIOUtils.write(getValidMeasurementNumber(), buffer); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowNode.java index b11c0f6784e3..ae5abf83377c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowNode.java @@ -237,6 +237,11 @@ void subSerialize(DataOutputStream stream) throws IOException { } } + @Override + protected int pipeSubSerializedSize() { + return super.pipeSubSerializedSize() + getValidMeasurementNumber() * Byte.BYTES; + } + @Override protected void subSerialize(IWALByteBufferView buffer) { super.subSerialize(buffer); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java index d41b078cb9b8..d4c373e57d7f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java @@ -324,6 +324,17 @@ protected void serializeAttributes(DataOutputStream stream) throws IOException { } } + @Override + protected int serializedAttributesSize() { + int size = super.serializedAttributesSize(); + for (int i = 0; measurements != null && i < measurements.length; i++) { + if (shouldSerializeMeasurement(i)) { + size += Byte.BYTES; + } + } + return size; + } + @Override public void subDeserialize(ByteBuffer buffer) { super.subDeserialize(buffer); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferSerializationSizeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferSerializationSizeTest.java new file mode 100644 index 000000000000..ebaafe4fbcd2 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/request/PipeTransferSerializationSizeTest.java @@ -0,0 +1,458 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.iotdb.db.pipe.sink.payload.evolvable.request; + +import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex; +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory; +import org.apache.iotdb.db.pipe.sink.protocol.iotconsensusv2.payload.request.IoTConsensusV2TabletInsertNodeReq; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedInsertNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertMultiTabletsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsOfOneDeviceNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertTabletNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowsNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.utils.BitMap; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.Assert; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +public class PipeTransferSerializationSizeTest { + + @Test + public void testTabletRequestLengths() throws Exception { + final Tablet tablet = createTablet(); + final String database = "pipe_db"; + Assert.assertEquals( + PipeTransferTabletRawReq.calculateSerializedSize(tablet), + PipeTransferTabletRawReq.toTPipeTransferReq(tablet, false).getBody().length); + Assert.assertEquals( + PipeTransferTabletRawReqV2.calculateSerializedSize(tablet, database), + PipeTransferTabletRawReqV2.toTPipeTransferReq(tablet, false, database).getBody().length); + Assert.assertEquals( + PipeTransferTabletRawReq.calculateAirGapSerializedSize(tablet), + PipeTransferTabletRawReq.toTPipeTransferBytes(tablet, false).length); + } + + @Test + public void testBinaryRequestLengths() throws Exception { + final ByteBuffer payload = ByteBuffer.wrap(new byte[] {1, 2, 3, 4}); + final String database = "pipe_db"; + Assert.assertEquals( + PipeTransferTabletBinaryReqV2.calculateSerializedSize(payload, database), + PipeTransferTabletBinaryReqV2.toTPipeTransferReq(payload, database).getBody().length); + Assert.assertEquals( + PipeTransferTabletBinaryReqV2.calculateAirGapSerializedSize(payload, database), + PipeTransferTabletBinaryReqV2.toTPipeTransferBytes(payload, database).length); + Assert.assertEquals( + PipeTransferTabletBinaryReq.calculateSerializedSize(payload), + PipeTransferTabletBinaryReq.toTPipeTransferBytes(payload).length); + } + + @Test + public void testBatchRequestLengths() throws Exception { + final ByteBuffer insertNode = ByteBuffer.wrap(new byte[] {1, 2}); + final ByteBuffer tablet = ByteBuffer.wrap(new byte[] {3, 4, 5}); + Assert.assertEquals( + PipeTransferTabletBatchReq.calculateSerializedSize( + Collections.singletonList(insertNode), Collections.singletonList(tablet)), + PipeTransferTabletBatchReq.toTPipeTransferReq( + Collections.singletonList(insertNode), Collections.singletonList(tablet)) + .getBody() + .length); + + final String database = "db"; + Assert.assertEquals( + PipeTransferTabletBatchReqV2.calculateSerializedSize( + Collections.singletonList(insertNode), + Collections.singletonList(tablet), + Collections.singletonList(database), + Collections.singletonList(database)), + PipeTransferTabletBatchReqV2.toTPipeTransferReq( + Collections.singletonList(insertNode), + Collections.singletonList(tablet), + Collections.singletonList(database), + Collections.singletonList(database)) + .getBody() + .length); + } + + @Test + public void testInsertNodeSerializedSize() throws Exception { + assertInsertNodeRequestSizes(createInsertRowNode(0), "tree_db"); + assertInsertNodeRequestSizes(createInsertRowNodeWithSchemas(1), "tree_db"); + assertInsertNodeRequestSizes(createInsertRowNodeWithNullValue(), "tree_db"); + assertInsertNodeRequestSizes(createInsertRowNodeWithInferredType(), "tree_db"); + final InsertRowNode partiallyFailedRowNode = createInsertRowNodeWithSchemas(2); + partiallyFailedRowNode.markFailedMeasurement(1); + assertInsertNodeRequestSizes(partiallyFailedRowNode, "tree_db"); + + final InsertTabletNode tabletNode = createInsertTabletNode(); + assertInsertNodeRequestSizes(tabletNode, "tree_db"); + assertInsertNodeRequestSizes(createInsertTabletNode(false, true), "tree_db"); + final InsertTabletNode partiallyFailedTabletNode = createInsertTabletNode(false, true); + partiallyFailedTabletNode.markFailedMeasurement(1); + assertInsertNodeRequestSizes(partiallyFailedTabletNode, "tree_db"); + + final RelationalInsertRowNode relationalRowNode = + new RelationalInsertRowNode( + new PlanNodeId("relational-row"), + new PartialPath("table"), + false, + measurements(), + dataTypes(), + 1, + rowValues(1), + false, + columnCategories()); + assertInsertNodeRequestSizes(relationalRowNode, "table_db_\u6d4b\u8bd5"); + + final RelationalInsertTabletNode relationalTabletNode = createRelationalInsertTabletNode(); + assertInsertNodeRequestSizes(relationalTabletNode, "table_db"); + + final List rows = new ArrayList<>(); + final List relationalRows = new ArrayList<>(); + for (int row = 0; row < 50; row++) { + rows.add(createInsertRowNode(row)); + relationalRows.add(createRelationalInsertRowNode(row)); + } + final InsertRowsNode insertRowsNode = new InsertRowsNode(new PlanNodeId("rows")); + insertRowsNode.setInsertRowNodeList(rows); + insertRowsNode.setInsertRowNodeIndexList(indexes(rows.size())); + assertInsertNodeRequestSizes(insertRowsNode, "tree_db"); + + final InsertRowsOfOneDeviceNode oneDeviceNode = + new InsertRowsOfOneDeviceNode(new PlanNodeId("one-device")); + oneDeviceNode.setInsertRowNodeList(rows); + oneDeviceNode.setInsertRowNodeIndexList(indexes(rows.size())); + assertInsertNodeRequestSizes(oneDeviceNode, "tree_db"); + + final InsertMultiTabletsNode multiTabletsNode = + new InsertMultiTabletsNode(new PlanNodeId("multi-tablets")); + multiTabletsNode.addInsertTabletNode(tabletNode, 0); + multiTabletsNode.addInsertTabletNode(relationalTabletNode, 1); + assertInsertNodeRequestSizes(multiTabletsNode, "tree_db"); + + final RelationalInsertRowsNode relationalRowsNode = + new RelationalInsertRowsNode( + new PlanNodeId("relational-rows"), indexes(relationalRows.size()), relationalRows); + assertInsertNodeRequestSizes(relationalRowsNode, "table_db"); + + final PipeEnrichedInsertNode pipeEnrichedInsertNode = + new PipeEnrichedInsertNode(createInsertRowNode(2)); + pipeEnrichedInsertNode.setPlanNodeId(new PlanNodeId("enriched-row")); + assertInsertNodeRequestSizes(pipeEnrichedInsertNode, "tree_db"); + } + + private static void assertInsertNodeRequestSizes( + final InsertNode insertNode, final String databaseName) throws Exception { + final ByteBuffer serializedInsertNode = insertNode.serializeToByteBuffer(); + Assert.assertEquals(insertNode.serializeToByteBufferSize(), serializedInsertNode.capacity()); + Assert.assertEquals(insertNode.serializeToByteBufferSize(), serializedInsertNode.remaining()); + Assert.assertEquals( + PipeTransferTabletInsertNodeReq.calculateSerializedSize(insertNode), + PipeTransferTabletInsertNodeReq.toTPipeTransferReq(insertNode).getBody().length); + Assert.assertEquals( + PipeTransferTabletInsertNodeReq.calculateAirGapSerializedSize(insertNode), + PipeTransferTabletInsertNodeReq.toTPipeTransferBytes(insertNode).length); + Assert.assertEquals( + PipeTransferTabletInsertNodeReqV2.calculateSerializedSize(insertNode, databaseName), + PipeTransferTabletInsertNodeReqV2.toTPipeTransferReq(insertNode, databaseName) + .getBody() + .length); + Assert.assertEquals( + PipeTransferTabletInsertNodeReqV2.calculateAirGapSerializedSize(insertNode, databaseName), + PipeTransferTabletInsertNodeReqV2.toTPipeTransferBytes(insertNode, databaseName).length); + Assert.assertEquals( + IoTConsensusV2TabletInsertNodeReq.calculateSerializedSize(insertNode), + IoTConsensusV2TabletInsertNodeReq.toTIoTConsensusV2TransferReq( + insertNode, null, null, MinimumProgressIndex.INSTANCE, 0) + .getBody() + .length); + } + + private static List indexes(final int size) { + final List indexes = new ArrayList<>(size); + for (int i = 0; i < size; i++) { + indexes.add(i); + } + return indexes; + } + + private static InsertRowNode createInsertRowNode(final int row) throws IllegalPathException { + return new InsertRowNode( + new PlanNodeId("row-" + row), + new PartialPath("root.sg.d"), + false, + measurements(), + dataTypes(), + row, + rowValues(row), + false); + } + + private static InsertRowNode createInsertRowNodeWithSchemas(final int row) + throws IllegalPathException { + return new InsertRowNode( + new PlanNodeId("row-with-schemas"), + new PartialPath("root.sg.d"), + false, + measurements(), + dataTypes(), + measurementSchemas(), + row, + rowValues(row), + false); + } + + private static InsertRowNode createInsertRowNodeWithNullValue() throws IllegalPathException { + return new InsertRowNode( + new PlanNodeId("row-with-null"), + new PartialPath("root.sg.d"), + false, + new String[] {"s"}, + new TSDataType[] {TSDataType.INT32}, + 1, + new Object[] {null}, + false); + } + + private static InsertRowNode createInsertRowNodeWithInferredType() throws IllegalPathException { + return new InsertRowNode( + new PlanNodeId("row-with-inferred-type"), + new PartialPath("root.sg.d"), + false, + new String[] {"s"}, + new TSDataType[] {null}, + 1, + new Object[] {"value"}, + true); + } + + private static RelationalInsertRowNode createRelationalInsertRowNode(final int row) + throws IllegalPathException { + return new RelationalInsertRowNode( + new PlanNodeId("relational-row-" + row), + new PartialPath("table"), + false, + measurements(), + dataTypes(), + row, + rowValues(row), + false, + columnCategories()); + } + + private static InsertTabletNode createInsertTabletNode() throws IllegalPathException { + return createInsertTabletNode(false, false); + } + + private static InsertTabletNode createInsertTabletNode(final boolean relational) + throws IllegalPathException { + return createInsertTabletNode(relational, false); + } + + private static InsertTabletNode createInsertTabletNode( + final boolean relational, final boolean withBitMaps) throws IllegalPathException { + final String[] measurements = measurements(); + final TSDataType[] types = dataTypes(); + final MeasurementSchema[] schemas = measurementSchemas(); + final Object[] columns = new Object[types.length]; + final int rowCount = 50; + final long[] times = new long[rowCount]; + for (int i = 0; i < rowCount; i++) { + times[i] = i; + } + for (int column = 0; column < types.length; column++) { + switch (types[column]) { + case BOOLEAN: + final boolean[] booleanValues = new boolean[rowCount]; + for (int row = 0; row < rowCount; row++) { + booleanValues[row] = row % 2 == 0; + } + columns[column] = booleanValues; + break; + case INT32: + case DATE: + final int[] intValues = new int[rowCount]; + for (int row = 0; row < rowCount; row++) { + intValues[row] = row; + } + columns[column] = intValues; + break; + case INT64: + case TIMESTAMP: + final long[] longValues = new long[rowCount]; + for (int row = 0; row < rowCount; row++) { + longValues[row] = row; + } + columns[column] = longValues; + break; + case FLOAT: + final float[] floatValues = new float[rowCount]; + for (int row = 0; row < rowCount; row++) { + floatValues[row] = row; + } + columns[column] = floatValues; + break; + case DOUBLE: + final double[] doubleValues = new double[rowCount]; + for (int row = 0; row < rowCount; row++) { + doubleValues[row] = row; + } + columns[column] = doubleValues; + break; + case TEXT: + case BLOB: + case STRING: + case OBJECT: + Binary[] values = new Binary[rowCount]; + for (int row = 1; row < rowCount; row++) { + values[row] = new Binary(("value-" + row).getBytes(StandardCharsets.UTF_8)); + } + columns[column] = values; + break; + default: + throw new AssertionError(types[column]); + } + } + final BitMap[] bitMaps = withBitMaps ? createBitMaps(types.length, rowCount) : null; + return relational + ? new RelationalInsertTabletNode( + new PlanNodeId("relational-tablet"), + new PartialPath("table"), + false, + measurements, + types, + schemas, + times, + bitMaps, + columns, + rowCount, + columnCategories()) + : new InsertTabletNode( + new PlanNodeId("tablet"), + new PartialPath("root.sg.d"), + false, + measurements, + types, + schemas, + times, + bitMaps, + columns, + rowCount); + } + + private static RelationalInsertTabletNode createRelationalInsertTabletNode() + throws IllegalPathException { + return (RelationalInsertTabletNode) createInsertTabletNode(true); + } + + private static String[] measurements() { + return new String[] {"b", "i", "l", "f", "d", "t", "ts", "date", "blob", "string", "object"}; + } + + private static TSDataType[] dataTypes() { + return new TSDataType[] { + TSDataType.BOOLEAN, + TSDataType.INT32, + TSDataType.INT64, + TSDataType.FLOAT, + TSDataType.DOUBLE, + TSDataType.TEXT, + TSDataType.TIMESTAMP, + TSDataType.DATE, + TSDataType.BLOB, + TSDataType.STRING, + TSDataType.OBJECT + }; + } + + private static MeasurementSchema[] measurementSchemas() { + final String[] measurements = measurements(); + final TSDataType[] types = dataTypes(); + final MeasurementSchema[] schemas = new MeasurementSchema[types.length]; + for (int i = 0; i < types.length; i++) { + schemas[i] = new MeasurementSchema(measurements[i], types[i], TSEncoding.PLAIN); + } + return schemas; + } + + private static BitMap[] createBitMaps(final int columnCount, final int rowCount) { + final BitMap[] bitMaps = new BitMap[columnCount]; + bitMaps[0] = new BitMap(rowCount); + bitMaps[0].mark(0); + bitMaps[columnCount - 1] = new BitMap(rowCount); + bitMaps[columnCount - 1].mark(rowCount - 1); + return bitMaps; + } + + private static Object[] rowValues(final int row) { + return new Object[] { + true, + row, + (long) row, + (float) row, + (double) row, + new Binary(("text-" + row).getBytes(StandardCharsets.UTF_8)), + (long) row, + row, + new Binary(("blob-" + row).getBytes(StandardCharsets.UTF_8)), + new Binary(("string-" + row).getBytes(StandardCharsets.UTF_8)), + new Binary(("object-" + row).getBytes(StandardCharsets.UTF_8)) + }; + } + + private static TsTableColumnCategory[] columnCategories() { + final TsTableColumnCategory[] categories = new TsTableColumnCategory[dataTypes().length]; + Arrays.fill(categories, TsTableColumnCategory.FIELD); + categories[0] = TsTableColumnCategory.TAG; + return categories; + } + + private static Tablet createTablet() { + final Tablet tablet = + new Tablet( + "table1", Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT32)), 1); + tablet.setColumnCategories(Collections.singletonList(ColumnCategory.FIELD)); + tablet.addTimestamp(0, 1L); + tablet.addValue(0, 0, 1); + tablet.setRowSize(1); + return tablet; + } +}