diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java index 1748121fa8fa..fc67f31ac585 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java @@ -23,6 +23,8 @@ import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.spi.block.ArrayBlock.createArrayBlockInternal; +import static io.trino.spi.block.BlockUtil.checkArrayRange; +import static io.trino.spi.block.BlockUtil.checkValidRegion; import static java.lang.Math.max; import static java.util.Objects.requireNonNull; @@ -42,6 +44,7 @@ public class ArrayBlockBuilder private int[] offsets = new int[1]; private boolean[] valueIsNull = new boolean[0]; private boolean hasNullValue; + private boolean hasNonNullRow; private final BlockBuilder values; private boolean currentEntryOpened; @@ -186,6 +189,7 @@ private void entryAdded(boolean isNull) offsets[positionCount + 1] = values.getPositionCount(); valueIsNull[positionCount] = isNull; hasNullValue |= isNull; + hasNonNullRow |= !isNull; positionCount++; if (blockBuilderStatus != null) { @@ -218,11 +222,14 @@ private void updateDataSize() } @Override - public ArrayBlock build() + public Block build() { if (currentEntryOpened) { throw new IllegalStateException("Current entry must be closed before the block can be built"); } + if (!hasNonNullRow) { + return nullRle(positionCount); + } return createArrayBlockInternal(0, positionCount, hasNullValue ? valueIsNull : null, offsets, values.build()); } @@ -240,4 +247,45 @@ public String toString() sb.append('}'); return sb.toString(); } + + @Override + public Block copyPositions(int[] positions, int offset, int length) + { + checkArrayRange(positions, offset, length); + + if (!hasNonNullRow) { + return nullRle(length); + } + return super.copyPositions(positions, offset, length); + } + + @Override + public Block getRegion(int position, int length) + { + int positionCount = getPositionCount(); + checkValidRegion(positionCount, position, length); + + if (!hasNonNullRow) { + return nullRle(length); + } + return super.getRegion(position, length); + } + + @Override + public Block copyRegion(int position, int length) + { + int positionCount = getPositionCount(); + checkValidRegion(positionCount, position, length); + + if (!hasNonNullRow) { + return nullRle(length); + } + return super.copyRegion(position, length); + } + + private RunLengthEncodedBlock nullRle(int positionCount) + { + ArrayBlock nullValueBlock = createArrayBlockInternal(0, 1, new boolean[] {true}, new int[] {0, 0}, values.newBlockBuilderLike(null).build()); + return new RunLengthEncodedBlock(nullValueBlock, positionCount); + } } diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java index 5d669ac76f4c..9b69753502b2 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java @@ -73,4 +73,36 @@ public void testConcurrentWriting() .isInstanceOf(IllegalStateException.class) .hasMessage("Expected current entry to be closed but was opened"); } + + @Test + public void testBuilderProducesNullRleForNullRows() + { + // empty block + assertIsNullRle(blockBuilder().build(), 0); + + // single null + assertIsNullRle(blockBuilder().appendNull().build(), 1); + + // multiple nulls + assertIsNullRle(blockBuilder().appendNull().appendNull().build(), 2); + + BlockBuilder blockBuilder = blockBuilder().appendNull().appendNull(); + assertIsNullRle(blockBuilder.copyPositions(new int[] {0}, 0, 1), 1); + assertIsNullRle(blockBuilder.getRegion(0, 1), 1); + assertIsNullRle(blockBuilder.copyRegion(0, 1), 1); + } + + private static BlockBuilder blockBuilder() + { + return new ArrayBlockBuilder(BIGINT, null, 10); + } + + private void assertIsNullRle(Block block, int expectedPositionCount) + { + assertEquals(block.getPositionCount(), expectedPositionCount); + assertEquals(block.getClass(), RunLengthEncodedBlock.class); + if (expectedPositionCount > 0) { + assertTrue(block.isNull(0)); + } + } }