Fix LOAD splitting of sparse batched aligned chunks (complete #19032)

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBLoadTsFileIT.java

@@ -26,14 +26,22 @@
 import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
 import org.apache.iotdb.itbase.env.BaseEnv;
 
+import org.apache.tsfile.encoding.encoder.PlainEncoder;
 import org.apache.tsfile.enums.ColumnCategory;
 import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.metadata.StringArrayDeviceID;
 import org.apache.tsfile.file.metadata.TableSchema;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.utils.Pair;
 import org.apache.tsfile.write.TsFileWriter;
+import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl;
+import org.apache.tsfile.write.chunk.ValueChunkWriter;
 import org.apache.tsfile.write.record.Tablet;
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.apache.tsfile.write.schema.Schema;
+import org.apache.tsfile.write.writer.TsFileIOWriter;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -51,6 +59,7 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
@@ -560,6 +569,119 @@ public void testLoadWithAllFieldsNullRows() throws Exception {
     }
   }
 
+  /**
+   * The compaction writer emits the series batches behind the first one as value chunks that follow
+   * all time chunks of the chunk group. When LOAD splits such a chunk by time partition it has to
+   * re-encode it, and the rows of the chunk keep their positions: a null value still occupies a
+   * bitmap position and an all-null page still occupies the page slot opposite its time page.
+   *
+   * <p>The rows span three time partitions of the interval configured in {@code setUp}, so LOAD has
+   * to decode the file and split it instead of storing it as it is.
+   */
+  @Test
+  public void testLoadSparseBatchedAlignedTsFile() throws Exception {
+    // The first page stays inside one time partition and is copied as it is, the second one crosses
+    // a partition boundary and holds a null row, and the third one is entirely null in both
+    // partitions it spans.
+    final long[][] pages = {
+      {0, 1, 2},
+      {PARTITION_INTERVAL - 2, PARTITION_INTERVAL - 1, PARTITION_INTERVAL, PARTITION_INTERVAL + 1},
+      {
+        2 * PARTITION_INTERVAL - 2,
+        2 * PARTITION_INTERVAL - 1,
+        2 * PARTITION_INTERVAL,
+        2 * PARTITION_INTERVAL + 1
+      }
+    };
+    final File file = new File(tmpDir, "1-0-0-0.tsfile");
+    writeSparseBatchedAlignedTsFile(file, pages);
+
+    try (final Connection connection =
+            EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        final Statement statement = connection.createStatement()) {
+      statement.execute(String.format("create database if not exists %s", SchemaConfig.DATABASE_0));
+      statement.execute(String.format("use %s", SchemaConfig.DATABASE_0));
+      statement.execute(
+          String.format(
+              "load '%s' with ('database'='%s')", file.getAbsolutePath(), SchemaConfig.DATABASE_0));
+
+      try (final ResultSet resultSet =
+          statement.executeQuery(
+              String.format("select time, f1, f2 from %s order by time", SchemaConfig.TABLE_0))) {
+        for (final long[] page : pages) {
+          for (final long time : page) {
+            Assert.assertTrue(resultSet.next());
+            Assert.assertEquals(time, resultSet.getLong("time"));
+            Assert.assertEquals(time, resultSet.getLong("f1"));
+            if (isNullField2(time)) {
+              resultSet.getLong("f2");
+              Assert.assertTrue(resultSet.wasNull());
+            } else {
+              Assert.assertEquals(time + 1, resultSet.getLong("f2"));
+            }
+          }
+        }
+        Assert.assertFalse(resultSet.next());
+      }
+    }
+  }
+
+  private void writeSparseBatchedAlignedTsFile(final File file, final long[][] pages)
+      throws Exception {
+    final Schema schema = new Schema();
+    schema.registerTableSchema(
+        new TableSchema(
+            SchemaConfig.TABLE_0,
+            Arrays.asList(
+                new MeasurementSchema("tag1", TSDataType.STRING),
+                new MeasurementSchema("f1", TSDataType.INT64),
+                new MeasurementSchema("f2", TSDataType.INT64)),
+            Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD)));
+
+    try (final TsFileIOWriter writer = new TsFileIOWriter(file)) {
+      writer.setSchema(schema);
+      writer.startChunkGroup(new StringArrayDeviceID(SchemaConfig.TABLE_0, "tagA"));
+
+      // The first field shares the chunk group of the time chunk.
+      final AlignedChunkWriterImpl alignedWriter =
+          new AlignedChunkWriterImpl(
+              Collections.singletonList(new MeasurementSchema("f1", TSDataType.INT64)));
+      for (final long[] page : pages) {
+        for (final long time : page) {
+          alignedWriter.getTimeChunkWriter().write(time);
+          alignedWriter.getValueChunkWriterByIndex(0).write(time, time, false);
+        }
+        alignedWriter.sealCurrentPage();
+      }
+      alignedWriter.writeToFileWriter(writer);
+
+      // The second field follows the time chunk as a value chunk of its own.
+      final ValueChunkWriter valueWriter =
+          new ValueChunkWriter(
+              "f2",
+              CompressionType.LZ4,
+              TSDataType.INT64,
+              TSEncoding.PLAIN,
+              new PlainEncoder(TSDataType.INT64, 8));
+      for (final long[] page : pages) {
+        for (final long time : page) {
+          valueWriter.write(time, time + 1, isNullField2(time));
+        }
+        valueWriter.sealCurrentPage();
+      }
+      valueWriter.writeToFileWriter(writer);
+
+      writer.endChunkGroup();
+      writer.endFile();
+    }
+  }
+
+  private static boolean isNullField2(final long time) {
+    // The row in the middle of the page that crosses the first boundary and the rows of the last
+    // page carry no value.
+    return time == PARTITION_INTERVAL - 1 || time >= 2 * PARTITION_INTERVAL - 2;
+  }
+
   private static class SchemaConfig {
     private static final String DATABASE_0 = "root";
     private static final String TABLE_0 = "test";

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

@@ -25,6 +25,7 @@
 
 import org.apache.tsfile.exception.write.PageException;
 import org.apache.tsfile.file.header.ChunkHeader;
+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;
@@ -33,7 +34,6 @@
 import org.apache.tsfile.write.schema.MeasurementSchema;
 
 import java.io.IOException;
-import java.io.Serializable;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -53,6 +53,12 @@ public class BatchedAlignedValueChunkData extends AlignedChunkData {
 
   private static final ByteBuffer EMPTY_CHUNK_DATA = ByteBuffer.allocate(0);
 
+  // An empty page holds no page body and no page statistics: it serializes to the single 0 varint
+  // that stands for it. The header must therefore not carry statistics, otherwise the page writer
+  // appends a compressed-size varint to a page without data and the reader sees a phantom page.
+  private static final ByteBuffer EMPTY_PAGE_DATA = ByteBuffer.allocate(0);
+  private static final PageHeader EMPTY_PAGE_HEADER = new PageHeader(0, 0, null);
+
   // Used for splitter
   public BatchedAlignedValueChunkData(AlignedChunkData alignedChunkData) {
     super(alignedChunkData);
@@ -64,9 +70,9 @@ public BatchedAlignedValueChunkData(IDeviceID device, TTimePartitionSlot timePar
   }
 
   /**
-   * Override the encoding process to handle standalone Value Chunks without a Time Chunk. It uses
-   * ValueChunkWriter directly, writes a dummy time (0), and manually corrects the start/end time in
-   * the statistics.
+   * Encodes standalone value chunks while preserving the row and the page positions of their time
+   * chunks: a row whose value is null still occupies a bitmap position, and an all-null page still
+   * occupies a page slot.
    */
   protected void encodeAndBuildChunks() throws IOException, PageException {
     final long startTime = timePartitionSlot.getStartTime();
@@ -112,28 +118,31 @@ protected void encodeAndBuildChunks() throws IOException, PageException {
               measurementSchema.getEncodingType(),
               measurementSchema.getValueEncoder());
 
-      // Filter and process the pages that belong to the current value chunk (index 'i').
-      // A chunk header can exist for a physical chunk whose points all belong to another time
-      // partition. Such a chunk stays empty here and must be emitted without any page: sealing an
-      // untouched writer would produce a zero-point page whose data size does not match the bytes
-      // actually written, and the resulting TsFile could not be parsed back.
-      boolean hasBufferedData = false;
-      for (PageBuffer page : chunkPages) {
+      // Whether this chunk covers any row of this time partition. Such a chunk writes one page for
+      // every page slot of its time chunk, all-null pages included: the aligned reader pairs the
+      // pages of a value chunk with the pages of its time chunk by position, so a missing page
+      // shifts every later page onto the wrong timestamps. A chunk that covers no row at all keeps
+      // the zero-page layout of an all-null column, which the metadata reader reports as null.
+      final boolean chunkHasRows = hasPageWithRows(chunkPages);
+      for (final PageBuffer page : chunkPages) {
         if (!page.needDecode) {
-          // An undecoded page may carry no data at all (for aligned series an all-null column is
-          // stored as an empty page). Writing its page header alone would produce a chunk whose
-          // data size does not describe a parsable page, so such a page is dropped.
-          if (page.pageData.remaining() > 0) {
-            valueChunkWriter.writePageHeaderAndDataIntoBuff(page.pageData, page.pageHeader);
-            hasBufferedData = true;
+          if (page.pageData.remaining() == 0) {
+            // An all-null page is stored as an empty page and keeps its slot. Its header is
+            // normalized because the page writer decides on the presence of the statistics instead
+            // of the page size: the placeholder page built by
+            // TsFileSplitter#handleEmptyValueChunk carries the statistics of an empty chunk, and
+            // those would turn the empty page into two bytes that the reader parses as a page
+            // followed by a phantom one.
+            if (chunkHasRows) {
+              valueChunkWriter.writePageHeaderAndDataIntoBuff(EMPTY_PAGE_DATA, EMPTY_PAGE_HEADER);
+            }
+            continue;
           }
+          valueChunkWriter.writePageHeaderAndDataIntoBuff(page.pageData, page.pageHeader);
           continue;
         }
 
-        // Keep track of actual bounds since we will write '0' as dummy time
-        long pageStartTime = Long.MAX_VALUE;
-        long pageEndTime = Long.MIN_VALUE;
-        boolean pageHasValue = false;
+        boolean pageHasRows = false;
 
         final TypeServices.TsPrimitiveValueChunkWriter valueChunkValueWriter =
             TypeServices.StorageEngine.TS_PRIMITIVE_VALUE_CHUNK_WRITER_SERVICE.call(
@@ -144,44 +153,34 @@ protected void encodeAndBuildChunks() throws IOException, PageException {
           if (time >= startTime && time <= endTime) {
             boolean isNull =
                 (page.valueBatch == null
-                    || page.valueBatch.length == 0
+                    || j >= page.valueBatch.length
                     || page.valueBatch[j] == null);
-            if (!isNull) {
-              hasBufferedData = true;
-              pageHasValue = true;
-              pageStartTime = Math.min(pageStartTime, time);
-              pageEndTime = Math.max(pageEndTime, time);
-
-              // Write 0 as a dummy time, valueChunkWriter will ignore it anyway for value columns
-              valueChunkValueWriter.write(valueChunkWriter, 0, page.valueBatch[j], false);
-            }
-            // A standalone value chunk cannot store nulls: for aligned series a missing value is
-            // represented by the absence of the point, the time chunk still carries the row.
+            // The bitmap has one position per timestamp, including nulls. Use the real timestamp
+            // so the writer also computes the statistics of the non-null values correctly.
+            valueChunkValueWriter.write(
+                valueChunkWriter, time, isNull ? null : page.valueBatch[j], isNull);
+            pageHasRows = true;
           }
         }
 
-        // Correct the page statistics with the actual Start/End time
-        Statistics<? extends Serializable> statistics =
-            valueChunkWriter.getPageWriter().getStatistics();
-        if (pageStartTime <= pageEndTime) {
-          statistics.setStartTime(pageStartTime);
-          statistics.setEndTime(pageEndTime);
-        }
-
-        if (pageHasValue) {
+        // Sealing an untouched writer would count a page that holds no byte.
+        if (pageHasRows) {
           valueChunkWriter.sealCurrentPage();
         }
       }
 
-      // Directly assemble the final standard Chunk into memory
+      // Directly assemble the final standard Chunk into memory. A chunk without any row of this
+      // partition stays empty: sealing its untouched writer would produce a page whose data size
+      // does not match the bytes actually written, and the resulting TsFile could not be parsed
+      // back.
       final ByteBuffer valueData =
-          hasBufferedData ? valueChunkWriter.getByteBuffer() : EMPTY_CHUNK_DATA;
+          chunkHasRows ? valueChunkWriter.getByteBuffer() : EMPTY_CHUNK_DATA;
       final Statistics<?> valueStat = valueChunkWriter.getStatistics();
       final ChunkHeader encodedHeader =
           createEncodedChunkHeader(
               chunkHeader,
               valueData.remaining(),
-              hasBufferedData ? valueChunkWriter.getNumOfPages() : 0,
+              chunkHasRows ? valueChunkWriter.getNumOfPages() : 0,
               VALUE_COLUMN_MASK);
 
       entireChunks.add(new Chunk(encodedHeader, valueData, null, valueStat));
@@ -190,4 +189,22 @@ protected void encodeAndBuildChunks() throws IOException, PageException {
       dataSize += valueData.remaining() + valueStat.getSerializedSize();
     }
   }
+
+  /**
+   * Whether any page of the chunk holds rows of this time partition. Only such a chunk mirrors the
+   * page slots of its time chunk, because only then does the time chunk have a page to pair with; a
+   * chunk that holds no row at all keeps the zero-page layout. This is the counterpart of the rule
+   * applied by the base class to the chunks it writes itself.
+   */
+  private static boolean hasPageWithRows(final List<PageBuffer> chunkPages) {
+    for (final PageBuffer page : chunkPages) {
+      // A decoded page already records how many of its rows belong to this partition, and a page
+      // copied as it is was read from this partition, so a body means rows and only an all-null
+      // page has no body.
+      if (page.needDecode ? page.satisfiedLength > 0 : page.pageData.remaining() > 0) {
+        return true;
+      }
+    }
+    return false;
+  }
 }