Java Code Examples for io.prestosql.spi.block.Block#getRegion()

@Override
public void addInput(Page page)
{
    checkState(needsInput());

    if (page.getPositionCount() <= remainingLimit) {
        remainingLimit -= page.getPositionCount();
        nextPage = page;
    }
    else {
        Block[] blocks = new Block[page.getChannelCount()];
        for (int channel = 0; channel < page.getChannelCount(); channel++) {
            Block block = page.getBlock(channel);
            blocks[channel] = block.getRegion(0, (int) remainingLimit);
        }
        nextPage = new Page((int) remainingLimit, blocks);
        remainingLimit = 0;
    }
}
 

private static void verifyBlock(Block block, Slice[][][] expectedValues)
{
    assertBlock(block, expectedValues);

    assertColumnarMap(block, expectedValues);
    assertDictionaryBlock(block, expectedValues);
    assertRunLengthEncodedBlock(block, expectedValues);

    int offset = 1;
    int length = expectedValues.length - 2;
    Block blockRegion = block.getRegion(offset, length);
    Slice[][][] expectedValuesRegion = Arrays.copyOfRange(expectedValues, offset, offset + length);

    assertBlock(blockRegion, expectedValuesRegion);

    assertColumnarMap(blockRegion, expectedValuesRegion);
    assertDictionaryBlock(blockRegion, expectedValuesRegion);
    assertRunLengthEncodedBlock(blockRegion, expectedValuesRegion);
}
 

private static <T> void verifyBlock(Block block, T[] expectedValues)
{
    assertBlock(block, expectedValues);

    assertColumnarRow(block, expectedValues);
    assertDictionaryBlock(block, expectedValues);
    assertRunLengthEncodedBlock(block, expectedValues);

    int offset = 1;
    int length = expectedValues.length - 2;
    Block blockRegion = block.getRegion(offset, length);
    T[] expectedValuesRegion = Arrays.copyOfRange(expectedValues, offset, offset + length);

    assertBlock(blockRegion, expectedValuesRegion);

    assertColumnarRow(blockRegion, expectedValuesRegion);
    assertDictionaryBlock(blockRegion, expectedValuesRegion);
    assertRunLengthEncodedBlock(blockRegion, expectedValuesRegion);
}
 

private static <T> void verifyBlock(Block block, T[] expectedValues)
{
    assertBlock(block, expectedValues);

    assertColumnarArray(block, expectedValues);
    assertDictionaryBlock(block, expectedValues);
    assertRunLengthEncodedBlock(block, expectedValues);

    int offset = 1;
    int length = expectedValues.length - 2;
    Block blockRegion = block.getRegion(offset, length);
    T[] expectedValuesRegion = Arrays.copyOfRange(expectedValues, offset, offset + length);

    assertBlock(blockRegion, expectedValuesRegion);

    assertColumnarArray(blockRegion, expectedValuesRegion);
    assertDictionaryBlock(blockRegion, expectedValuesRegion);
    assertRunLengthEncodedBlock(blockRegion, expectedValuesRegion);
}
 

private boolean writeDictionaryRowGroup(Block dictionary, int valueCount, IntBigArray dictionaryIndexes, int maxDirectBytes)
{
    int[][] segments = dictionaryIndexes.getSegments();
    for (int i = 0; valueCount > 0 && i < segments.length; i++) {
        int[] segment = segments[i];
        int positionCount = Math.min(valueCount, segment.length);
        Block block = new DictionaryBlock(positionCount, dictionary, segment);

        while (block != null) {
            int chunkPositionCount = block.getPositionCount();
            Block chunk = block.getRegion(0, chunkPositionCount);

            // avoid chunk with huge logical size
            while (chunkPositionCount > 1 && chunk.getLogicalSizeInBytes() > DIRECT_CONVERSION_CHUNK_MAX_LOGICAL_BYTES) {
                chunkPositionCount /= 2;
                chunk = chunk.getRegion(0, chunkPositionCount);
            }

            directColumnWriter.writeBlock(chunk);
            if (directColumnWriter.getBufferedBytes() > maxDirectBytes) {
                return false;
            }

            // slice block to only unconverted rows
            if (chunkPositionCount < block.getPositionCount()) {
                block = block.getRegion(chunkPositionCount, block.getPositionCount() - chunkPositionCount);
            }
            else {
                block = null;
            }
        }

        valueCount -= positionCount;
    }
    checkState(valueCount == 0);
    return true;
}
 

