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

@TypeParameter("E")
@SqlType("array(E)")
public Block reverse(
        @TypeParameter("E") Type type,
        @SqlType("array(E)") Block block)
{
    int arrayLength = block.getPositionCount();

    if (arrayLength < 2) {
        return block;
    }

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int i = arrayLength - 1; i >= 0; i--) {
        type.appendTo(block, i, blockBuilder);
    }
    pageBuilder.declarePositions(arrayLength);

    return blockBuilder.getRegion(blockBuilder.getPositionCount() - arrayLength, arrayLength);
}
 

@SqlType("array(bigint)")
public Block bigintDistinct(@SqlType("array(bigint)") Block array)
{
    if (array.getPositionCount() == 0) {
        return array;
    }

    boolean containsNull = false;
    LongSet set = new LongOpenHashSet(array.getPositionCount());
    int distinctCount = 0;

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder distinctElementBlockBuilder = pageBuilder.getBlockBuilder(0);
    for (int i = 0; i < array.getPositionCount(); i++) {
        if (array.isNull(i)) {
            if (!containsNull) {
                containsNull = true;
                distinctElementBlockBuilder.appendNull();
                distinctCount++;
            }
            continue;
        }
        long value = BIGINT.getLong(array, i);
        if (!set.contains(value)) {
            set.add(value);
            distinctCount++;
            BIGINT.appendTo(array, i, distinctElementBlockBuilder);
        }
    }

    pageBuilder.declarePositions(distinctCount);

    return distinctElementBlockBuilder.getRegion(distinctElementBlockBuilder.getPositionCount() - distinctCount, distinctCount);
}
 

@UsedByGeneratedCode
public static Block concat(Type elementType, Object state, Block[] blocks)
{
    int resultPositionCount = 0;

    // fast path when there is at most one non empty block
    Block nonEmptyBlock = null;
    for (int i = 0; i < blocks.length; i++) {
        resultPositionCount += blocks[i].getPositionCount();
        if (blocks[i].getPositionCount() > 0) {
            nonEmptyBlock = blocks[i];
        }
    }
    if (nonEmptyBlock == null) {
        return blocks[0];
    }
    if (resultPositionCount == nonEmptyBlock.getPositionCount()) {
        return nonEmptyBlock;
    }

    PageBuilder pageBuilder = (PageBuilder) state;
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int blockIndex = 0; blockIndex < blocks.length; blockIndex++) {
        Block block = blocks[blockIndex];
        for (int i = 0; i < block.getPositionCount(); i++) {
            elementType.appendTo(block, i, blockBuilder);
        }
    }
    pageBuilder.declarePositions(resultPositionCount);
    return blockBuilder.getRegion(blockBuilder.getPositionCount() - resultPositionCount, resultPositionCount);
}
 

@TypeParameter("E")
@SqlType("array(E)")
public Block shuffle(
        @TypeParameter("E") Type type,
        @SqlType("array(E)") Block block)
{
    int length = block.getPositionCount();
    if (positions.length < length) {
        positions = new int[length];
    }
    for (int i = 0; i < length; i++) {
        positions[i] = i;
    }

    // Fisher-Yates shuffle
    // Randomly swap a pair of positions
    for (int i = length - 1; i > 0; i--) {
        int index = ThreadLocalRandom.current().nextInt(i + 1);
        int swap = positions[i];
        positions[i] = positions[index];
        positions[index] = swap;
    }

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    for (int i = 0; i < length; i++) {
        type.appendTo(block, positions[i], blockBuilder);
    }
    pageBuilder.declarePositions(length);

    return blockBuilder.getRegion(blockBuilder.getPositionCount() - length, length);
}
 

private Block computeResultBlock(Type type, Block block, int arrayLength)
{
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    for (int i = 0; i < arrayLength; ++i) {
        type.appendTo(block, positions.get(i), blockBuilder);
    }
    pageBuilder.declarePositions(arrayLength);

    return blockBuilder.getRegion(blockBuilder.getPositionCount() - arrayLength, arrayLength);
}
 

