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 53879c21bf27..38e0c2bb7039 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 @@ -119,6 +119,45 @@ public void buildEntry(ArrayValueBuilder builder) currentEntryOpened = false; } + public ArrayEntryBuilder buildEntry() + { + return new ArrayEntryBuilderImplementation(); + } + + private class ArrayEntryBuilderImplementation + implements ArrayEntryBuilder + { + private boolean entryBuilt; + + public ArrayEntryBuilderImplementation() + { + if (currentEntryOpened) { + throw new IllegalStateException("Expected current entry to be closed but was opened"); + } + currentEntryOpened = true; + } + + @Override + public BlockBuilder getElementBuilder() + { + if (entryBuilt || !currentEntryOpened) { + throw new IllegalStateException("Entry has already been built"); + } + return values; + } + + @Override + public void build() + { + if (entryBuilt || !currentEntryOpened) { + throw new IllegalStateException("Entry has already been built"); + } + entryBuilt = true; + entryAdded(false); + currentEntryOpened = false; + } + } + @Override public void append(ValueBlock block, int position) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayEntryBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayEntryBuilder.java new file mode 100644 index 000000000000..c37adce7f855 --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayEntryBuilder.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.spi.block; + +public interface ArrayEntryBuilder +{ + /** + * Get the element BlockBuilder. + * The block builder must not be retained or used after build() is called. + */ + BlockBuilder getElementBuilder(); + + /** + * Finalize the entry the elements have been appended to the element builder. + * This method MUST be called exactly once. + */ + void build(); +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java index 003e95f069d6..35ecf22e8e1e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java @@ -113,6 +113,45 @@ public void buildEntry(RowValueBuilder builder) currentEntryOpened = false; } + public RowEntryBuilder buildEntry() + { + return new RowEntryBuilderImplementation(); + } + + private class RowEntryBuilderImplementation + implements RowEntryBuilder + { + private boolean entryBuilt; + + public RowEntryBuilderImplementation() + { + if (currentEntryOpened) { + throw new IllegalStateException("Expected current entry to be closed but was opened"); + } + currentEntryOpened = true; + } + + @Override + public BlockBuilder getFieldBuilder(int fieldId) + { + if (entryBuilt || !currentEntryOpened) { + throw new IllegalStateException("Entry has already been built"); + } + return fieldBlockBuilders[fieldId]; + } + + @Override + public void build() + { + if (entryBuilt || !currentEntryOpened) { + throw new IllegalStateException("Entry has already been built"); + } + entryBuilt = true; + entryAdded(false); + currentEntryOpened = false; + } + } + @Override public void append(ValueBlock block, int position) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowEntryBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowEntryBuilder.java new file mode 100644 index 000000000000..db80682498bf --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowEntryBuilder.java @@ -0,0 +1,30 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.spi.block; + +public interface RowEntryBuilder +{ + /** + * Get the BlockBuilder for a specific field in the row. + * Only a single value should be appended to the returned BlockBuilder before calling build(). + * The block builder must not be retained or used after build() is called. + */ + BlockBuilder getFieldBuilder(int fieldId); + + /** + * Finalize the entry after ALL field values have been appended to the field builders. + * This method MUST be called exactly once after all field builders have been used. + */ + void build(); +} 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 689504f8f927..86c2d967d9f4 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 @@ -94,6 +94,56 @@ public void testBuilderProducesNullRleForNullRows() assertIsAllNulls(blockBuilder().appendNull().appendNull().build(), 2); } + @Test + void buildEntry() + { + List> values = getTestValues(); + assertThat(values) + .hasSize(5) + .doesNotHaveDuplicates() + .doesNotContainNull(); + + ArrayBlockBuilder blockBuilder = (ArrayBlockBuilder) createBlockBuilder(); + for (List array : values) { + blockBuilder.buildEntry(elementBuilder -> { + for (String element : array) { + if (element == null) { + elementBuilder.appendNull(); + } + else { + VARCHAR.writeString(elementBuilder, element); + } + } + }); + } + assertThat(blockToValues(blockBuilder.buildValueBlock())).isEqualTo(values); + + blockBuilder = (ArrayBlockBuilder) createBlockBuilder(); + for (List array : values) { + ArrayEntryBuilder arrayEntryBuilder = blockBuilder.buildEntry(); + for (String element : array) { + if (element == null) { + arrayEntryBuilder.getElementBuilder().appendNull(); + } + else { + VARCHAR.writeString(arrayEntryBuilder.getElementBuilder(), element); + } + } + arrayEntryBuilder.build(); + } + assertThat(blockToValues(blockBuilder.buildValueBlock())).isEqualTo(values); + + blockBuilder = (ArrayBlockBuilder) createBlockBuilder(); + blockBuilder.buildEntry(); + assertThatThrownBy(blockBuilder::buildEntry).isInstanceOf(IllegalStateException.class); + + blockBuilder = (ArrayBlockBuilder) createBlockBuilder(); + ArrayEntryBuilder multipleEntryBuilder = blockBuilder.buildEntry(); + multipleEntryBuilder.build(); + assertThatThrownBy(multipleEntryBuilder::getElementBuilder).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(multipleEntryBuilder::build).isInstanceOf(IllegalStateException.class); + } + private static BlockBuilder blockBuilder() { return new ArrayBlockBuilder(BIGINT, null, 10); diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockBuilder.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockBuilder.java index 5a0ff445388f..b5d0087087d7 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockBuilder.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockBuilder.java @@ -24,6 +24,7 @@ import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.VarcharType.VARCHAR; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; public class TestRowBlockBuilder extends AbstractTestBlockBuilder @@ -41,6 +42,62 @@ public void testBuilderProducesNullRleForNullRows() assertIsAllNulls(blockBuilder().appendNull().appendNull().build(), 2); } + @Test + void buildEntry() + { + List values = getTestValues(); + assertThat(values) + .hasSize(5) + .doesNotHaveDuplicates() + .doesNotContainNull() + .doesNotContain(getUnusedTestValue()); + + RowBlockBuilder blockBuilder = new RowBlockBuilder(List.of(VARCHAR, INTEGER, BOOLEAN), null, 1); + for (TestRow row : values) { + blockBuilder.buildEntry(fieldBuilders -> { + if (row.name() == null) { + fieldBuilders.getFirst().appendNull(); + } + else { + VARCHAR.writeString(fieldBuilders.getFirst(), row.name()); + } + INTEGER.writeLong(fieldBuilders.get(1), row.number()); + BOOLEAN.writeBoolean(fieldBuilders.get(2), row.flag()); + }); + } + assertThat(blockToValues(blockBuilder.buildValueBlock())).isEqualTo(values); + + blockBuilder = new RowBlockBuilder(List.of(VARCHAR, INTEGER, BOOLEAN), null, 1); + for (TestRow row : values) { + RowEntryBuilder rowEntryBuilder = blockBuilder.buildEntry(); + if (row.name() == null) { + rowEntryBuilder.getFieldBuilder(0).appendNull(); + } + else { + VARCHAR.writeString(rowEntryBuilder.getFieldBuilder(0), row.name()); + } + INTEGER.writeLong(rowEntryBuilder.getFieldBuilder(1), row.number()); + BOOLEAN.writeBoolean(rowEntryBuilder.getFieldBuilder(2), row.flag()); + rowEntryBuilder.build(); + } + assertThat(blockToValues(blockBuilder.buildValueBlock())).isEqualTo(values); + + blockBuilder = new RowBlockBuilder(List.of(VARCHAR, INTEGER, BOOLEAN), null, 1); + blockBuilder.buildEntry(); + assertThatThrownBy(blockBuilder::buildEntry).isInstanceOf(IllegalStateException.class); + + blockBuilder = new RowBlockBuilder(List.of(VARCHAR, INTEGER, BOOLEAN), null, 1); + RowEntryBuilder incompleteEntryBuilder = blockBuilder.buildEntry(); + assertThatThrownBy(incompleteEntryBuilder::build).isInstanceOf(IllegalStateException.class); + + blockBuilder = new RowBlockBuilder(List.of(BOOLEAN), null, 1); + RowEntryBuilder multipleEntryBuilder = blockBuilder.buildEntry(); + BOOLEAN.writeBoolean(multipleEntryBuilder.getFieldBuilder(0), true); + multipleEntryBuilder.build(); + assertThatThrownBy(() -> multipleEntryBuilder.getFieldBuilder(0)).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(multipleEntryBuilder::build).isInstanceOf(IllegalStateException.class); + } + private static BlockBuilder blockBuilder() { return new RowBlockBuilder(ImmutableList.of(BIGINT), null, 10);