Load: Fix LoadPieceNode reading aligned Chunk causing OOM (#16132)

* Load: Fix LoadPieceNode reading aligned Chunk causing OOM

* update alignedChunkData

* add ByteBufferInputStream

Author: luoluoyuyu

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

@@ -34,7 +34,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.File;
@@ -152,7 +151,9 @@ public List<WritePlanNode> splitByPartition(IAnalysis analysis) {
   }
 
   public static PlanNode deserialize(ByteBuffer buffer) {
-    InputStream stream = new ByteArrayInputStream(buffer.array());
+    buffer = buffer.duplicate();
+    buffer.position(0);
+    ByteBufferInputStream stream = new ByteBufferInputStream(buffer);
     try {
       ReadWriteIOUtils.readShort(stream); // read PlanNodeType
       final File tsFile = new File(ReadWriteIOUtils.readString(stream));
@@ -193,4 +194,41 @@ public int hashCode() {
   public String toString() {
     return "LoadTsFilePieceNode{" + "tsFile=" + tsFile + ", dataSize=" + dataSize + '}';
   }
+
+  public static class ByteBufferInputStream extends InputStream {
+    private final ByteBuffer buffer;
+
+    public ByteBufferInputStream(ByteBuffer buffer) {
+      this.buffer = buffer;
+    }
+
+    @Override
+    public int read() {
+      if (!buffer.hasRemaining()) {
+        return -1;
+      }
+      return buffer.get() & 0xFF;
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) {
+      if (!buffer.hasRemaining()) {
+        return -1;
+      }
+      int toRead = Math.min(len, buffer.remaining());
+      buffer.get(b, off, toRead);
+      return toRead;
+    }
+
+    public ByteBuffer read(int length) {
+      if (length < 0 || length > buffer.remaining()) {
+        throw new IllegalArgumentException("Invalid length for slicing: " + length);
+      }
+      ByteBuffer slicedBuffer = buffer.slice();
+      slicedBuffer.limit(length);
+
+      buffer.position(buffer.position() + length);
+      return slicedBuffer;
+    }
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java

@@ -81,6 +81,7 @@
 import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.apache.commons.io.FileUtils;
+import org.apache.tsfile.exception.write.PageException;
 import org.apache.tsfile.utils.FilePathUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -947,7 +948,7 @@ public TSStatus writeLoadTsFileNode(
 
     try {
       loadTsFileManager.writeToDataRegion(dataRegion, pieceNode, uuid);
-    } catch (IOException e) {
+    } catch (IOException | PageException e) {
       LOGGER.warn(
           "IO error when writing piece node of TsFile {} to DataRegion {}.",
           pieceNode.getTsFile(),

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java

@@ -56,6 +56,7 @@
 
 import org.apache.tsfile.common.constant.TsFileConstant;
 import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.exception.write.PageException;
 import org.apache.tsfile.file.metadata.ChunkGroupMetadata;
 import org.apache.tsfile.file.metadata.ChunkMetadata;
 import org.apache.tsfile.file.metadata.IDeviceID;
@@ -203,7 +204,7 @@ private void recover() {
   }
 
   public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNode, String uuid)
-      throws IOException {
+      throws IOException, PageException {
     if (!uuid2WriterManager.containsKey(uuid)) {
       synchronized (uuid2CleanupTask) {
         final CleanupTask cleanupTask =
@@ -414,7 +415,8 @@ private void clearDir(File dir) {
      * BatchedAlignedChunkData, it may result in no data for the time column in the new file.
      */
     @SuppressWarnings("squid:S3824")
-    private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) throws IOException {
+    private void write(DataPartitionInfo partitionInfo, ChunkData chunkData)
+        throws IOException, PageException {
       if (isClosed) {
         throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir));
       }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.utils.TimePartitionUtils;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode;
 
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.exception.write.PageException;
@@ -66,16 +67,15 @@ public class AlignedChunkData implements ChunkData {
   protected final IDeviceID device;
   protected List<ChunkHeader> chunkHeaderList;
 
-  protected final PublicBAOS byteStream;
-  protected final DataOutputStream stream;
+  protected PublicBAOS byteStream;
+  protected DataOutputStream stream;
   protected List<long[]> timeBatch;
   protected long dataSize;
   protected boolean needDecodeChunk;
   protected List<Integer> pageNumbers;
   protected Queue<Integer> satisfiedLengthQueue;
 
-  private AlignedChunkWriterImpl chunkWriter;
-  protected List<Chunk> chunkList;
+  protected ByteBuffer chunkData;
 
   public AlignedChunkData(
       @Nonnull final IDeviceID device,
@@ -143,14 +143,8 @@ public boolean isAligned() {
   }
 
   @Override
-  public void writeToFileWriter(final TsFileIOWriter writer) throws IOException {
-    if (chunkList != null) {
-      for (final Chunk chunk : chunkList) {
-        writer.writeChunk(chunk);
-      }
-    } else {
-      chunkWriter.writeToFileWriter(writer);
-    }
+  public void writeToFileWriter(final TsFileIOWriter writer) throws IOException, PageException {
+    writeTsFileData(writer);
   }
 
   public void addValueChunk(final ChunkHeader chunkHeader) {
@@ -167,6 +161,7 @@ public void serialize(final DataOutputStream stream) throws IOException {
     ReadWriteIOUtils.write(getType().ordinal(), stream);
     ReadWriteIOUtils.write(isAligned(), stream);
     serializeAttr(stream);
+    ReadWriteIOUtils.write(byteStream.size(), stream);
     byteStream.writeTo(stream);
     close();
   }
@@ -291,26 +286,41 @@ public void writeDecodeValuePage(
     }
   }
 
-  protected void deserializeTsFileData(final InputStream stream) throws IOException, PageException {
+  protected void writeTsFileData(TsFileIOWriter writer) throws IOException, PageException {
+    final InputStream stream = new LoadTsFilePieceNode.ByteBufferInputStream(chunkData);
     if (needDecodeChunk) {
-      buildChunkWriter(stream);
+      writeChunkToWriter(stream, writer);
+    } else {
+      writeEntireChunkToWriter(stream, writer);
+    }
+  }
+
+  protected void deserializeTsFileDataByte(final InputStream stream) throws IOException {
+    final int size = ReadWriteIOUtils.readInt(stream);
+    if (stream instanceof LoadTsFilePieceNode.ByteBufferInputStream) {
+      this.chunkData = ((LoadTsFilePieceNode.ByteBufferInputStream) stream).read(size);
     } else {
-      deserializeEntireChunk(stream);
+      byte[] data = new byte[size];
+      if (size != stream.read(data)) {
+        throw new IOException("TsFileData byte array read error, size mismatch.");
+      }
+      this.chunkData = ByteBuffer.wrap(data);
     }
   }
 
-  private void deserializeEntireChunk(final InputStream stream) throws IOException {
-    chunkList = new ArrayList<>();
+  private void writeEntireChunkToWriter(final InputStream stream, final TsFileIOWriter writer)
+      throws IOException {
     for (final ChunkHeader chunkHeader : chunkHeaderList) {
       final ByteBuffer chunkData =
           ByteBuffer.wrap(ReadWriteIOUtils.readBytesWithSelfDescriptionLength(stream));
       final Statistics<? extends Serializable> statistics =
           Statistics.deserialize(stream, chunkHeader.getDataType());
-      chunkList.add(new Chunk(chunkHeader, chunkData, null, statistics));
+      writer.writeChunk(new Chunk(chunkHeader, chunkData, null, statistics));
     }
   }
 
-  protected void buildChunkWriter(final InputStream stream) throws IOException, PageException {
+  protected void writeChunkToWriter(final InputStream stream, final TsFileIOWriter writer)
+      throws IOException, PageException {
     final List<IMeasurementSchema> measurementSchemaList = new ArrayList<>();
     IMeasurementSchema timeSchema = null;
     for (final ChunkHeader chunkHeader : chunkHeaderList) {
@@ -330,17 +340,20 @@ protected void buildChunkWriter(final InputStream stream) throws IOException, Pa
               chunkHeader.getEncodingType(),
               chunkHeader.getCompressionType()));
     }
-    chunkWriter = new AlignedChunkWriterImpl(timeSchema, measurementSchemaList);
+    AlignedChunkWriterImpl chunkWriter =
+        new AlignedChunkWriterImpl(timeSchema, measurementSchemaList);
     timeBatch = new ArrayList<>();
     final int chunkHeaderSize = chunkHeaderList.size();
     for (int i = 0; i < chunkHeaderSize; i++) {
-      buildChunk(stream, chunkHeaderList.get(i), pageNumbers.get(i), i - 1, i == 0);
+      buildChunk(chunkWriter, stream, chunkHeaderList.get(i), pageNumbers.get(i), i - 1, i == 0);
     }
     timeBatch = null;
+    chunkWriter.writeToFileWriter(writer);
   }
 
   @SuppressWarnings({"squid:S6541", "squid:S3776"})
   private void buildChunk(
+      final AlignedChunkWriterImpl chunkWriter,
       final InputStream stream,
       final ChunkHeader chunkHeader,
       final int pageNumber,
@@ -463,7 +476,7 @@ public static AlignedChunkData deserialize(final InputStream stream)
     chunkData.needDecodeChunk = needDecodeChunk;
     chunkData.chunkHeaderList = chunkHeaderList;
     chunkData.pageNumbers = pageNumbers;
-    chunkData.deserializeTsFileData(stream);
+    chunkData.deserializeTsFileDataByte(stream);
     chunkData.dataSize = dataSize;
     chunkData.close();
     return chunkData;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/BatchedAlignedValueChunkData.java

@@ -28,7 +28,6 @@
 import org.apache.tsfile.file.header.PageHeader;
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.file.metadata.statistics.Statistics;
-import org.apache.tsfile.read.common.Chunk;
 import org.apache.tsfile.utils.Binary;
 import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.apache.tsfile.utils.TsPrimitiveType;
@@ -41,8 +40,6 @@
 import java.io.InputStream;
 import java.io.Serializable;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
 
 /**
  * This class is used to be compatible with the new distribution of aligned series in chunk group.
@@ -52,8 +49,6 @@
  */
 public class BatchedAlignedValueChunkData extends AlignedChunkData {
 
-  private List<ValueChunkWriter> valueChunkWriters;
-
   // Used for splitter
   public BatchedAlignedValueChunkData(AlignedChunkData alignedChunkData) {
     super(alignedChunkData);
@@ -62,7 +57,6 @@ public BatchedAlignedValueChunkData(AlignedChunkData alignedChunkData) {
   // Used for deserialize
   public BatchedAlignedValueChunkData(IDeviceID device, TTimePartitionSlot timePartitionSlot) {
     super(device, timePartitionSlot);
-    valueChunkWriters = new ArrayList<>();
   }
 
   @Override
@@ -130,7 +124,8 @@ public void writeDecodeValuePage(long[] times, TsPrimitiveType[] values, TSDataT
   }
 
   @Override
-  protected void buildChunkWriter(final InputStream stream) throws IOException, PageException {
+  protected void writeChunkToWriter(final InputStream stream, final TsFileIOWriter writer)
+      throws IOException, PageException {
     for (int i = 0; i < chunkHeaderList.size(); i++) {
       ChunkHeader chunkHeader = chunkHeaderList.get(i);
       MeasurementSchema measurementSchema =
@@ -146,8 +141,8 @@ protected void buildChunkWriter(final InputStream stream) throws IOException, Pa
               measurementSchema.getType(),
               measurementSchema.getEncodingType(),
               measurementSchema.getValueEncoder());
-      valueChunkWriters.add(valueChunkWriter);
       buildValueChunkWriter(stream, chunkHeader, pageNumbers.get(i), valueChunkWriter);
+      valueChunkWriter.writeToFileWriter(writer);
     }
   }
 
@@ -226,17 +221,4 @@ private void buildValueChunkWriter(
       valueChunkWriter.sealCurrentPage();
     }
   }
-
-  @Override
-  public void writeToFileWriter(TsFileIOWriter writer) throws IOException {
-    if (chunkList != null) {
-      for (final Chunk chunk : chunkList) {
-        writer.writeChunk(chunk);
-      }
-    } else {
-      for (ValueChunkWriter valueChunkWriter : valueChunkWriters) {
-        valueChunkWriter.writeToFileWriter(writer);
-      }
-    }
-  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java

@@ -48,7 +48,7 @@ public interface ChunkData extends TsFileData {
 
   void writeDecodePage(long[] times, Object[] values, int satisfiedLength) throws IOException;
 
-  void writeToFileWriter(TsFileIOWriter writer) throws IOException;
+  void writeToFileWriter(TsFileIOWriter writer) throws IOException, PageException;
 
   @Override
   default TsFileDataType getType() {