Search in sources :

Example 11 with TypedSet

use of com.facebook.presto.operator.aggregation.TypedSet in project presto by prestodb.

the class MultimapFromEntriesFunction method multimapFromEntries.

@TypeParameter("K")
@TypeParameter("V")
@SqlType("map(K,array(V))")
@SqlNullable
public Block multimapFromEntries(@TypeParameter("map(K,array(V))") MapType mapType, @SqlType("array(row(K,V))") Block block) {
    Type keyType = mapType.getKeyType();
    Type valueType = ((ArrayType) mapType.getValueType()).getElementType();
    RowType rowType = RowType.anonymous(ImmutableList.of(keyType, valueType));
    int entryCount = block.getPositionCount();
    IntList[] entryIndicesList = new IntList[entryCount];
    for (int i = 0; i < entryIndicesList.length; i++) {
        entryIndicesList[i] = new IntArrayList();
    }
    TypedSet keySet = new TypedSet(keyType, entryCount, NAME);
    for (int i = 0; i < entryCount; i++) {
        if (block.isNull(i)) {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null");
        }
        Block rowBlock = rowType.getObject(block, i);
        if (rowBlock.isNull(0)) {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }
        if (keySet.contains(rowBlock, 0)) {
            entryIndicesList[keySet.positionOf(rowBlock, 0)].add(i);
        } else {
            keySet.add(rowBlock, 0);
            entryIndicesList[keySet.size() - 1].add(i);
        }
    }
    BlockBuilder multimapBlockBuilder = mapType.createBlockBuilder(null, keySet.size());
    BlockBuilder singleMapWriter = multimapBlockBuilder.beginBlockEntry();
    for (int i = 0; i < keySet.size(); i++) {
        keyType.appendTo(rowType.getObject(block, entryIndicesList[i].getInt(0)), 0, singleMapWriter);
        BlockBuilder singleArrayWriter = singleMapWriter.beginBlockEntry();
        for (int entryIndex : entryIndicesList[i]) {
            valueType.appendTo(rowType.getObject(block, entryIndex), 1, singleArrayWriter);
        }
        singleMapWriter.closeEntry();
    }
    multimapBlockBuilder.closeEntry();
    return mapType.getObject(multimapBlockBuilder, multimapBlockBuilder.getPositionCount() - 1);
}
Also used : ArrayType(com.facebook.presto.common.type.ArrayType) MapType(com.facebook.presto.common.type.MapType) ArrayType(com.facebook.presto.common.type.ArrayType) Type(com.facebook.presto.common.type.Type) SqlType(com.facebook.presto.spi.function.SqlType) RowType(com.facebook.presto.common.type.RowType) RowType(com.facebook.presto.common.type.RowType) TypedSet(com.facebook.presto.operator.aggregation.TypedSet) Block(com.facebook.presto.common.block.Block) PrestoException(com.facebook.presto.spi.PrestoException) IntArrayList(it.unimi.dsi.fastutil.ints.IntArrayList) IntList(it.unimi.dsi.fastutil.ints.IntList) BlockBuilder(com.facebook.presto.common.block.BlockBuilder) SqlNullable(com.facebook.presto.spi.function.SqlNullable) TypeParameter(com.facebook.presto.spi.function.TypeParameter) SqlType(com.facebook.presto.spi.function.SqlType)

Example 12 with TypedSet

use of com.facebook.presto.operator.aggregation.TypedSet in project presto by prestodb.

the class ArrayDistinctFunction method distinct.

@TypeParameter("E")
@SqlType("array(E)")
public static Block distinct(@TypeParameter("E") Type type, @SqlType("array(E)") Block array) {
    int arrayLength = array.getPositionCount();
    if (arrayLength < 2) {
        return array;
    }
    if (arrayLength == 2) {
        if (TypeUtils.positionEqualsPosition(type, array, 0, array, 1)) {
            return array.getSingleValueBlock(0);
        } else {
            return array;
        }
    }
    TypedSet typedSet = new TypedSet(type, arrayLength, "array_distinct");
    BlockBuilder distinctElementBlockBuilder;
    if (array.mayHaveNull()) {
        int firstDuplicatePosition = 0;
        // Keep adding the element to the set as long as there are no dupes.
        while (firstDuplicatePosition < arrayLength && typedSet.add(array, firstDuplicatePosition)) {
            firstDuplicatePosition++;
        }
        if (firstDuplicatePosition == arrayLength) {
            // All elements are distinct, so just return the original.
            return array;
        }
        int position = 0;
        distinctElementBlockBuilder = type.createBlockBuilder(null, arrayLength);
        while (position < firstDuplicatePosition) {
            type.appendTo(array, position, distinctElementBlockBuilder);
            position++;
        }
        while (position < arrayLength) {
            if (typedSet.add(array, position)) {
                type.appendTo(array, position, distinctElementBlockBuilder);
            }
            position++;
        }
    } else {
        int firstDuplicatePosition = 0;
        while (firstDuplicatePosition < arrayLength && typedSet.addNonNull(array, firstDuplicatePosition)) {
            firstDuplicatePosition++;
        }
        if (firstDuplicatePosition == arrayLength) {
            // All elements are distinct, so just return the original.
            return array;
        }
        int position = 0;
        distinctElementBlockBuilder = type.createBlockBuilder(null, arrayLength);
        while (position < firstDuplicatePosition) {
            type.appendTo(array, position, distinctElementBlockBuilder);
            position++;
        }
        while (position < arrayLength) {
            if (typedSet.addNonNull(array, position)) {
                type.appendTo(array, position, distinctElementBlockBuilder);
            }
            position++;
        }
    }
    return distinctElementBlockBuilder.build();
}
Also used : TypedSet(com.facebook.presto.operator.aggregation.TypedSet) BlockBuilder(com.facebook.presto.common.block.BlockBuilder) TypeParameter(com.facebook.presto.spi.function.TypeParameter) SqlType(com.facebook.presto.spi.function.SqlType)

Aggregations

TypedSet (com.facebook.presto.operator.aggregation.TypedSet)12 BlockBuilder (com.facebook.presto.common.block.BlockBuilder)6 Block (com.facebook.presto.common.block.Block)5 PrestoException (com.facebook.presto.spi.PrestoException)5 SqlType (com.facebook.presto.spi.function.SqlType)5 Type (com.facebook.presto.common.type.Type)4 TypeParameter (com.facebook.presto.spi.function.TypeParameter)4 BlockBuilder (com.facebook.presto.spi.block.BlockBuilder)3 UsedByGeneratedCode (com.facebook.presto.annotation.UsedByGeneratedCode)2 ArrayType (com.facebook.presto.common.type.ArrayType)2 MapType (com.facebook.presto.common.type.MapType)2 RowType (com.facebook.presto.common.type.RowType)2 Block (com.facebook.presto.spi.block.Block)2 BlockBuilderStatus (com.facebook.presto.spi.block.BlockBuilderStatus)2 InterleavedBlockBuilder (com.facebook.presto.spi.block.InterleavedBlockBuilder)2 SqlNullable (com.facebook.presto.spi.function.SqlNullable)2 ObjectBigArray (com.facebook.presto.common.array.ObjectBigArray)1 CastType (com.facebook.presto.metadata.CastType)1 PageBuilder (com.facebook.presto.spi.PageBuilder)1 InterleavedBlock (com.facebook.presto.spi.block.InterleavedBlock)1