Examples with BlockBuilder com.facebook.presto.common.block.BlockBuilder used on opensource projects

Example 26 with BlockBuilder

use of com.facebook.presto.common.block.BlockBuilder in project presto by prestodb.

the class TestSliceDictionaryColumnWriter method getDictionaryKeys.

private List<String> getDictionaryKeys(List<String> values, OrcEncoding orcEncoding, boolean sortDictionaryKeys) throws IOException {
    DictionaryColumnWriter writer = getDictionaryColumnWriter(orcEncoding, sortDictionaryKeys);
    for (int index = 0; index < values.size(); ) {
        int endIndex = Math.min(index + 10_000, values.size());
        BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, 10_000);
        while (index < endIndex) {
            VARCHAR.writeSlice(blockBuilder, utf8Slice(values.get(index++)));
        }
        writer.beginRowGroup();
        writer.writeBlock(blockBuilder);
        writer.finishRowGroup();
    }
    writer.close();
    List<StreamDataOutput> streams = writer.getDataStreams();
    int dictionarySize = writer.getColumnEncodings().get(COLUMN_ID).getDictionarySize();
    ByteArrayInputStream dictionaryDataStream = new ByteArrayInputStream(getOrcInputStream(streams, DICTIONARY_DATA));
    LongInputStream dictionaryLengthStream = getDictionaryLengthStream(streams, orcEncoding);
    List<String> dictionaryKeys = new ArrayList<>(dictionarySize);
    for (int i = 0; i < dictionarySize; i++) {
        int length = toIntExact(dictionaryLengthStream.next());
        String dictionaryKey = new String(dictionaryDataStream.next(length), UTF_8);
        dictionaryKeys.add(dictionaryKey);
    }
    return dictionaryKeys;
}

Example 27 with BlockBuilder

use of com.facebook.presto.common.block.BlockBuilder in project presto by prestodb.

the class OrcTester method writeOrcColumnsPresto.

public static void writeOrcColumnsPresto(File outputFile, Format format, CompressionKind compression, Optional<DwrfWriterEncryption> dwrfWriterEncryption, List<Type> types, List<List<?>> values, WriterStats stats) throws Exception {
    OrcWriter writer = createOrcWriter(outputFile, format.orcEncoding, compression, dwrfWriterEncryption, types, OrcWriterOptions.builder().build(), stats);
    Block[] blocks = new Block[types.size()];
    for (int i = 0; i < types.size(); i++) {
        Type type = types.get(i);
        BlockBuilder blockBuilder = type.createBlockBuilder(null, values.size());
        for (Object value : values.get(i)) {
            writeValue(type, blockBuilder, value);
        }
        blocks[i] = blockBuilder.build();
    }
    writer.write(new Page(blocks));
    writer.close();
    writer.validate(new FileOrcDataSource(outputFile, new DataSize(1, MEGABYTE), new DataSize(1, MEGABYTE), new DataSize(1, MEGABYTE), true));
}

Example 28 with BlockBuilder

use of com.facebook.presto.common.block.BlockBuilder in project presto by prestodb.

the class TestOrcWriter method testStreamOrder.

