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

public static void output(ArrayType outputType, MinMaxByNState state, BlockBuilder out)
{
    TypedKeyValueHeap heap = state.getTypedKeyValueHeap();
    if (heap == null || heap.isEmpty()) {
        out.appendNull();
        return;
    }

    Type elementType = outputType.getElementType();

    BlockBuilder arrayBlockBuilder = out.beginBlockEntry();
    BlockBuilder reversedBlockBuilder = elementType.createBlockBuilder(null, heap.getCapacity());
    long startSize = heap.getEstimatedSize();
    heap.popAll(reversedBlockBuilder);
    state.addMemoryUsage(heap.getEstimatedSize() - startSize);

    for (int i = reversedBlockBuilder.getPositionCount() - 1; i >= 0; i--) {
        elementType.appendTo(reversedBlockBuilder, i, arrayBlockBuilder);
    }
    out.closeEntry();
}
 

public static void output(ArrayType outputType, MinMaxNState state, BlockBuilder out)
{
    TypedHeap heap = state.getTypedHeap();
    if (heap == null || heap.isEmpty()) {
        out.appendNull();
        return;
    }

    Type elementType = outputType.getElementType();

    BlockBuilder reversedBlockBuilder = elementType.createBlockBuilder(null, heap.getCapacity());
    long startSize = heap.getEstimatedSize();
    heap.popAll(reversedBlockBuilder);
    state.addMemoryUsage(heap.getEstimatedSize() - startSize);

    BlockBuilder arrayBlockBuilder = out.beginBlockEntry();
    for (int i = reversedBlockBuilder.getPositionCount() - 1; i >= 0; i--) {
        elementType.appendTo(reversedBlockBuilder, i, arrayBlockBuilder);
    }
    out.closeEntry();
}
 

@Test
public void testGetElementPosition()
{
    int elementCount = 100;
    // Set initialTypedSetEntryCount to a small number to trigger rehash()
    int initialTypedSetEntryCount = 10;
    TypedSet typedSet = new TypedSet(BIGINT, initialTypedSetEntryCount, FUNCTION_NAME);
    BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(elementCount);
    for (int i = 0; i < elementCount; i++) {
        BIGINT.writeLong(blockBuilder, i);
        typedSet.add(blockBuilder, i);
    }

    assertEquals(typedSet.size(), elementCount);

    for (int j = 0; j < blockBuilder.getPositionCount(); j++) {
        assertEquals(typedSet.positionOf(blockBuilder, j), j);
    }
}
 

public static void output(Type keyType, Type valueType, MultimapAggregationState state, BlockBuilder out)
{
    if (state.isEmpty()) {
        out.appendNull();
    }
    else {
        // TODO: Avoid copy value block associated with the same key by using strategy similar to multimap_from_entries
        ObjectBigArray<BlockBuilder> valueArrayBlockBuilders = new ObjectBigArray<>();
        valueArrayBlockBuilders.ensureCapacity(state.getEntryCount());
        BlockBuilder distinctKeyBlockBuilder = keyType.createBlockBuilder(null, state.getEntryCount(), expectedValueSize(keyType, 100));
        TypedSet keySet = new TypedSet(keyType, state.getEntryCount(), MultimapAggregationFunction.NAME);

        state.forEach((key, value, keyValueIndex) -> {
            // Merge values of the same key into an array
            if (!keySet.contains(key, keyValueIndex)) {
                keySet.add(key, keyValueIndex);
                keyType.appendTo(key, keyValueIndex, distinctKeyBlockBuilder);
                BlockBuilder valueArrayBuilder = valueType.createBlockBuilder(null, 10, expectedValueSize(valueType, EXPECTED_ENTRY_SIZE));
                valueArrayBlockBuilders.set(keySet.positionOf(key, keyValueIndex), valueArrayBuilder);
            }
            valueType.appendTo(value, keyValueIndex, valueArrayBlockBuilders.get(keySet.positionOf(key, keyValueIndex)));
        });

        // Write keys and value arrays into one Block
        Type valueArrayType = new ArrayType(valueType);
        BlockBuilder multimapBlockBuilder = out.beginBlockEntry();
        for (int i = 0; i < distinctKeyBlockBuilder.getPositionCount(); i++) {
            keyType.appendTo(distinctKeyBlockBuilder, i, multimapBlockBuilder);
            valueArrayType.writeObject(multimapBlockBuilder, valueArrayBlockBuilders.get(i).build());
        }
        out.closeEntry();
    }
}
 