@LiteralParameters("x")
@SqlType("varchar")
@SqlNullable
public Slice regexpReplace(
        @SqlType("varchar") Slice source,
        @SqlType(JoniRegexpType.NAME) JoniRegexp pattern,
        @SqlType("function(array(varchar), varchar(x))") UnaryFunctionInterface replaceFunction)
{
    // If there is no match we can simply return the original source without doing copy.
    Matcher matcher = pattern.matcher(source.getBytes());
    if (matcher.search(0, source.length(), Option.DEFAULT) == -1) {
        return source;
    }

    SliceOutput output = new DynamicSliceOutput(source.length());

    // Prepare a BlockBuilder that will be used to create the target block
    // that will be passed to the lambda function.
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    int groupCount = pattern.regex().numberOfCaptures();
    int appendPosition = 0;
    int nextStart;

    do {
        // nextStart is the same as the last appendPosition, unless the last match was zero-width.
        if (matcher.getEnd() == matcher.getBegin()) {
            if (matcher.getBegin() < source.length()) {
                nextStart = matcher.getEnd() + lengthOfCodePointFromStartByte(source.getByte(matcher.getBegin()));
            }
            else {
                // last match is empty and we matched end of source, move past the source length to terminate the loop
                nextStart = matcher.getEnd() + 1;
            }
        }
        else {
            nextStart = matcher.getEnd();
        }

        // Append the un-matched part
        Slice unmatched = source.slice(appendPosition, matcher.getBegin() - appendPosition);
        appendPosition = matcher.getEnd();
        output.appendBytes(unmatched);

        // Append the capturing groups to the target block that will be passed to lambda
        Region matchedRegion = matcher.getEagerRegion();
        for (int i = 1; i <= groupCount; i++) {
            // Add to the block builder if the matched region is not null. In Joni null is represented as [-1, -1]
            if (matchedRegion.beg[i] >= 0 && matchedRegion.end[i] >= 0) {
                VARCHAR.writeSlice(blockBuilder, source, matchedRegion.beg[i], matchedRegion.end[i] - matchedRegion.beg[i]);
            }
            else {
                blockBuilder.appendNull();
            }
        }
        pageBuilder.declarePositions(groupCount);
        Block target = blockBuilder.getRegion(blockBuilder.getPositionCount() - groupCount, groupCount);

        // Call the lambda function to replace the block, and append the result to output
        Slice replaced = (Slice) replaceFunction.apply(target);
        if (replaced == null) {
            // replacing a substring with null (unknown) makes the entire string null
            return null;
        }
        output.appendBytes(replaced);
    }
    while (matcher.search(nextStart, source.length(), Option.DEFAULT) != -1);

    // Append the last un-matched part
    output.writeBytes(source, appendPosition, source.length() - appendPosition);
    return output.slice();
}
 

@TypeParameter("E")
@SqlType("array(E)")
public Block sort(
        @OperatorDependency(operator = LESS_THAN, argumentTypes = {"E", "E"}) MethodHandle lessThanFunction,
        @TypeParameter("E") Type type,
        @SqlType("array(E)") Block block)
{
    int arrayLength = block.getPositionCount();
    if (positions.size() < arrayLength) {
        positions = Ints.asList(new int[arrayLength]);
    }
    for (int i = 0; i < arrayLength; i++) {
        positions.set(i, i);
    }

    Collections.sort(positions.subList(0, arrayLength), new Comparator<Integer>()
    {
        @Override
        public int compare(Integer p1, Integer p2)
        {
            boolean nullLeft = block.isNull(p1);
            boolean nullRight = block.isNull(p2);
            if (nullLeft && nullRight) {
                return 0;
            }
            if (nullLeft) {
                return 1;
            }
            if (nullRight) {
                return -1;
            }

            //TODO: This could be quite slow, it should use parametric equals
            return type.compareTo(block, p1, block, p2);
        }
    });

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    for (int i = 0; i < arrayLength; i++) {
        type.appendTo(block, positions.get(i), blockBuilder);
    }
    pageBuilder.declarePositions(arrayLength);

    return blockBuilder.getRegion(blockBuilder.getPositionCount() - arrayLength, arrayLength);
}
 