@Override
public Work<Block> project(ConnectorSession session, DriverYieldSignal yieldSignal, Page page, SelectedPositions selectedPositions)
{
    Block block = requireNonNull(page, "page is null").getBlock(0);
    requireNonNull(selectedPositions, "selectedPositions is null");

    // TODO: make it lazy when MergePages have better merging heuristics for small lazy pages
    if (selectedPositions.isList()) {
        block = block.copyPositions(selectedPositions.getPositions(), selectedPositions.getOffset(), selectedPositions.size());
    }
    else {
        block = block.getRegion(selectedPositions.getOffset(), selectedPositions.size());
    }
    return new CompletedWork<>(block);
}
 

public Page build(JoinProbe probe)
{
    int[] probeIndices = probeIndexBuilder.toIntArray();
    int length = probeIndices.length;
    verify(buildPageBuilder.getPositionCount() == length);

    int[] probeOutputChannels = probe.getOutputChannels();
    Block[] blocks = new Block[probeOutputChannels.length + buildOutputChannelCount];
    for (int i = 0; i < probeOutputChannels.length; i++) {
        Block probeBlock = probe.getPage().getBlock(probeOutputChannels[i]);
        if (!isSequentialProbeIndices || length == 0) {
            blocks[i] = probeBlock.getPositions(probeIndices, 0, probeIndices.length);
        }
        else if (length == probeBlock.getPositionCount()) {
            // probeIndices are a simple covering of the block
            verify(probeIndices[0] == 0);
            verify(probeIndices[length - 1] == length - 1);
            blocks[i] = probeBlock;
        }
        else {
            // probeIndices are sequential without holes
            verify(probeIndices[length - 1] - probeIndices[0] == length - 1);
            blocks[i] = probeBlock.getRegion(probeIndices[0], length);
        }
    }

    Page buildPage = buildPageBuilder.build();
    int offset = probeOutputChannels.length;
    for (int i = 0; i < buildOutputChannelCount; i++) {
        blocks[offset + i] = buildPage.getBlock(i);
    }
    return new Page(buildPageBuilder.getPositionCount(), blocks);
}
 

@TypeParameter("E")
@SqlType("array(E)")
public static Block slice(
        @TypeParameter("E") Type type,
        @SqlType("array(E)") Block array,
        @SqlType(StandardTypes.BIGINT) long fromIndex,
        @SqlType(StandardTypes.BIGINT) long length)
{
    checkCondition(length >= 0, INVALID_FUNCTION_ARGUMENT, "length must be greater than or equal to 0");
    checkCondition(fromIndex != 0, INVALID_FUNCTION_ARGUMENT, "SQL array indices start at 1");

    int size = array.getPositionCount();
    if (size == 0) {
        return array;
    }

    if (fromIndex < 0) {
        fromIndex = size + fromIndex + 1;
    }

    long toIndex = Math.min(fromIndex + length, size + 1);

    if (fromIndex >= toIndex || fromIndex < 1) {
        return type.createBlockBuilder(null, 0).build();
    }

    return array.getRegion((int) (fromIndex - 1), (int) (toIndex - fromIndex));
}
 

@Test
public void testTrimmedBlocks()
{
    int[] unnestedLengths = {2, 1, 2, 3, 1};

    Slice[][][] elements = column(
            array(toSlices("0.0.0"), toSlices("0.1.0")),
            array(toSlices("1.0.0")),
            array(toSlices("2.0.0"), toSlices("2.1.0")),
            array(toSlices("3.0.0"), toSlices("3.1.0"), toSlices("3.2.0")),
            array(toSlices("4.0.0")));

    Slice[][] slices = getFieldElements(elements, 0);
    Block arrayBlock = createArrayBlock(slices);

    // Remove the first element from the arrayBlock for testing
    int startElement = 1;

    Slice[][] truncatedSlices = Arrays.copyOfRange(slices, startElement, slices.length - startElement + 1);
    int[] truncatedUnnestedLengths = Arrays.copyOfRange(unnestedLengths, startElement, slices.length - startElement + 1);
    Block truncatedBlock = arrayBlock.getRegion(startElement, truncatedSlices.length);
    assertBlock(truncatedBlock, truncatedSlices);

    Unnester arrayUnnester = new ArrayUnnester(VARCHAR);
    arrayUnnester.resetInput(truncatedBlock);

    arrayUnnester.startNewOutput(new PageBuilderStatus(), 20);

    // Process all input entries in the truncated block
    for (int i = 0; i < truncatedBlock.getPositionCount(); i++) {
        arrayUnnester.processCurrentAndAdvance(truncatedUnnestedLengths[i]);
    }

    Block[] output = arrayUnnester.buildOutputBlocksAndFlush();
    assertEquals(Arrays.asList(truncatedSlices).stream().mapToInt(slice -> slice.length).sum(), output[0].getPositionCount());

    Slice[] expectedOutput = computeExpectedUnnestedOutput(truncatedSlices, truncatedUnnestedLengths, 0, truncatedSlices.length);
    assertBlock(output[0], expectedOutput);
}
 

