perf: add encodeBatch interface to accelerate flushing for sequential insert (apache#15243)

* perf: batch encode to improve insert performance

* fix IoTDBSimpleQueryIT

* move batchEncodeInfo / times as MemTableFlushTask member

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java

@@ -27,12 +27,15 @@
 import org.apache.iotdb.db.service.metrics.WritingMetrics;
 import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
 import org.apache.iotdb.db.storageengine.dataregion.flush.pool.FlushSubTaskPoolManager;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.AlignedWritableMemChunk;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunk;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.IWritableMemChunkGroup;
 import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo;
+import org.apache.iotdb.db.utils.datastructure.BatchEncodeInfo;
 import org.apache.iotdb.metrics.utils.MetricLevel;
 
+import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.file.metadata.IDeviceID;
 import org.apache.tsfile.write.chunk.IChunkWriter;
 import org.apache.tsfile.write.writer.RestorableTsFileIOWriter;
@@ -62,6 +65,9 @@ public class MemTableFlushTask {
       FlushSubTaskPoolManager.getInstance();
   private static final WritingMetrics WRITING_METRICS = WritingMetrics.getInstance();
   private static final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+  private final int MAX_NUMBER_OF_POINTS_IN_PAGE =
+      TSFileDescriptor.getInstance().getConfig().getMaxNumberOfPointsInPage();
+
   /* storage group name -> last time */
   private static final Map<String, Long> flushPointsCache = new ConcurrentHashMap<>();
   private final Future<?> encodingTaskFuture;
@@ -82,6 +88,9 @@ public class MemTableFlushTask {
   private volatile long memSerializeTime = 0L;
   private volatile long ioTime = 0L;
 
+  private final BatchEncodeInfo encodeInfo;
+  private long[] times;
+
   /**
    * @param memTable the memTable to flush
    * @param writer the writer where memTable will be flushed to (current tsfile writer or vm writer)
@@ -98,6 +107,7 @@ public MemTableFlushTask(
     this.dataRegionId = dataRegionId;
     this.encodingTaskFuture = SUB_TASK_POOL_MANAGER.submit(encodingTask);
     this.ioTaskFuture = SUB_TASK_POOL_MANAGER.submit(ioTask);
+    this.encodeInfo = new BatchEncodeInfo(0, 0, 0);
     LOGGER.debug(
         "flush task of database {} memtable is created, flushing to file {}.",
         storageGroup,
@@ -248,7 +258,10 @@ public void run() {
             } else {
               long starTime = System.currentTimeMillis();
               IWritableMemChunk writableMemChunk = (IWritableMemChunk) task;
-              writableMemChunk.encode(ioTaskQueue);
+              if (writableMemChunk instanceof AlignedWritableMemChunk && times == null) {
+                times = new long[MAX_NUMBER_OF_POINTS_IN_PAGE];
+              }
+              writableMemChunk.encode(ioTaskQueue, encodeInfo, times);
               long subTaskTime = System.currentTimeMillis() - starTime;
               WRITING_METRICS.recordFlushSubTaskCost(WritingMetrics.ENCODING_TASK, subTaskTime);
               memSerializeTime += subTaskTime;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractWritableMemChunk.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
 import org.apache.iotdb.db.queryengine.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView;
+import org.apache.iotdb.db.utils.datastructure.BatchEncodeInfo;
 import org.apache.iotdb.db.utils.datastructure.TVList;
 
 import org.apache.tsfile.enums.TSDataType;
@@ -192,7 +193,8 @@ public abstract void writeAlignedTablet(
   public abstract IChunkWriter createIChunkWriter();
 
   @Override
-  public abstract void encode(BlockingQueue<Object> ioTaskQueue);
+  public abstract void encode(
+      BlockingQueue<Object> ioTaskQueue, BatchEncodeInfo encodeInfo, long[] times);
 
   @Override
   public abstract void release();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedWritableMemChunk.java

@@ -25,21 +25,20 @@
 import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView;
 import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALWriteUtils;
 import org.apache.iotdb.db.utils.datastructure.AlignedTVList;
+import org.apache.iotdb.db.utils.datastructure.BatchEncodeInfo;
 import org.apache.iotdb.db.utils.datastructure.MemPointIterator;
 import org.apache.iotdb.db.utils.datastructure.MemPointIteratorFactory;
 import org.apache.iotdb.db.utils.datastructure.TVList;
 
 import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.read.common.TimeRange;
-import org.apache.tsfile.read.common.block.TsBlock;
 import org.apache.tsfile.utils.Binary;
 import org.apache.tsfile.utils.BitMap;
 import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.write.UnSupportedDataTypeException;
 import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl;
 import org.apache.tsfile.write.chunk.IChunkWriter;
-import org.apache.tsfile.write.chunk.ValueChunkWriter;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.apache.tsfile.write.schema.MeasurementSchema;
 
@@ -65,6 +64,7 @@ public class AlignedWritableMemChunk extends AbstractWritableMemChunk {
   private final List<IMeasurementSchema> schemaList;
   private AlignedTVList list;
   private List<AlignedTVList> sortedList;
+  private long sortedRowCount = 0;
   private final boolean ignoreAllNullRows;
 
   private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig();
@@ -197,6 +197,7 @@ protected void handoverAlignedTvList() {
       list.sort();
     }
     sortedList.add(list);
+    this.sortedRowCount += list.rowCount();
     this.list = AlignedTVList.newAlignedList(new ArrayList<>(dataTypes));
     this.dataTypes = list.getTsDataTypes();
   }
@@ -352,15 +353,11 @@ public long count() {
 
   @Override
   public long rowCount() {
-    return alignedListSize();
+    return sortedRowCount + list.rowCount();
   }
 
   public int alignedListSize() {
-    int rowCount = list.rowCount();
-    for (AlignedTVList alignedTvList : sortedList) {
-      rowCount += alignedTvList.rowCount();
-    }
-    return rowCount;
+    return (int) rowCount();
   }
 
   @Override
@@ -638,107 +635,56 @@ private void handleEncoding(
   }
 
   @Override
-  public synchronized void encode(BlockingQueue<Object> ioTaskQueue) {
+  public synchronized void encode(
+      BlockingQueue<Object> ioTaskQueue, BatchEncodeInfo encodeInfo, long[] times) {
     if (TVLIST_SORT_THRESHOLD == 0) {
       encodeWorkingAlignedTVList(ioTaskQueue);
       return;
     }
 
     AlignedChunkWriterImpl alignedChunkWriter = new AlignedChunkWriterImpl(schemaList);
+
     // create MergeSortAlignedTVListIterator.
     List<AlignedTVList> alignedTvLists = new ArrayList<>(sortedList);
     alignedTvLists.add(list);
+    List<Integer> columnIndexList = buildColumnIndexList(schemaList);
     MemPointIterator timeValuePairIterator =
-        MemPointIteratorFactory.create(dataTypes, null, alignedTvLists, ignoreAllNullRows);
-
-    int pointNumInPage = 0;
-    int pointNumInChunk = 0;
-    long[] times = new long[MAX_NUMBER_OF_POINTS_IN_PAGE];
+        MemPointIteratorFactory.create(
+            dataTypes, columnIndexList, alignedTvLists, ignoreAllNullRows);
 
     while (timeValuePairIterator.hasNextBatch()) {
-      TsBlock tsBlock = timeValuePairIterator.nextBatch();
-      if (tsBlock == null) {
-        continue;
+      timeValuePairIterator.encodeBatch(alignedChunkWriter, encodeInfo, times);
+      if (encodeInfo.pointNumInPage >= MAX_NUMBER_OF_POINTS_IN_PAGE) {
+        alignedChunkWriter.write(times, encodeInfo.pointNumInPage, 0);
+        encodeInfo.pointNumInPage = 0;
       }
-      for (int rowIndex = 0; rowIndex < tsBlock.getPositionCount(); rowIndex++) {
-        long time = tsBlock.getTimeByIndex(rowIndex);
-        times[pointNumInPage] = time;
 
-        for (int columnIndex = 0; columnIndex < dataTypes.size(); columnIndex++) {
-          ValueChunkWriter valueChunkWriter =
-              alignedChunkWriter.getValueChunkWriterByIndex(columnIndex);
-          if (tsBlock.getColumn(columnIndex).isNull(rowIndex)) {
-            valueChunkWriter.write(time, null, true);
-            continue;
-          }
-          switch (schemaList.get(columnIndex).getType()) {
-            case BOOLEAN:
-              valueChunkWriter.write(
-                  time, tsBlock.getColumn(columnIndex).getBoolean(rowIndex), false);
-              break;
-            case INT32:
-            case DATE:
-              valueChunkWriter.write(time, tsBlock.getColumn(columnIndex).getInt(rowIndex), false);
-              break;
-            case INT64:
-            case TIMESTAMP:
-              valueChunkWriter.write(time, tsBlock.getColumn(columnIndex).getLong(rowIndex), false);
-              break;
-            case FLOAT:
-              valueChunkWriter.write(
-                  time, tsBlock.getColumn(columnIndex).getFloat(rowIndex), false);
-              break;
-            case DOUBLE:
-              valueChunkWriter.write(
-                  time, tsBlock.getColumn(columnIndex).getDouble(rowIndex), false);
-              break;
-            case TEXT:
-            case BLOB:
-            case STRING:
-              valueChunkWriter.write(
-                  time, tsBlock.getColumn(columnIndex).getBinary(rowIndex), false);
-              break;
-            default:
-              break;
-          }
-        }
-        pointNumInPage++;
-        pointNumInChunk++;
-
-        // new page
-        if (pointNumInPage == MAX_NUMBER_OF_POINTS_IN_PAGE
-            || pointNumInChunk >= maxNumberOfPointsInChunk) {
-          alignedChunkWriter.write(times, pointNumInPage, 0);
-          pointNumInPage = 0;
-        }
-
-        // new chunk
-        if (pointNumInChunk >= maxNumberOfPointsInChunk) {
-          alignedChunkWriter.sealCurrentPage();
-          alignedChunkWriter.clearPageWriter();
-          try {
-            ioTaskQueue.put(alignedChunkWriter);
-          } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-          }
-          alignedChunkWriter = new AlignedChunkWriterImpl(schemaList);
-          pointNumInChunk = 0;
+      if (encodeInfo.pointNumInChunk >= maxNumberOfPointsInChunk) {
+        alignedChunkWriter.sealCurrentPage();
+        alignedChunkWriter.clearPageWriter();
+        try {
+          ioTaskQueue.put(alignedChunkWriter);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
         }
+        alignedChunkWriter = new AlignedChunkWriterImpl(schemaList);
+        encodeInfo.reset();
       }
     }
 
     // last batch of points
-    if (pointNumInChunk > 0) {
-      if (pointNumInPage > 0) {
-        alignedChunkWriter.write(times, pointNumInPage, 0);
-        alignedChunkWriter.sealCurrentPage();
-        alignedChunkWriter.clearPageWriter();
+    if (encodeInfo.pointNumInChunk > 0) {
+      if (encodeInfo.pointNumInPage > 0) {
+        alignedChunkWriter.write(times, encodeInfo.pointNumInPage, 0);
       }
+      alignedChunkWriter.sealCurrentPage();
+      alignedChunkWriter.clearPageWriter();
       try {
         ioTaskQueue.put(alignedChunkWriter);
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
       }
+      encodeInfo.reset();
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/IWritableMemChunk.java

@@ -20,6 +20,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryValue;
+import org.apache.iotdb.db.utils.datastructure.BatchEncodeInfo;
 import org.apache.iotdb.db.utils.datastructure.TVList;
 
 import org.apache.tsfile.enums.TSDataType;
@@ -111,7 +112,7 @@ default long getMinTime() {
 
   IChunkWriter createIChunkWriter();
 
-  void encode(BlockingQueue<Object> ioTaskQueue);
+  void encode(BlockingQueue<Object> ioTaskQueue, BatchEncodeInfo encodeInfo, long[] times);
 
   void release();