Java Code Examples for io.prestosql.spi.type.Type#getObject()

static Block getBlockFromMap(Type mapType, Map<?, ?> map)
{
    // on functions like COUNT() the Type won't be a MapType
    if (!(mapType instanceof MapType)) {
        return null;
    }
    Type keyType = mapType.getTypeParameters().get(0);
    Type valueType = mapType.getTypeParameters().get(1);

    BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1);
    BlockBuilder builder = mapBlockBuilder.beginBlockEntry();

    for (Map.Entry<?, ?> entry : map.entrySet()) {
        writeObject(builder, keyType, entry.getKey());
        writeObject(builder, valueType, entry.getValue());
    }

    mapBlockBuilder.closeEntry();
    return (Block) mapType.getObject(mapBlockBuilder, 0);
}
 

/**
 * Encodes the given map into a Block.
 *
 * @param mapType Presto type of the map
 * @param map Map of key/value pairs to encode
 * @return Presto Block
 */
static Block getBlockFromMap(Type mapType, Map<?, ?> map)
{
    Type keyType = mapType.getTypeParameters().get(0);
    Type valueType = mapType.getTypeParameters().get(1);

    BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1);
    BlockBuilder builder = mapBlockBuilder.beginBlockEntry();

    for (Entry<?, ?> entry : map.entrySet()) {
        writeObject(builder, keyType, entry.getKey());
        writeObject(builder, valueType, entry.getValue());
    }

    mapBlockBuilder.closeEntry();
    return (Block) mapType.getObject(mapBlockBuilder, 0);
}
 

public static Block flatten(Type type, Type arrayType, Block array)
{
    if (array.getPositionCount() == 0) {
        return type.createBlockBuilder(null, 0).build();
    }

    BlockBuilder builder = type.createBlockBuilder(null, array.getPositionCount(), toIntExact(array.getSizeInBytes() / array.getPositionCount()));
    for (int i = 0; i < array.getPositionCount(); i++) {
        if (!array.isNull(i)) {
            Block subArray = (Block) arrayType.getObject(array, i);
            for (int j = 0; j < subArray.getPositionCount(); j++) {
                type.appendTo(subArray, j, builder);
            }
        }
    }
    return builder.build();
}
 

@ScalarFunction
@SqlType("map(bigint,smallint)")
public static Block hashCounts(@TypeParameter("map(bigint,smallint)") Type mapType, @SqlType(SetDigestType.NAME) Slice slice)
{
    SetDigest digest = SetDigest.newInstance(slice);

    // Maybe use static BlockBuilderStatus in order avoid `new`?
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 1);
    BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry();
    for (Map.Entry<Long, Short> entry : digest.getHashCounts().entrySet()) {
        BIGINT.writeLong(singleMapBlockBuilder, entry.getKey());
        SMALLINT.writeLong(singleMapBlockBuilder, entry.getValue());
    }
    blockBuilder.closeEntry();

    return (Block) mapType.getObject(blockBuilder, 0);
}
 

private static Object getNativeContainerValue(Type type, Block block, int position)
{
    if (block.isNull(position)) {
        return null;
    }
    else if (type.getJavaType() == boolean.class) {
        return type.getBoolean(block, position);
    }
    else if (type.getJavaType() == long.class) {
        return type.getLong(block, position);
    }
    else if (type.getJavaType() == double.class) {
        return type.getDouble(block, position);
    }
    else if (type.getJavaType() == Slice.class) {
        return type.getSlice(block, position);
    }
    else if (type.getJavaType() == Block.class) {
        return type.getObject(block, position);
    }
    else {
        throw new AssertionError("Unimplemented type: " + type);
    }
}
 

@Override
public Object getObject(int field)
{
    checkState(position >= 0, "Not yet advanced");
    checkState(position < page.getPositionCount(), "Already finished");
    Type type = types.get(field);
    return type.getObject(page.getBlock(field), position);
}
 