@Test
public void testTrimmedBlocks()
{
    int[] unnestedLengths = {1, 2, 1};

    Slice[][][] elements = column(
            array(toSlices("0.0.0", "0.0.1")),
            array(toSlices("1.0.0", "1.0.1"), toSlices("1.1.0", "1.1.1")),
            array(toSlices("2.0.0", "2.0.1")));

    Block mapBlock = createBlockBuilderWithValues(elements).build();

    // Remove the first element from the arrayBlock for testing
    int startElement = 1;

    Slice[][][] truncatedSlices = Arrays.copyOfRange(elements, startElement, elements.length - startElement + 1);
    int[] truncatedUnnestedLengths = Arrays.copyOfRange(unnestedLengths, startElement, elements.length - startElement + 1);
    Block truncatedBlock = mapBlock.getRegion(startElement, elements.length - startElement);
    assertBlock(truncatedBlock, truncatedSlices);

    Unnester mapUnnester = new MapUnnester(VARCHAR, VARCHAR);
    mapUnnester.resetInput(truncatedBlock);

    mapUnnester.startNewOutput(new PageBuilderStatus(), 20);

    // Process all input entries in the truncated block
    for (int i = 0; i < truncatedBlock.getPositionCount(); i++) {
        mapUnnester.processCurrentAndAdvance(truncatedUnnestedLengths[i]);
    }

    Block[] output = mapUnnester.buildOutputBlocksAndFlush();
    assertEquals(Arrays.asList(truncatedSlices).stream().mapToInt(slice -> slice.length).sum(), output[0].getPositionCount());

    Slice[] expectedOutput0 = computeExpectedUnnestedOutput(getFieldElements(truncatedSlices, 0), truncatedUnnestedLengths, 0, truncatedSlices.length);
    assertBlock(output[0], expectedOutput0);

    Slice[] expectedOutput1 = computeExpectedUnnestedOutput(getFieldElements(truncatedSlices, 1), truncatedUnnestedLengths, 0, truncatedSlices.length);
    assertBlock(output[1], expectedOutput1);
}
 

@Test
public void testTrimmedBlocks()
{
    int fieldCount = 3;
    int[] unnestedLengths = {1, 2, 1};

    Slice[][][] elements = column(
            array(toSlices("0.0.0", "0.0.1", "0.0.2")),
            array(toSlices("1.0.0", "1.0.1", "1.0.2"), toSlices("1.1.0", "1.1.1", "1.1.2")),
            array(toSlices("2.0.0", "2.0.1", "2.0.2")));

    RowType rowType = RowType.anonymous(Collections.nCopies(fieldCount, VARCHAR));
    Block arrayOfRowBlock = createArrayBlockOfRowBlocks(elements, rowType);

    // Remove the first element from the arrayBlock for testing
    int startElement = 1;

    Slice[][][] truncatedSlices = Arrays.copyOfRange(elements, startElement, elements.length - startElement + 1);
    int[] truncatedUnnestedLengths = Arrays.copyOfRange(unnestedLengths, startElement, elements.length - startElement + 1);
    Block truncatedBlock = arrayOfRowBlock.getRegion(startElement, truncatedSlices.length);

    Unnester arrayOfRowsUnnester = new ArrayOfRowsUnnester(rowType);
    arrayOfRowsUnnester.resetInput(truncatedBlock);

    arrayOfRowsUnnester.startNewOutput(new PageBuilderStatus(), 20);

    // Process all input entries in the truncated block
    for (int i = 0; i < truncatedBlock.getPositionCount(); i++) {
        arrayOfRowsUnnester.processCurrentAndAdvance(truncatedUnnestedLengths[i]);
    }

    Block[] output = arrayOfRowsUnnester.buildOutputBlocksAndFlush();
    assertEquals(Arrays.asList(truncatedSlices).stream().mapToInt(slice -> slice.length).sum(), output[0].getPositionCount());

    for (int i = 0; i < fieldCount; i++) {
        Slice[] expectedOutput = computeExpectedUnnestedOutput(getFieldElements(truncatedSlices, i), truncatedUnnestedLengths, 0, truncatedSlices.length);
        assertBlock(output[i], expectedOutput);
    }
}
 

public static RunLengthEncodedBlock createTestRleBlock(Block block, int position)
{
    return new RunLengthEncodedBlock(block.getRegion(position, 1), 10);
}