Improve VariableWidthBlock encoding

Avoids converting between offsets and lengths when serializing and
deserializing VariableWidthBlock instances, which enables a fast-path
conversion for blocks without nulls present. When nulls are present, the
compaction and expansion of offsets still outperforms the length to
offset conversion.

Author: pettyjamesm

core/trino-main/src/test/java/io/trino/execution/buffer/TestPagesSerde.java

@@ -218,13 +218,13 @@ public void testVarcharSerializedSize()
         pageSize = 44; // Now we have moved to the normal block implementation so the page size overhead is 44
         page = new Page(builder.build());
         int firstValueSize = serializedSize(ImmutableList.of(VARCHAR), page) - pageSize;
-        assertThat(firstValueSize).isEqualTo(8 + 5); // length + nonNullsCount + "alice"
+        assertThat(firstValueSize).isEqualTo(8 + 5); // ending offsets + nonNullsCount + "alice"
 
         // page with two values
         VARCHAR.writeString(builder, "bob");
         page = new Page(builder.build());
         int secondValueSize = serializedSize(ImmutableList.of(VARCHAR), page) - (pageSize + firstValueSize);
-        assertThat(secondValueSize).isEqualTo(4 + 3); // length + "bob" (null shared with first entry)
+        assertThat(secondValueSize).isEqualTo(4 + 3); // one additional ending offset + "bob" (null shared with first entry)
     }
 
     private int serializedSize(List<? extends Type> types, Page expectedPage)

core/trino-spi/src/main/java/io/trino/spi/block/VariableWidthBlockEncoding.java

@@ -19,8 +19,6 @@
 import io.airlift.slice.Slices;
 import jakarta.annotation.Nullable;
 
-import java.util.Arrays;
-
 import static io.trino.spi.block.EncoderUtil.decodeNullBits;
 import static io.trino.spi.block.EncoderUtil.encodeNullsAsBits;
 import static java.lang.String.format;