@Override
public void forEach(MultimapAggregationStateConsumer consumer)
{
    BlockBuilder keyBlockBuilder = keyBlockBuilders.get(getGroupId());
    BlockBuilder valueBlockBuilder = valueBlockBuilders.get(getGroupId());
    if (keyBlockBuilder == null) {
        return;
    }
    for (int i = 0; i < keyBlockBuilder.getPositionCount(); i++) {
        consumer.accept(keyBlockBuilder, valueBlockBuilder, i);
    }
}
 

@Override
public void forEach(ArrayAggregationStateConsumer consumer)
{
    BlockBuilder blockBuilder = blockBuilders.get(getGroupId());
    if (blockBuilder == null) {
        return;
    }

    for (int i = 0; i < blockBuilder.getPositionCount(); i++) {
        consumer.accept(blockBuilder, i);
    }
}
 

@ScalarFunction
@LiteralParameters({"x", "y"})
@SqlType("array(varchar(x))")
public static Block split(@SqlType("varchar(x)") Slice string, @SqlType("varchar(y)") Slice delimiter, @SqlType(StandardTypes.BIGINT) long limit)
{
    checkCondition(limit > 0, INVALID_FUNCTION_ARGUMENT, "Limit must be positive");
    checkCondition(limit <= Integer.MAX_VALUE, INVALID_FUNCTION_ARGUMENT, "Limit is too large");
    checkCondition(delimiter.length() > 0, INVALID_FUNCTION_ARGUMENT, "The delimiter may not be the empty string");
    BlockBuilder parts = VARCHAR.createBlockBuilder(null, 1, string.length());
    // If limit is one, the last and only element is the complete string
    if (limit == 1) {
        VARCHAR.writeSlice(parts, string);
        return parts.build();
    }

    int index = 0;
    while (index < string.length()) {
        int splitIndex = string.indexOf(delimiter, index);
        // Found split?
        if (splitIndex < 0) {
            break;
        }
        // Add the part from current index to found split
        VARCHAR.writeSlice(parts, string, index, splitIndex - index);
        // Continue searching after delimiter
        index = splitIndex + delimiter.length();
        // Reached limit-1 parts so we can stop
        if (parts.getPositionCount() == limit - 1) {
            break;
        }
    }
    // Rest of string
    VARCHAR.writeSlice(parts, string, index, string.length() - index);

    return parts.build();
}
 

@Test
public void testGetElementPositionWithNull()
{
    int elementCount = 100;
    // Set initialTypedSetEntryCount to a small number to trigger rehash()
    int initialTypedSetEntryCount = 10;
    TypedSet typedSet = new TypedSet(BIGINT, initialTypedSetEntryCount, FUNCTION_NAME);
    BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(elementCount);
    for (int i = 0; i < elementCount; i++) {
        if (i % 10 == 0) {
            blockBuilder.appendNull();
        }
        else {
            BIGINT.writeLong(blockBuilder, i);
        }
        typedSet.add(blockBuilder, i);
    }

    // The internal elementBlock and hashtable of the typedSet should contain
    // all distinct non-null elements plus one null
    assertEquals(typedSet.size(), elementCount - elementCount / 10 + 1);

    int nullCount = 0;
    for (int j = 0; j < blockBuilder.getPositionCount(); j++) {
        // The null is only added to typedSet once, so the internal elementBlock subscript is shifted by nullCountMinusOne
        if (!blockBuilder.isNull(j)) {
            assertEquals(typedSet.positionOf(blockBuilder, j), j - nullCount + 1);
        }
        else {
            // The first null added to typedSet is at position 0
            assertEquals(typedSet.positionOf(blockBuilder, j), 0);
            nullCount++;
        }
    }
}
 

private void testGetElementPositionRandomFor(TypedSet set)
{
    BlockBuilder keys = VARCHAR.createBlockBuilder(null, 5);
    VARCHAR.writeSlice(keys, utf8Slice("hello"));
    VARCHAR.writeSlice(keys, utf8Slice("bye"));
    VARCHAR.writeSlice(keys, utf8Slice("abc"));

    for (int i = 0; i < keys.getPositionCount(); i++) {
        set.add(keys, i);
    }

    BlockBuilder values = VARCHAR.createBlockBuilder(null, 5);
    VARCHAR.writeSlice(values, utf8Slice("bye"));
    VARCHAR.writeSlice(values, utf8Slice("abc"));
    VARCHAR.writeSlice(values, utf8Slice("hello"));
    VARCHAR.writeSlice(values, utf8Slice("bad"));
    values.appendNull();

    assertEquals(set.positionOf(values, 4), -1);
    assertEquals(set.positionOf(values, 2), 0);
    assertEquals(set.positionOf(values, 1), 2);
    assertEquals(set.positionOf(values, 0), 1);
    assertFalse(set.contains(values, 3));

    set.add(values, 4);
    assertTrue(set.contains(values, 4));
}