From 2e4ed63db0b9a124d597f072b7f7689baf27caa8 Mon Sep 17 00:00:00 2001 From: Vigneshwar Selvaraj Date: Fri, 1 Dec 2023 18:32:04 -0800 Subject: [PATCH] Add BlockBuilderUtil class to support null map keys for OrcReader --- .../common/block/BlockBuilderUtils.java | 129 ++++++++++++++++ .../common/block/TestBlockBuilderUtils.java | 146 ++++++++++++++++++ 2 files changed, 275 insertions(+) create mode 100644 presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilderUtils.java create mode 100644 presto-common/src/test/java/com/facebook/presto/common/block/TestBlockBuilderUtils.java diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilderUtils.java b/presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilderUtils.java new file mode 100644 index 000000000000..5994bf2e6bd0 --- /dev/null +++ b/presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilderUtils.java @@ -0,0 +1,129 @@ +/* + * 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 com.facebook.presto.common.block; + +public class BlockBuilderUtils +{ + private BlockBuilderUtils() + { + // utility class + } + + public static void writePositionToBlockBuilder(Block block, int position, BlockBuilder blockBuilder) + { + if (block instanceof DictionaryBlock) { + position = ((DictionaryBlock) block).getId(position); + block = ((DictionaryBlock) block).getDictionary(); + } + + if (blockBuilder instanceof MapBlockBuilder) { + writePositionToMapBuilder(block, position, (MapBlockBuilder) blockBuilder); + } + else if (blockBuilder instanceof ArrayBlockBuilder) { + writePositionToArrayBuilder(block, position, (ArrayBlockBuilder) blockBuilder); + } + else if (blockBuilder instanceof RowBlockBuilder) { + writePositionToRowBuilder(block, position, (RowBlockBuilder) blockBuilder); + } + else { + block.writePositionTo(position, blockBuilder); + } + } + + public static void writePositionToMapBuilder(Block block, int position, MapBlockBuilder mapBlockBuilder) + { + if (!(block instanceof AbstractMapBlock)) { + throw new IllegalArgumentException("Expected AbstractMapBlock"); + } + + mapBlockBuilder.beginBlockEntry(); + BlockBuilder keyBlockBuilder = mapBlockBuilder.getKeyBlockBuilder(); + BlockBuilder valueBlockBuilder = mapBlockBuilder.getValueBlockBuilder(); + + AbstractMapBlock mapBlock = (AbstractMapBlock) block; + int startOffset = mapBlock.getOffset(position); + int endOffset = mapBlock.getOffset(position + 1); + + Block keyBlock = mapBlock.getRawKeyBlock(); + Block valueBlock = mapBlock.getRawValueBlock(); + + for (int i = startOffset; i < endOffset; i++) { + if (keyBlock.isNull(i)) { + keyBlockBuilder.appendNull(); + } + else { + writePositionToBlockBuilder(keyBlock, i, keyBlockBuilder); + } + } + + for (int i = startOffset; i < endOffset; i++) { + if (valueBlock.isNull(i)) { + valueBlockBuilder.appendNull(); + } + else { + writePositionToBlockBuilder(valueBlock, i, valueBlockBuilder); + } + } + mapBlockBuilder.closeEntry(); + } + + private static void writePositionToArrayBuilder(Block block, int position, ArrayBlockBuilder arrayBlockBuilder) + { + if (!(block instanceof AbstractArrayBlock)) { + throw new IllegalArgumentException("Expected AbstractArrayBlock"); + } + + arrayBlockBuilder.beginBlockEntry(); + BlockBuilder elementBlockBuilder = arrayBlockBuilder.getElementBlockBuilder(); + + AbstractArrayBlock arrayBlock = (AbstractArrayBlock) block; + int startOffset = arrayBlock.getOffset(position); + int endOffset = arrayBlock.getOffset(position + 1); + + Block elementBlock = arrayBlock.getRawElementBlock(); + + for (int i = startOffset; i < endOffset; i++) { + if (elementBlock.isNull(i)) { + elementBlockBuilder.appendNull(); + } + else { + writePositionToBlockBuilder(elementBlock, i, elementBlockBuilder); + } + } + arrayBlockBuilder.closeEntry(); + } + + private static void writePositionToRowBuilder(Block block, int position, RowBlockBuilder rowBlockBuilder) + { + if (!(block instanceof AbstractRowBlock)) { + throw new IllegalArgumentException("Expected AbstractRowBlock"); + } + + rowBlockBuilder.beginBlockEntry(); + AbstractRowBlock rowBlock = (AbstractRowBlock) block; + + int offset = rowBlock.getFieldBlockOffset(position); + for (int fieldIndex = 0; fieldIndex < rowBlock.numFields; fieldIndex++) { + BlockBuilder fieldBlockBuilder = rowBlockBuilder.getBlockBuilder(fieldIndex); + Block fieldBlock = rowBlock.getRawFieldBlocks()[fieldIndex]; + if (fieldBlock.isNull(offset)) { + fieldBlockBuilder.appendNull(); + } + else { + writePositionToBlockBuilder(fieldBlock, offset, fieldBlockBuilder); + } + } + rowBlockBuilder.closeEntry(); + } +} diff --git a/presto-common/src/test/java/com/facebook/presto/common/block/TestBlockBuilderUtils.java b/presto-common/src/test/java/com/facebook/presto/common/block/TestBlockBuilderUtils.java new file mode 100644 index 000000000000..bd0b3ab32599 --- /dev/null +++ b/presto-common/src/test/java/com/facebook/presto/common/block/TestBlockBuilderUtils.java @@ -0,0 +1,146 @@ +/* + * 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 com.facebook.presto.common.block; + +import com.facebook.presto.common.type.MapType; +import com.facebook.presto.common.type.RowType; +import com.facebook.presto.common.type.Type; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.common.block.BlockBuilderUtils.writePositionToBlockBuilder; +import static com.facebook.presto.common.type.BigintType.BIGINT; +import static com.facebook.presto.common.type.VarcharType.VARCHAR; +import static org.testng.Assert.assertEquals; + +public class TestBlockBuilderUtils +{ + private static final MapType TEST_MAP_TYPE = new MapType( + BIGINT, + VARCHAR, + MethodHandleUtil.methodHandle(TestBlockBuilderUtils.class, "throwUnsupportedOperation"), + MethodHandleUtil.methodHandle(TestBlockBuilderUtils.class, "throwUnsupportedOperation")); + private static final Map TEST_MAP_VALUES = createTestMap(); + + // Presto Query Engine does not support Map with Null keys. + // Presto ORC reader and Writer are used as library in some other + // projects, and it requires null keys to be supported in the Map. + private static Map createTestMap() + { + Map testMap = new HashMap<>(); + testMap.put(1L, "ONE"); + testMap.put(2L, "TWO"); + testMap.put(null, "NULL"); + testMap.put(4L, null); + return testMap; + } + + @Test + public void testArrayBlockBuilder() + { + long[] values = new long[]{1, 2, 3, 4, 5}; + + ArrayBlockBuilder blockBuilder1 = new ArrayBlockBuilder(BIGINT, null, 1); + BlockBuilder elementBuilder = blockBuilder1.beginBlockEntry(); + for (long value : values) { + BIGINT.writeLong(elementBuilder, value); + } + Block expectedBlock = blockBuilder1.closeEntry().build(); + + // write values to a new block using BlockBuilderUtil + BlockBuilder blockBuilder2 = new ArrayBlockBuilder(BIGINT, null, 1); + writePositionToBlockBuilder(expectedBlock, 0, blockBuilder2); + Block newBlock = blockBuilder2.build(); + assertEquals(newBlock, expectedBlock); + } + + @Test + public void testMapBlockBuilder() + { + BlockBuilder blockBuilder1 = TEST_MAP_TYPE.createBlockBuilder(null, 1); + BlockBuilder mapBlockBuilder = blockBuilder1.beginBlockEntry(); + writeValuesToMapBuilder(mapBlockBuilder); + Block expectedBlock = blockBuilder1.closeEntry().build(); + + // write values to a new block using BlockBuilderUtil + BlockBuilder blockBuilder2 = TEST_MAP_TYPE.createBlockBuilder(null, 1); + writePositionToBlockBuilder(expectedBlock, 0, blockBuilder2); + Block newBlock = blockBuilder2.build(); + assertEquals(newBlock, expectedBlock); + } + + @Test + public void testRowBlockBuilder() + { + RowType rowType = rowType(VARCHAR, BIGINT, TEST_MAP_TYPE); + BlockBuilder blockBuilder = rowType.createBlockBuilder(null, 1); + + BlockBuilder rowBlockBuilder = blockBuilder.beginBlockEntry(); + VARCHAR.writeString(rowBlockBuilder, "TEST_ROW"); + BIGINT.writeLong(rowBlockBuilder, 10L); + BlockBuilder mapBlockBuilder = rowBlockBuilder.beginBlockEntry(); + writeValuesToMapBuilder(mapBlockBuilder); + rowBlockBuilder.closeEntry(); + Block expectedBlock = blockBuilder.closeEntry().build(); + + // write values to a new block using BlockBuilderUtil + BlockBuilder blockBuilder2 = rowType.createBlockBuilder(null, 1); + writePositionToBlockBuilder(expectedBlock, 0, blockBuilder2); + Block newBlock = blockBuilder2.build(); + assertEquals(newBlock, expectedBlock); + } + + private static void writeValuesToMapBuilder(BlockBuilder mapBlockBuilder) + { + for (Map.Entry entry : TEST_MAP_VALUES.entrySet()) { + Long key = entry.getKey(); + if (key == null) { + mapBlockBuilder.appendNull(); + } + else { + BIGINT.writeLong(mapBlockBuilder, entry.getKey()); + } + + String value = entry.getValue(); + if (value == null) { + mapBlockBuilder.appendNull(); + } + else { + VARCHAR.writeString(mapBlockBuilder, entry.getValue()); + } + } + } + + public static RowType rowType(Type... fieldTypes) + { + List fields = new ArrayList<>(); + for (int i = 0; i < fieldTypes.length; i++) { + Type fieldType = fieldTypes[i]; + String fieldName = "field_" + i; + fields.add(new RowType.Field(Optional.of(fieldName), fieldType)); + } + return RowType.from(fields); + } + + // Used via reflection for creating mapType. + public static void throwUnsupportedOperation() + { + throw new UnsupportedOperationException(); + } +}