This is an automated email from the ASF dual-hosted git repository. Wei-hao-Li pushed a commit to branch mppEx in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 280c432188bf740ad779e020ee7cd2a6a2d9a86d Author: Weihao Li <[email protected]> AuthorDate: Mon Sep 14 14:36:43 2026 +0800 draft Signed-off-by: Weihao Li <[email protected]> --- .../java/org/apache/iotdb/db/conf/IoTDBConfig.java | 12 +++ .../org/apache/iotdb/db/conf/IoTDBDescriptor.java | 5 ++ .../execution/exchange/MPPDataExchangeManager.java | 35 ++++++++ .../execution/exchange/sink/SinkChannel.java | 32 ++++++- .../execution/exchange/source/SourceHandle.java | 98 +++++++++++++++++++++- .../conf/iotdb-system.properties.template | 5 ++ .../src/main/thrift/datanode.thrift | 4 + 7 files changed, 189 insertions(+), 2 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index 989031fc0fd..7d7f3652d38 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -956,6 +956,8 @@ public class IoTDBConfig { /** Core pool size of mpp data exchange. */ private int mppDataExchangeCorePoolSize = 10; + private int mppDataExchangeMaxPayloadSizeInBytes = 8 * 1024 * 1024; + /** Max pool size of mpp data exchange. */ private int mppDataExchangeMaxPoolSize = 10; @@ -3421,6 +3423,16 @@ public class IoTDBConfig { this.mppDataExchangeKeepAliveTimeInMs = mppDataExchangeKeepAliveTimeInMs; } + public int getMppDataExchangeMaxPayloadSizeInBytes() { + return mppDataExchangeMaxPayloadSizeInBytes; + } + + public void setMppDataExchangeMaxPayloadSizeInBytes( + int mppDataExchangeMaxPayloadSizeInBytes) { + this.mppDataExchangeMaxPayloadSizeInBytes = + Math.max(1, Math.min(mppDataExchangeMaxPayloadSizeInBytes, thriftMaxFrameSize - 1024)); + } + public int getConnectionTimeoutInMS() { return connectionTimeoutInMS; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 9f48bac0e0c..8c3efb84b83 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -3070,6 +3070,11 @@ public class IoTDBDescriptor { properties.getProperty( "mpp_data_exchange_keep_alive_time_in_ms", Integer.toString(conf.getMppDataExchangeKeepAliveTimeInMs())))); + conf.setMppDataExchangeMaxPayloadSizeInBytes( + Integer.parseInt( + properties.getProperty( + "mpp_data_exchange_max_payload_size_in_bytes", + Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes())))); conf.setPartitionCacheSize( Integer.parseInt( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index ceba3880122..f8327673a80 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -24,6 +24,7 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient; import org.apache.iotdb.commons.utils.TestOnly; +import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.exception.exchange.GetTsBlockFromClosedOrAbortedChannelException; import org.apache.iotdb.db.queryengine.execution.driver.DriverContext; @@ -176,6 +177,40 @@ public class MPPDataExchangeManager implements IMPPDataExchangeManager { } // index of the channel must be a SinkChannel SinkChannel sinkChannel = (SinkChannel) (sinkHandle.getChannel(req.getIndex())); + if (req.isSetOffset()) { + int remainingPayloadSize = + IoTDBDescriptor.getInstance() + .getConfig() + .getMppDataExchangeMaxPayloadSizeInBytes(); + long offset = req.getOffset(); + for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { + try { + ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); + long blockOffset = i == req.getStartSequenceId() ? offset : 0L; + long remainingBlockSize = serializedTsBlock.remaining() - blockOffset; + if (remainingBlockSize <= remainingPayloadSize) { + resp.addToTsBlocks( + sinkChannel.getSerializedTsBlockFragment( + i, blockOffset, Math.toIntExact(remainingBlockSize))); + remainingPayloadSize -= Math.toIntExact(remainingBlockSize); + if (remainingPayloadSize == 0) { + break; + } + } else { + resp.addToTsBlocks( + sinkChannel.getSerializedTsBlockFragment( + i, blockOffset, remainingPayloadSize)); + resp.setOffset(blockOffset + remainingPayloadSize); + break; + } + } catch (GetTsBlockFromClosedOrAbortedChannelException e) { + return new TGetDataBlockResponse(new ArrayList<>()); + } catch (IllegalArgumentException | IllegalStateException | IOException e) { + throw new TException(e); + } + } + return resp; + } for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index daaff9fdbcd..533387afdf0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -96,6 +96,10 @@ public class SinkChannel implements ISinkChannel { private final LinkedHashMap<Integer, Pair<TsBlock, Long>> sequenceIdToTsBlock = new LinkedHashMap<>(); + /** Serialized blocks are cached so fragmented requests do not serialize the same block again. */ + private final LinkedHashMap<Integer, ByteBuffer> sequenceIdToSerializedTsBlock = + new LinkedHashMap<>(); + // size for current TsBlock to reserve and free private long currentTsBlockSize; @@ -305,6 +309,7 @@ public class SinkChannel implements ISinkChannel { return false; } sequenceIdToTsBlock.clear(); + sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -335,6 +340,7 @@ public class SinkChannel implements ISinkChannel { return false; } sequenceIdToTsBlock.clear(); + sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -407,6 +413,10 @@ public class SinkChannel implements ISinkChannel { throw new GetTsBlockFromClosedOrAbortedChannelException( DataNodeQueryMessages.SINKCHANNEL_IS_ABORTED_OR_CLOSED); } + ByteBuffer serializedTsBlock = sequenceIdToSerializedTsBlock.get(sequenceId); + if (serializedTsBlock != null) { + return serializedTsBlock.duplicate(); + } Pair<TsBlock, Long> pair = sequenceIdToTsBlock.get(sequenceId); if (pair == null || pair.left == null) { LOGGER.warn( @@ -416,7 +426,26 @@ public class SinkChannel implements ISinkChannel { throw new IllegalStateException( DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); } - return serde.serialize(pair.left); + serializedTsBlock = serde.serialize(pair.left); + sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock.asReadOnlyBuffer()); + return serializedTsBlock.duplicate(); + } + + public synchronized ByteBuffer getSerializedTsBlockFragment( + int sequenceId, long offset, int maxBytes) throws IOException { + ByteBuffer serializedTsBlock = getSerializedTsBlock(sequenceId); + if (offset < 0 || offset > serializedTsBlock.remaining() || maxBytes <= 0) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_ARG_ARG_2946DBE5, + "serialized TsBlock", + "fragment range")); + } + int length = (int) Math.min(maxBytes, serializedTsBlock.remaining() - offset); + ByteBuffer fragment = serializedTsBlock.duplicate(); + fragment.position(Math.toIntExact(offset)); + fragment.limit(Math.toIntExact(offset + length)); + return fragment.slice(); } public void acknowledgeTsBlock(int startSequenceId, int endSequenceId) { @@ -439,6 +468,7 @@ public class SinkChannel implements ISinkChannel { freedBytes += entry.getValue().right; bufferRetainedSizeInBytes -= entry.getValue().right; iterator.remove(); + sequenceIdToSerializedTsBlock.remove(entry.getKey()); if (LOGGER.isDebugEnabled()) { LOGGER.debug(DataNodeQueryMessages.ACK_TSBLOCK, entry.getKey()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 67a5defb09a..1deba48e8ad 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -51,6 +51,7 @@ import org.apache.tsfile.utils.RamUsageEstimator; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.ByteArrayOutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; @@ -640,6 +641,8 @@ public class SourceHandle implements ISourceHandle { startSequenceId, endSequenceId, indexOfUpstreamSinkHandle); + DataBlockFetchProgress fetchProgress = + new DataBlockFetchProgress(startSequenceId, endSequenceId); int attempt = 0; while (attempt < MAX_ATTEMPT_TIMES) { attempt += 1; @@ -648,7 +651,8 @@ public class SourceHandle implements ISourceHandle { boolean transferAttemptRecorded = false; try (SyncDataNodeMPPDataExchangeServiceClient client = mppDataExchangeServiceClientManager.borrowClient(remoteEndpoint)) { - TGetDataBlockResponse resp = client.getDataBlock(req); + TGetDataBlockResponse resp = + getDataBlockWithFragments(client, req, fetchProgress); int tsBlockNum = resp.getTsBlocks().size(); if (tsBlockNum != endSequenceId - startSequenceId) { recordTransferAttempt( @@ -767,6 +771,98 @@ public class SourceHandle implements ISourceHandle { sourceHandleListener.onFailure(SourceHandle.this, t); } } + + private TGetDataBlockResponse getDataBlockWithFragments( + SyncDataNodeMPPDataExchangeServiceClient client, + TGetDataBlockRequest request, + DataBlockFetchProgress fetchProgress) + throws TException { + while (!fetchProgress.isFinished()) { + TGetDataBlockRequest fragmentRequest = request.deepCopy(); + fragmentRequest.setStartSequenceId(fetchProgress.nextSequenceId); + fragmentRequest.setOffset(fetchProgress.offset); + TGetDataBlockResponse response = client.getDataBlock(fragmentRequest); + if (response.getTsBlocks().isEmpty()) { + return response; + } + fetchProgress.addResponse(response); + } + return new TGetDataBlockResponse(fetchProgress.tsBlocks); + } + + private class DataBlockFetchProgress { + private final int endSequenceId; + private final List<ByteBuffer> tsBlocks; + private int nextSequenceId; + private long offset; + private ByteArrayOutputStream partialTsBlock; + + private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { + this.nextSequenceId = startSequenceId; + this.endSequenceId = endSequenceId; + this.tsBlocks = new ArrayList<>(endSequenceId - startSequenceId); + } + + private boolean isFinished() { + return nextSequenceId == endSequenceId && partialTsBlock == null; + } + + private void addResponse(TGetDataBlockResponse response) throws TException { + List<ByteBuffer> responseBlocks = response.getTsBlocks(); + boolean lastBlockIsFragment = response.isSetOffset(); + int blockIndex = 0; + + if (partialTsBlock != null) { + appendFragment(responseBlocks.get(blockIndex++)); + if (lastBlockIsFragment && blockIndex == responseBlocks.size()) { + updateOffset(response.getOffset()); + return; + } + tsBlocks.add(ByteBuffer.wrap(partialTsBlock.toByteArray())); + partialTsBlock = null; + offset = 0; + nextSequenceId++; + } + + int lastCompleteBlockIndex = + lastBlockIsFragment ? responseBlocks.size() - 1 : responseBlocks.size(); + while (blockIndex < lastCompleteBlockIndex) { + tsBlocks.add(responseBlocks.get(blockIndex++)); + nextSequenceId++; + } + + if (lastBlockIsFragment) { + partialTsBlock = new ByteArrayOutputStream(); + appendFragment(responseBlocks.get(blockIndex)); + updateOffset(response.getOffset()); + } + + if (nextSequenceId > endSequenceId + || (!lastBlockIsFragment && nextSequenceId == endSequenceId && partialTsBlock != null)) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + } + + private void appendFragment(ByteBuffer fragment) throws TException { + ByteBuffer duplicate = fragment.duplicate(); + if (!duplicate.hasRemaining()) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + byte[] bytes = new byte[duplicate.remaining()]; + duplicate.get(bytes); + partialTsBlock.writeBytes(bytes); + } + + private void updateOffset(long nextOffset) throws TException { + if (nextOffset <= offset || nextOffset != partialTsBlock.size()) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + offset = nextOffset; + } + } } class SendAcknowledgeDataBlockEventTask implements Runnable { diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index c53c95fd016..f7665c30098 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -1194,6 +1194,11 @@ mpp_data_exchange_max_pool_size=10 # Datatype: int mpp_data_exchange_keep_alive_time_in_ms=1000 +# The maximum payload size of one MPP data exchange RPC response +# effectiveMode: restart +# Datatype: int, Unit: byte +mpp_data_exchange_max_payload_size_in_bytes=8388608 + # The max execution time of a DriverTask # effectiveMode: restart # Datatype: int, Unit: ms diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index c5c2bdac4a6..77dacf07a6b 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -90,10 +90,14 @@ struct TGetDataBlockRequest { 3: required i32 endSequenceId // Index of upstream SinkChannel 4: required i32 index + // Optional byte range for fetching one serialized TsBlock in fragments. + 5: optional i64 offset } struct TGetDataBlockResponse { 1: required list<binary> tsBlocks + // The start offset of the next fragment. It is set only when the last element in tsBlocks is a fragment. + 2: optional i64 offset } struct TAcknowledgeDataBlockEvent {
