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 a2beae65cccb484841328f60f1a50033367f76ae Author: Weihao Li <[email protected]> AuthorDate: Tue Sep 15 16:51:39 2026 +0800 modify some Signed-off-by: Weihao Li <[email protected]> --- .../execution/exchange/MPPDataExchangeManager.java | 5 ++- .../execution/exchange/sink/SinkChannel.java | 3 +- .../execution/exchange/source/SourceHandle.java | 39 +++++++--------------- .../execution/exchange/SinkChannelTest.java | 16 ++++++++- .../db/queryengine/execution/exchange/Utils.java | 7 ++++ .../src/main/thrift/datanode.thrift | 2 ++ 6 files changed, 42 insertions(+), 30 deletions(-) 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 e60a7b3d89f..284aa1976ce 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 @@ -183,7 +183,7 @@ public class MPPDataExchangeManager implements IMPPDataExchangeManager { int offset = req.getOffset(); for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { - ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i).asReadOnlyBuffer(); + ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); int blockOffset = i == req.getStartSequenceId() ? offset : 0; int serializedTsBlockSize = serializedTsBlock.remaining(); if (blockOffset < 0 || blockOffset > serializedTsBlockSize) { @@ -209,6 +209,9 @@ public class MPPDataExchangeManager implements IMPPDataExchangeManager { fragment.limit(blockOffset + remainingPayloadSize); resp.addToTsBlocks(fragment.slice()); resp.setOffset(blockOffset + remainingPayloadSize); + if (blockOffset == 0) { + resp.setTotalLength(serializedTsBlockSize); + } break; } } catch (GetTsBlockFromClosedOrAbortedChannelException e) { 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 86c66b6677e..c98be7d55bd 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 @@ -427,7 +427,8 @@ public class SinkChannel implements ISinkChannel { DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); } serializedTsBlock = serde.serialize(pair.left); - sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock.asReadOnlyBuffer()); + sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock); + pair.left = null; return serializedTsBlock.duplicate(); } 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 3afc800c2af..780c523b007 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 @@ -47,11 +47,11 @@ import org.apache.tsfile.external.commons.lang3.Validate; import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.PublicBAOS; 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; @@ -60,6 +60,7 @@ import java.util.Map; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; import static org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.createFullIdFrom; import static org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet.GET_DATA_BLOCK_TASK_CALLER; @@ -794,7 +795,7 @@ public class SourceHandle implements ISourceHandle { private final List<ByteBuffer> tsBlocks; private int nextSequenceId; private int offset; - private ByteArrayOutputStream partialTsBlock; + private PublicBAOS partialTsBlock; private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { this.nextSequenceId = startSequenceId; @@ -806,7 +807,7 @@ public class SourceHandle implements ISourceHandle { return nextSequenceId == endSequenceId && partialTsBlock == null; } - private void addResponse(TGetDataBlockResponse response) throws TException { + private void addResponse(TGetDataBlockResponse response) { List<ByteBuffer> responseBlocks = response.getTsBlocks(); boolean lastBlockIsFragment = response.isSetOffset(); int blockIndex = 0; @@ -817,7 +818,7 @@ public class SourceHandle implements ISourceHandle { updateOffset(response.getOffset()); return; } - tsBlocks.add(ByteBuffer.wrap(partialTsBlock.toByteArray())); + tsBlocks.add(ByteBuffer.wrap(partialTsBlock.getBuf())); partialTsBlock = null; offset = 0; nextSequenceId++; @@ -831,36 +832,20 @@ public class SourceHandle implements ISourceHandle { } if (lastBlockIsFragment) { - partialTsBlock = new ByteArrayOutputStream(); + checkArgument(response.isSetTotalLength(), "xxx"); + partialTsBlock = new PublicBAOS(response.getTotalLength()); 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 appendFragment(ByteBuffer fragment) { + checkArgument(fragment.hasRemaining(), "xxx"); + partialTsBlock.writeBytes(fragment.array()); } - private void updateOffset(int nextOffset) throws TException { - if (nextOffset <= offset || nextOffset != partialTsBlock.size()) { - throw new TException( - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); - } + private void updateOffset(int nextOffset) { + checkArgument(nextOffset > offset && nextOffset == partialTsBlock.size(), "xxx"); offset = nextOffset; } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java index f59dd9925aa..f22725c0fea 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java @@ -37,6 +37,7 @@ import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent; import org.apache.thrift.TException; import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; @@ -87,6 +88,7 @@ public class SinkChannelTest { SinkListener mockSinkListener = Mockito.mock(SinkListener.class); // Construct several mock TsBlock(s). List<TsBlock> mockTsBlocks = Utils.createMockTsBlocks(numOfMockTsBlock, mockTsBlockSize); + TsBlockSerde mockTsBlockSerde = Utils.createMockTsBlockSerde(mockTsBlockSize); // Construct SinkChannel. SinkChannel sinkChannel = @@ -98,7 +100,7 @@ public class SinkChannelTest { localFragmentInstanceId, mockLocalMemoryManager, Executors.newSingleThreadExecutor(), - Utils.createMockTsBlockSerde(mockTsBlockSize), + mockTsBlockSerde, mockSinkListener, mockClientManager); sinkChannel.open(); @@ -145,6 +147,8 @@ public class SinkChannelTest { for (int i = 0; i < numOfMockTsBlock; i++) { try { sinkChannel.getSerializedTsBlock(i); + sinkChannel.getSerializedTsBlock(i); + Mockito.verify(mockTsBlockSerde, Mockito.times(1)).serialize(mockTsBlocks.get(i)); } catch (IOException e) { e.printStackTrace(); Assert.fail(); @@ -166,6 +170,16 @@ public class SinkChannelTest { Assert.assertTrue(sinkChannel.isFinished()); Assert.assertFalse(sinkChannel.isAborted()); Assert.assertEquals(mockTsBlockSize, sinkChannel.getBufferRetainedSizeInBytes()); + for (int i = 0; i < numOfMockTsBlock; i++) { + try { + sinkChannel.getSerializedTsBlock(i); + Assert.fail("The acknowledged serialized TsBlock should have been released"); + } catch (IllegalStateException expected) { + // Both the original entry and serialized cache must be removed by acknowledgement. + } catch (IOException e) { + Assert.fail(e.getMessage()); + } + } Mockito.verify(mockMemoryPool, Mockito.timeout(10_0000).times(1)) .free( queryId, diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java index b09498ad949..1a701b9eded 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java @@ -28,6 +28,7 @@ import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.mockito.Mockito; import org.mockito.stubbing.Answer; +import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; @@ -155,6 +156,12 @@ public class Utils { TsBlock mockTsBlock = Mockito.mock(TsBlock.class); Mockito.when(mockTsBlock.getRetainedSizeInBytes()).thenReturn(mockTsBlockSize); Mockito.when(mockTsBlock.getSizeInBytes()).thenReturn(mockTsBlockSize); + try { + Mockito.when(mockTsBlockSerde.serialize(Mockito.any(TsBlock.class))) + .thenReturn(ByteBuffer.allocate(Math.toIntExact(mockTsBlockSize))); + } catch (IOException e) { + throw new AssertionError(e); + } Mockito.when(mockTsBlockSerde.deserialize(Mockito.any(ByteBuffer.class))) .thenReturn(mockTsBlock); return mockTsBlockSerde; diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index cf9aca63b17..ceb794989ce 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -98,6 +98,8 @@ 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 i32 offset + // Total serialized length of the TsBlock when the response starts its first fragment. + 3: optional i32 totalLength } struct TAcknowledgeDataBlockEvent {