@LiteralParameters("x")
@SqlType("varchar")
@SqlNullable
public Slice regexpReplace(
        @SqlType("varchar") Slice source,
        @SqlType(Re2JRegexpType.NAME) Re2JRegexp pattern,
        @SqlType("function(array(varchar), varchar(x))") UnaryFunctionInterface replaceFunction)
{
    // If there is no match we can simply return the original source without doing copy.
    Matcher matcher = pattern.matcher(source);
    if (!matcher.find()) {
        return source;
    }

    SliceOutput output = new DynamicSliceOutput(source.length());

    // Prepare a BlockBuilder that will be used to create the target block
    // that will be passed to the lambda function.
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    int groupCount = matcher.groupCount();
    int appendPosition = 0;

    do {
        int start = matcher.start();
        int end = matcher.end();

        // Append the un-matched part
        if (appendPosition < start) {
            output.writeBytes(source, appendPosition, start - appendPosition);
        }
        appendPosition = end;

        // Append the capturing groups to the target block that will be passed to lambda
        for (int i = 1; i <= groupCount; i++) {
            Slice matchedGroupSlice = matcher.group(i);
            if (matchedGroupSlice != null) {
                VARCHAR.writeSlice(blockBuilder, matchedGroupSlice);
            }
            else {
                blockBuilder.appendNull();
            }
        }
        pageBuilder.declarePositions(groupCount);
        Block target = blockBuilder.getRegion(blockBuilder.getPositionCount() - groupCount, groupCount);

        // Call the lambda function to replace the block, and append the result to output
        Slice replaced = (Slice) replaceFunction.apply(target);
        if (replaced == null) {
            // replacing a substring with null (unknown) makes the entire string null
            return null;
        }
        output.appendBytes(replaced);
    }
    while (matcher.find());

    // Append the rest of un-matched
    output.writeBytes(source, appendPosition, source.length() - appendPosition);
    return output.slice();
}
 

@TypeParameter("E")
@SqlType("array(E)")
public Block intersect(
        @TypeParameter("E") Type type,
        @OperatorDependency(
                operator = IS_DISTINCT_FROM,
                argumentTypes = {"E", "E"},
                convention = @Convention(arguments = {BLOCK_POSITION, BLOCK_POSITION}, result = FAIL_ON_NULL)) MethodHandle elementIsDistinctFrom,
        @SqlType("array(E)") Block leftArray,
        @SqlType("array(E)") Block rightArray)
{
    if (leftArray.getPositionCount() < rightArray.getPositionCount()) {
        Block tempArray = leftArray;
        leftArray = rightArray;
        rightArray = tempArray;
    }

    int leftPositionCount = leftArray.getPositionCount();
    int rightPositionCount = rightArray.getPositionCount();

    if (rightPositionCount == 0) {
        return rightArray;
    }

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    TypedSet rightTypedSet = new TypedSet(type, elementIsDistinctFrom, rightPositionCount, "array_intersect");
    for (int i = 0; i < rightPositionCount; i++) {
        rightTypedSet.add(rightArray, i);
    }

    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);

    // The intersected set can have at most rightPositionCount elements
    TypedSet intersectTypedSet = new TypedSet(type, Optional.of(elementIsDistinctFrom), blockBuilder, rightPositionCount, "array_intersect");
    for (int i = 0; i < leftPositionCount; i++) {
        if (rightTypedSet.contains(leftArray, i)) {
            intersectTypedSet.add(leftArray, i);
        }
    }

    pageBuilder.declarePositions(intersectTypedSet.size());

    return blockBuilder.getRegion(blockBuilder.getPositionCount() - intersectTypedSet.size(), intersectTypedSet.size());
}