@TypeParameter("T")
@SqlType("T")
@SqlNullable
public static Block blockArrayMin(
        @OperatorDependency(operator = LESS_THAN, argumentTypes = {"T", "T"}) MethodHandle compareMethodHandle,
        @TypeParameter("T") Type elementType,
        @SqlType("array(T)") Block block)
{
    try {
        if (block.getPositionCount() == 0) {
            return null;
        }

        Block selectedValue = (Block) elementType.getObject(block, 0);
        for (int i = 0; i < block.getPositionCount(); i++) {
            if (block.isNull(i)) {
                return null;
            }
            Block value = (Block) elementType.getObject(block, i);
            if ((boolean) compareMethodHandle.invokeExact(value, selectedValue)) {
                selectedValue = value;
            }
        }

        return selectedValue;
    }
    catch (Throwable t) {
        throw internalError(t);
    }
}
 

@TypeParameter("T")
@SqlType("T")
@SqlNullable
public static Block blockArrayMax(
        @OperatorDependency(operator = GREATER_THAN, argumentTypes = {"T", "T"}) MethodHandle compareMethodHandle,
        @TypeParameter("T") Type elementType,
        @SqlType("array(T)") Block block)
{
    try {
        if (block.getPositionCount() == 0) {
            return null;
        }

        Block selectedValue = (Block) elementType.getObject(block, 0);
        for (int i = 0; i < block.getPositionCount(); i++) {
            if (block.isNull(i)) {
                return null;
            }
            Block value = (Block) elementType.getObject(block, i);
            if ((boolean) compareMethodHandle.invokeExact(value, selectedValue)) {
                selectedValue = value;
            }
        }

        return selectedValue;
    }
    catch (Throwable t) {
        throw internalError(t);
    }
}
 

@TypeParameter("E")
@SqlNullable
@SqlType("E")
public static Block blockElementAt(@TypeParameter("E") Type elementType, @SqlType("array(E)") Block array, @SqlType("bigint") long index)
{
    int position = checkedIndexToBlockPosition(array, index);
    if (position == -1) {
        return null;
    }
    if (array.isNull(position)) {
        return null;
    }

    return (Block) elementType.getObject(array, position);
}
 

@UsedByGeneratedCode
public static Object objectSubscript(Type elementType, Block array, long index)
{
    checkIndex(array, index);
    int position = toIntExact(index - 1);
    if (array.isNull(position)) {
        return null;
    }

    return elementType.getObject(array, position);
}
 