@@ -54,61 +52,26 @@ public void writeBlock(BlockEncodingSerde blockEncodingSerde, SliceOutput sliceO
         int arrayBaseOffset = variableWidthBlock.getRawArrayBase();
         @Nullable
         boolean[] isNull = variableWidthBlock.getRawValueIsNull();
-        int[] rawOffsets = variableWidthBlock.getRawOffsets();
-        checkFromIndexSize(arrayBaseOffset, positionCount + 1, rawOffsets.length);
-
-        // lengths
-        int[] lengths = new int[positionCount];
-        int totalLength = 0;
-        int nonNullsCount = 0;
-
-        for (int position = 0; position < positionCount; position++) {
-            int length = rawOffsets[position + arrayBaseOffset + 1] - rawOffsets[position + arrayBaseOffset];
-            totalLength += length;
-            lengths[nonNullsCount] = length;
-            nonNullsCount += isNull != null && isNull[position + arrayBaseOffset] ? 0 : 1;
-        }
+        encodeNullsAsBits(sliceOutput, isNull, arrayBaseOffset, positionCount);
 
-        sliceOutput
-                .appendInt(nonNullsCount)
-                .writeInts(lengths, 0, nonNullsCount);
+        int[] rawOffsets = variableWidthBlock.getRawOffsets();
+        writeOffsetsWithNullsCompacted(sliceOutput, rawOffsets, isNull, arrayBaseOffset, positionCount);
 
-        encodeNullsAsBits(sliceOutput, isNull, arrayBaseOffset, positionCount);
+        int startingOffset = rawOffsets[arrayBaseOffset];
+        int totalLength = rawOffsets[positionCount + arrayBaseOffset] - startingOffset;
 
         sliceOutput
                 .appendInt(totalLength)
-                .writeBytes(variableWidthBlock.getRawSlice(), variableWidthBlock.getPositionOffset(0), totalLength);
+                .writeBytes(variableWidthBlock.getRawSlice(), startingOffset, totalLength);
     }
 
     @Override
     public Block readBlock(BlockEncodingSerde blockEncodingSerde, SliceInput sliceInput)
     {
         int positionCount = sliceInput.readInt();
-        int nonNullsCount = sliceInput.readInt();
-
-        if (nonNullsCount > positionCount) {
-            throw new IllegalArgumentException(format("nonNullsCount must be <= positionCount, found: %s > %s", nonNullsCount, positionCount));
-        }
-
-        int[] offsets = new int[positionCount + 1];
-        // Read the lengths array into the end of the offsets array, since nonNullsCount <= positionCount
-        int lengthIndex = offsets.length - nonNullsCount;
-        sliceInput.readInts(offsets, lengthIndex, nonNullsCount);
-
         boolean[] valueIsNull = decodeNullBits(sliceInput, positionCount).orElse(null);
-        // Transform lengths back to offsets
-        if (valueIsNull == null) {
-            if (positionCount != nonNullsCount || lengthIndex != 1) {
-                throw new IllegalArgumentException(format("nonNullsCount must equal positionCount, found: %s <> %s", nonNullsCount, positionCount));
-            }
-            // Simplified loop for no nulls present
-            for (int i = 1; i < offsets.length; i++) {
-                offsets[i] += offsets[i - 1];
-            }
-        }
-        else {
-            computeOffsetsFromLengths(offsets, valueIsNull, lengthIndex);
-        }
+
+        int[] offsets = readOffsetsWithNullsCompacted(sliceInput, valueIsNull, positionCount);
 
         int blockSize = sliceInput.readInt();
         Slice slice = Slices.allocate(blockSize);
@@ -117,28 +80,69 @@ public Block readBlock(BlockEncodingSerde blockEncodingSerde, SliceInput sliceIn
         return new VariableWidthBlock(0, positionCount, slice, offsets, valueIsNull);
     }
 
-    private static void computeOffsetsFromLengths(int[] offsets, boolean[] valueIsNull, int lengthIndex)
+    private static void writeOffsetsWithNullsCompacted(SliceOutput sliceOutput, int[] rawOffsets, @Nullable boolean[] valueIsNull, int baseOffset, int positionCount)
+    {
+        checkFromIndexSize(baseOffset, positionCount + 1, rawOffsets.length);
+
+        int startingOffset = rawOffsets[baseOffset];
+        if (valueIsNull == null && startingOffset == 0) {
+            // No translation of offsets required, write the range of raw offsets directly to the output
+            sliceOutput
+                    .appendInt(positionCount)
+                    .writeInts(rawOffsets, baseOffset + 1, positionCount);
+        }
+        else {
+            int[] nonNullOffsets;
+            int nonNullOffsetsCount;
+            if (valueIsNull == null) {
+                // Subtract starting offset from each ending offset to translate them to start from zero, no null suppression required
+                nonNullOffsets = new int[positionCount];
+                for (int i = 0; i < nonNullOffsets.length; i++) {
+                    nonNullOffsets[i] = rawOffsets[i + baseOffset + 1] - startingOffset;
+                }
+                nonNullOffsetsCount = nonNullOffsets.length;
+            }
+            else {
+                // Translate ending offsets and suppress null values from the output
+                nonNullOffsets = new int[positionCount];
+                nonNullOffsetsCount = 0;
+                for (int i = 0; i < positionCount; i++) {
+                    nonNullOffsets[nonNullOffsetsCount] = rawOffsets[i + baseOffset + 1] - startingOffset;
+                    nonNullOffsetsCount += valueIsNull[i + baseOffset] ? 0 : 1;
+                }
+            }
+            sliceOutput
+                    .appendInt(nonNullOffsetsCount)
+                    .writeInts(nonNullOffsets, 0, nonNullOffsetsCount);
+        }
+    }
+
+    private static int[] readOffsetsWithNullsCompacted(SliceInput sliceInput, @Nullable boolean[] valueIsNull, int positionCount)
     {
-        if (lengthIndex < 0 || lengthIndex > offsets.length) {
-            throw new IllegalArgumentException(format("Invalid lengthIndex %s for offsets %s", lengthIndex, offsets.length));
+        if (valueIsNull != null && valueIsNull.length != positionCount) {
+            throw new IllegalArgumentException(format("valueIsNull length must match positionCount, found %s <> %s", valueIsNull.length, positionCount));
+        }
+        int nonNullOffsetCount = sliceInput.readInt();
+        if (nonNullOffsetCount > positionCount) {
+            throw new IllegalArgumentException(format("nonNullOffsetCount must be <= positionCount, found: %s > %s", nonNullOffsetCount, positionCount));
         }
-        int currentOffset = 0;
-        for (int i = 1; i < offsets.length; i++) {
-            if (lengthIndex == offsets.length) {
-                // Populate remaining null elements
-                Arrays.fill(offsets, i, offsets.length, currentOffset);
-                break;
+        // Offsets are read into the end of the array, expansion will pull values down into the lower range until null positions are expanded in place
+        int[] offsets = new int[positionCount + 1];
+        int compactIndex = offsets.length - nonNullOffsetCount;
+        sliceInput.readInts(offsets, compactIndex, nonNullOffsetCount);
+        if (valueIsNull == null || compactIndex == 1) {
+            if (positionCount != nonNullOffsetCount) {
+                throw new IllegalArgumentException(format("nonNullOffsetCount must match positionCount, found %s <> %s", nonNullOffsetCount, positionCount));
             }
-            boolean isNull = valueIsNull[i - 1];
-            // must be accessed unconditionally, otherwise CMOV optimization isn't applied due to
-            // ArrayIndexOutOfBoundsException checks
-            int length = offsets[lengthIndex];
-            lengthIndex += isNull ? 0 : 1;
-            currentOffset += isNull ? 0 : length;
-            offsets[i] = currentOffset;
+            return offsets;
         }
-        if (lengthIndex != offsets.length) {
-            throw new IllegalArgumentException(format("Failed to consume all length entries, found %s <> %s", lengthIndex, offsets.length));
+        // Shift the offsets from the end of the offsets array downwards, repeating the previous offset when nulls are encountered
+        // until no more nulls are present
+        int readFrom = compactIndex - 1;
+        for (int position = 0; position < readFrom; position++) {
+            offsets[position] = offsets[readFrom];
+            readFrom += valueIsNull[position] ? 0 : 1;
         }
+        return offsets;
     }
 }