private void testStreamOrder(OrcEncoding encoding, CompressionKind kind, OptionalInt level, StreamLayoutFactory streamLayoutFactory, Supplier<Consumer<Stream>> streamConsumerFactory) throws IOException {
    OrcWriterOptions orcWriterOptions = OrcWriterOptions.builder().withFlushPolicy(DefaultOrcWriterFlushPolicy.builder().withStripeMinSize(new DataSize(0, MEGABYTE)).withStripeMaxSize(new DataSize(32, MEGABYTE)).withStripeMaxRowCount(ORC_STRIPE_SIZE).build()).withRowGroupMaxRowCount(ORC_ROW_GROUP_SIZE).withDictionaryMaxMemory(new DataSize(32, MEGABYTE)).withCompressionLevel(level).withStreamLayoutFactory(streamLayoutFactory).build();
    for (OrcWriteValidationMode validationMode : OrcWriteValidationMode.values()) {
        TempFile tempFile = new TempFile();
        OrcWriter writer = new OrcWriter(new OutputStreamDataSink(new FileOutputStream(tempFile.getFile())), ImmutableList.of("test1", "test2", "test3", "test4", "test5"), ImmutableList.of(VARCHAR, VARCHAR, VARCHAR, VARCHAR, VARCHAR), encoding, kind, Optional.empty(), NO_ENCRYPTION, orcWriterOptions, ImmutableMap.of(), HIVE_STORAGE_TIME_ZONE, true, validationMode, new OrcWriterStats());
        // write down some data with unsorted streams
        String[] data = new String[] { "a", "bbbbb", "ccc", "dd", "eeee" };
        Block[] blocks = new Block[data.length];
        int entries = 65536;
        BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, entries);
        for (int i = 0; i < data.length; i++) {
            byte[] bytes = data[i].getBytes();
            for (int j = 0; j < entries; j++) {
                // force to write different data
                bytes[0] = (byte) ((bytes[0] + 1) % 128);
                blockBuilder.writeBytes(Slices.wrappedBuffer(bytes, 0, bytes.length), 0, bytes.length);
                blockBuilder.closeEntry();
            }
            blocks[i] = blockBuilder.build();
            blockBuilder = blockBuilder.newBlockBuilderLike(null);
        }
        writer.write(new Page(blocks));
        writer.close();
        for (StripeFooter stripeFooter : OrcTester.getStripes(tempFile.getFile(), encoding)) {
            Consumer<Stream> streamConsumer = streamConsumerFactory.get();
            boolean dataStreamStarted = false;
            for (Stream stream : stripeFooter.getStreams()) {
                if (isIndexStream(stream)) {
                    assertFalse(dataStreamStarted);
                    continue;
                }
                dataStreamStarted = true;
                streamConsumer.accept(stream);
            }
        }
    }
}

Example 29 with BlockBuilder

use of com.facebook.presto.common.block.BlockBuilder in project presto by prestodb.

the class TestOrcWriter method testVerifyNoIllegalStateException.

@Test(expectedExceptions = IOException.class, expectedExceptionsMessageRegExp = "Dummy exception from mocked instance")
public void testVerifyNoIllegalStateException() throws IOException {
    OrcWriter writer = new OrcWriter(new MockDataSink(), ImmutableList.of("test1"), ImmutableList.of(VARCHAR), ORC, NONE, Optional.empty(), NO_ENCRYPTION, OrcWriterOptions.builder().withFlushPolicy(DefaultOrcWriterFlushPolicy.builder().withStripeMinSize(new DataSize(0, MEGABYTE)).withStripeMaxSize(new DataSize(32, MEGABYTE)).withStripeMaxRowCount(10).build()).withRowGroupMaxRowCount(ORC_ROW_GROUP_SIZE).withDictionaryMaxMemory(new DataSize(32, MEGABYTE)).build(), ImmutableMap.of(), HIVE_STORAGE_TIME_ZONE, false, null, new OrcWriterStats());
    int entries = 65536;
    BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, entries);
    byte[] bytes = "dummyString".getBytes();
    for (int j = 0; j < entries; j++) {
        // force to write different data
        bytes[0] = (byte) ((bytes[0] + 1) % 128);
        blockBuilder.writeBytes(Slices.wrappedBuffer(bytes, 0, bytes.length), 0, bytes.length);
        blockBuilder.closeEntry();
    }
    Block[] blocks = new Block[] { blockBuilder.build() };
    try {
        // Throw IOException after first flush
        writer.write(new Page(blocks));
    } catch (IOException e) {
        writer.close();
    }
}

Example 30 with BlockBuilder

use of com.facebook.presto.common.block.BlockBuilder in project presto by prestodb.

the class PrestoSparkShufflePageInput method createPage.

private static Page createPage(int rowCount, BasicSliceInput input, List<Type> types) {
    checkArgument(rowCount > 0, "rowCount must be greater than zero: %s", rowCount);
    if (input.length() == 0) {
        // zero column page
        verify(types.isEmpty(), "types is expected to be empty");
        return new Page(rowCount);
    }
    PageBuilder pageBuilder = new PageBuilder(types);
    while (input.isReadable()) {
        pageBuilder.declarePosition();
        for (int channel = 0; channel < types.size(); channel++) {
            BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(channel);
            blockBuilder.readPositionFrom(input);
        }
    }
    Page page = pageBuilder.build();
    verify(page.getPositionCount() == rowCount, "unexpected row count: %s != %s", page.getPositionCount(), rowCount);
    return page;
}