private static Block serializeMap(Type type, BlockBuilder builder, Object object, MapObjectInspector inspector, boolean filterNullMapKeys)
{
    Map<?, ?> map = inspector.getMap(object);
    if (map == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    List<Type> typeParameters = type.getTypeParameters();
    checkArgument(typeParameters.size() == 2, "map must have exactly 2 type parameter");
    Type keyType = typeParameters.get(0);
    Type valueType = typeParameters.get(1);
    ObjectInspector keyInspector = inspector.getMapKeyObjectInspector();
    ObjectInspector valueInspector = inspector.getMapValueObjectInspector();
    BlockBuilder currentBuilder;

    boolean builderSynthesized = false;
    if (builder == null) {
        builderSynthesized = true;
        builder = type.createBlockBuilder(null, 1);
    }
    currentBuilder = builder.beginBlockEntry();

    for (Map.Entry<?, ?> entry : map.entrySet()) {
        // Hive skips map entries with null keys
        if (!filterNullMapKeys || entry.getKey() != null) {
            serializeObject(keyType, currentBuilder, entry.getKey(), keyInspector);
            serializeObject(valueType, currentBuilder, entry.getValue(), valueInspector);
        }
    }

    builder.closeEntry();
    if (builderSynthesized) {
        return (Block) type.getObject(builder, 0);
    }
    else {
        return null;
    }
}
 

private static Block serializeStruct(Type type, BlockBuilder builder, Object object, StructObjectInspector inspector)
{
    if (object == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    List<Type> typeParameters = type.getTypeParameters();
    List<? extends StructField> allStructFieldRefs = inspector.getAllStructFieldRefs();
    checkArgument(typeParameters.size() == allStructFieldRefs.size());
    BlockBuilder currentBuilder;

    boolean builderSynthesized = false;
    if (builder == null) {
        builderSynthesized = true;
        builder = type.createBlockBuilder(null, 1);
    }
    currentBuilder = builder.beginBlockEntry();

    for (int i = 0; i < typeParameters.size(); i++) {
        StructField field = allStructFieldRefs.get(i);
        serializeObject(typeParameters.get(i), currentBuilder, inspector.getStructFieldData(object, field), field.getFieldObjectInspector());
    }

    builder.closeEntry();
    if (builderSynthesized) {
        return (Block) type.getObject(builder, 0);
    }
    else {
        return null;
    }
}
 

private static Block serializeUnion(Type type, BlockBuilder builder, Object object, UnionObjectInspector inspector)
{
    if (object == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    boolean builderSynthesized = false;
    if (builder == null) {
        builderSynthesized = true;
        builder = type.createBlockBuilder(null, 1);
    }

    BlockBuilder currentBuilder = builder.beginBlockEntry();

    byte tag = inspector.getTag(object);
    TINYINT.writeLong(currentBuilder, tag);

    List<Type> typeParameters = type.getTypeParameters();
    for (int i = 1; i < typeParameters.size(); i++) {
        if (i == tag + 1) {
            serializeObject(typeParameters.get(i), currentBuilder, inspector.getField(object), inspector.getObjectInspectors().get(tag));
        }
        else {
            currentBuilder.appendNull();
        }
    }

    builder.closeEntry();
    if (builderSynthesized) {
        return (Block) type.getObject(builder, 0);
    }
    return null;
}
 

@SqlType("map(varchar,varchar)")
public Block splitToMap(@TypeParameter("map(varchar,varchar)") Type mapType, @SqlType(StandardTypes.VARCHAR) Slice string, @SqlType(StandardTypes.VARCHAR) Slice entryDelimiter, @SqlType(StandardTypes.VARCHAR) Slice keyValueDelimiter)
{
    checkCondition(entryDelimiter.length() > 0, INVALID_FUNCTION_ARGUMENT, "entryDelimiter is empty");
    checkCondition(keyValueDelimiter.length() > 0, INVALID_FUNCTION_ARGUMENT, "keyValueDelimiter is empty");
    checkCondition(!entryDelimiter.equals(keyValueDelimiter), INVALID_FUNCTION_ARGUMENT, "entryDelimiter and keyValueDelimiter must not be the same");

    Map<Slice, Slice> map = new HashMap<>();
    int entryStart = 0;
    while (entryStart < string.length()) {
        // Extract key-value pair based on current index
        // then add the pair if it can be split by keyValueDelimiter
        Slice keyValuePair;
        int entryEnd = string.indexOf(entryDelimiter, entryStart);
        if (entryEnd >= 0) {
            keyValuePair = string.slice(entryStart, entryEnd - entryStart);
        }
        else {
            // The rest of the string is the last possible pair.
            keyValuePair = string.slice(entryStart, string.length() - entryStart);
        }

        int keyEnd = keyValuePair.indexOf(keyValueDelimiter);
        if (keyEnd >= 0) {
            int valueStart = keyEnd + keyValueDelimiter.length();
            Slice key = keyValuePair.slice(0, keyEnd);
            Slice value = keyValuePair.slice(valueStart, keyValuePair.length() - valueStart);

            if (value.indexOf(keyValueDelimiter) >= 0) {
                throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key-value delimiter must appear exactly once in each entry. Bad input: '" + keyValuePair.toStringUtf8() + "'");
            }
            if (map.containsKey(key)) {
                throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Duplicate keys (%s) are not allowed", key.toStringUtf8()));
            }

            map.put(key, value);
        }
        else {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key-value delimiter must appear exactly once in each entry. Bad input: '" + keyValuePair.toStringUtf8() + "'");
        }

        if (entryEnd < 0) {
            // No more pairs to add
            break;
        }
        // Next possible pair is placed next to the current entryDelimiter
        entryStart = entryEnd + entryDelimiter.length();
    }

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry();
    for (Map.Entry<Slice, Slice> entry : map.entrySet()) {
        VARCHAR.writeSlice(singleMapBlockBuilder, entry.getKey());
        VARCHAR.writeSlice(singleMapBlockBuilder, entry.getValue());
    }
    blockBuilder.closeEntry();
    pageBuilder.declarePosition();

    return (Block) mapType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1);
}
 

public static Object readBlockValue(BlockEncodingSerde blockEncodingSerde, Type type, Slice slice)
{
    Block block = readBlock(blockEncodingSerde, slice.getInput());
    return type.getObject(block, 0);
}