Examples with MapType io.trino.spi.type.MapType used on opensource projects

Example 16 with MapType

use of io.trino.spi.type.MapType in project trino by trinodb.

the class MapToMapCast method specialize.

@Override
public ScalarFunctionImplementation specialize(BoundSignature boundSignature, FunctionDependencies functionDependencies) {
    checkArgument(boundSignature.getArity() == 1, "Expected arity to be 1");
    MapType fromMapType = (MapType) boundSignature.getArgumentType(0);
    Type fromKeyType = fromMapType.getKeyType();
    Type fromValueType = fromMapType.getValueType();
    MapType toMapType = (MapType) boundSignature.getReturnType();
    Type toKeyType = toMapType.getKeyType();
    Type toValueType = toMapType.getValueType();
    MethodHandle keyProcessor = buildProcessor(functionDependencies, fromKeyType, toKeyType, true);
    MethodHandle valueProcessor = buildProcessor(functionDependencies, fromValueType, toValueType, false);
    BlockPositionEqual keyEqual = blockTypeOperators.getEqualOperator(toKeyType);
    BlockPositionHashCode keyHashCode = blockTypeOperators.getHashCodeOperator(toKeyType);
    MethodHandle target = MethodHandles.insertArguments(METHOD_HANDLE, 0, keyProcessor, valueProcessor, toMapType, keyEqual, keyHashCode);
    return new ChoicesScalarFunctionImplementation(boundSignature, NULLABLE_RETURN, ImmutableList.of(NEVER_NULL), target);
}

Example 17 with MapType

use of io.trino.spi.type.MapType in project trino by trinodb.

the class MapTransformKeysFunction method specialize.

@Override
protected ScalarFunctionImplementation specialize(BoundSignature boundSignature) {
    MapType inputMapType = (MapType) boundSignature.getArgumentType(0);
    Type inputKeyType = inputMapType.getKeyType();
    MapType outputMapType = (MapType) boundSignature.getReturnType();
    Type outputKeyType = outputMapType.getKeyType();
    Type valueType = outputMapType.getValueType();
    return new ChoicesScalarFunctionImplementation(boundSignature, FAIL_ON_NULL, ImmutableList.of(NEVER_NULL, FUNCTION), ImmutableList.of(BinaryFunctionInterface.class), generateTransformKey(inputKeyType, outputKeyType, valueType, outputMapType), Optional.of(STATE_FACTORY.bindTo(outputMapType)));
}

Example 18 with MapType

use of io.trino.spi.type.MapType in project trino by trinodb.

the class MapZipWithFunction method mapZipWith.

public static Block mapZipWith(Type keyType, Type leftValueType, Type rightValueType, MapType outputMapType, Object state, Block leftBlock, Block rightBlock, MapZipWithLambda function) {
    SingleMapBlock leftMapBlock = (SingleMapBlock) leftBlock;
    SingleMapBlock rightMapBlock = (SingleMapBlock) rightBlock;
    Type outputValueType = outputMapType.getValueType();
    PageBuilder pageBuilder = (PageBuilder) state;
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry();
    // seekKey() can take non-trivial time when key is complicated value, such as a long VARCHAR or ROW.
    boolean[] keyFound = new boolean[rightMapBlock.getPositionCount()];
    for (int leftKeyPosition = 0; leftKeyPosition < leftMapBlock.getPositionCount(); leftKeyPosition += 2) {
        Object key = readNativeValue(keyType, leftMapBlock, leftKeyPosition);
        Object leftValue = readNativeValue(leftValueType, leftMapBlock, leftKeyPosition + 1);
        int rightValuePosition = rightMapBlock.seekKey(key);
        Object rightValue = null;
        if (rightValuePosition != -1) {
            rightValue = readNativeValue(rightValueType, rightMapBlock, rightValuePosition);
            keyFound[rightValuePosition / 2] = true;
        }
        Object outputValue;
        try {
            outputValue = function.apply(key, leftValue, rightValue);
        } catch (Throwable throwable) {
            // Restore pageBuilder into a consistent state.
            mapBlockBuilder.closeEntry();
            pageBuilder.declarePosition();
            throwIfUnchecked(throwable);
            throw new RuntimeException(throwable);
        }
        keyType.appendTo(leftMapBlock, leftKeyPosition, blockBuilder);
        writeNativeValue(outputValueType, blockBuilder, outputValue);
    }
    // iterate over keys that only exists in rightMapBlock
    for (int rightKeyPosition = 0; rightKeyPosition < rightMapBlock.getPositionCount(); rightKeyPosition += 2) {
        if (!keyFound[rightKeyPosition / 2]) {
            Object key = readNativeValue(keyType, rightMapBlock, rightKeyPosition);
            Object rightValue = readNativeValue(rightValueType, rightMapBlock, rightKeyPosition + 1);
            Object outputValue;
            try {
                outputValue = function.apply(key, null, rightValue);
            } catch (Throwable throwable) {
                // Restore pageBuilder into a consistent state.
                mapBlockBuilder.closeEntry();
                pageBuilder.declarePosition();
                throwIfUnchecked(throwable);
                throw new RuntimeException(throwable);
            }
            keyType.appendTo(rightMapBlock, rightKeyPosition, blockBuilder);
            writeNativeValue(outputValueType, blockBuilder, outputValue);
        }
    }
    mapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return outputMapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
}

Example 19 with MapType

use of io.trino.spi.type.MapType in project trino by trinodb.

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, @OperatorDependency(operator = EQUAL, argumentTypes = { "K", "K" }, convention = @Convention(arguments = { BLOCK_POSITION, BLOCK_POSITION }, result = NULLABLE_RETURN)) BlockPositionEqual keyEqual, @OperatorDependency(operator = HASH_CODE, argumentTypes = "K", convention = @Convention(arguments = BLOCK_POSITION, result = FAIL_ON_NULL)) BlockPositionHashCode keyHashCode, @SqlType("array(row(K,V))") Block mapEntries) {
    Type keyType = mapType.getKeyType();
    Type valueType = ((ArrayType) mapType.getValueType()).getElementType();
    RowType mapEntryType = RowType.anonymous(ImmutableList.of(keyType, valueType));
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    int entryCount = mapEntries.getPositionCount();
    if (entryCount > entryIndicesList.length) {
        initializeEntryIndicesList(entryCount);
    }
    TypedSet keySet = createEqualityTypedSet(keyType, keyEqual, keyHashCode, entryCount, NAME);
    for (int i = 0; i < entryCount; i++) {
        if (mapEntries.isNull(i)) {
            clearEntryIndices(keySet.size());
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null");
        }
        Block mapEntryBlock = mapEntryType.getObject(mapEntries, i);
        if (mapEntryBlock.isNull(0)) {
            clearEntryIndices(keySet.size());
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }
        if (keySet.add(mapEntryBlock, 0)) {
            entryIndicesList[keySet.size() - 1].add(i);
        } else {
            entryIndicesList[keySet.positionOf(mapEntryBlock, 0)].add(i);
        }
    }
    BlockBuilder multimapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder mapWriter = multimapBlockBuilder.beginBlockEntry();
    for (int i = 0; i < keySet.size(); i++) {
        keyType.appendTo(mapEntryType.getObject(mapEntries, entryIndicesList[i].getInt(0)), 0, mapWriter);
        BlockBuilder valuesArray = mapWriter.beginBlockEntry();
        for (int entryIndex : entryIndicesList[i]) {
            valueType.appendTo(mapEntryType.getObject(mapEntries, entryIndex), 1, valuesArray);
        }
        mapWriter.closeEntry();
    }
    multimapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    clearEntryIndices(keySet.size());
    return mapType.getObject(multimapBlockBuilder, multimapBlockBuilder.getPositionCount() - 1);
}

Example 20 with MapType

use of io.trino.spi.type.MapType in project trino by trinodb.

the class MapConcatFunction method mapConcat.

@UsedByGeneratedCode
public static Block mapConcat(MapType mapType, BlockPositionEqual keyEqual, BlockPositionHashCode keyHashCode, Object state, Block[] maps) {
    int entries = 0;
    int lastMapIndex = maps.length - 1;
    int firstMapIndex = lastMapIndex;
    for (int i = 0; i < maps.length; i++) {
        entries += maps[i].getPositionCount();
        if (maps[i].getPositionCount() > 0) {
            lastMapIndex = i;
            firstMapIndex = min(firstMapIndex, i);
        }
    }
    if (lastMapIndex == firstMapIndex) {
        return maps[lastMapIndex];
    }
    PageBuilder pageBuilder = (PageBuilder) state;
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    // TODO: we should move TypedSet into user state as well
    Type keyType = mapType.getKeyType();
    Type valueType = mapType.getValueType();
    TypedSet typedSet = createEqualityTypedSet(keyType, keyEqual, keyHashCode, entries / 2, FUNCTION_NAME);
    BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry();
    // the last map
    Block map = maps[lastMapIndex];
    for (int i = 0; i < map.getPositionCount(); i += 2) {
        typedSet.add(map, i);
        keyType.appendTo(map, i, blockBuilder);
        valueType.appendTo(map, i + 1, blockBuilder);
    }
    // the map between the last and the first
    for (int idx = lastMapIndex - 1; idx > firstMapIndex; idx--) {
        map = maps[idx];
        for (int i = 0; i < map.getPositionCount(); i += 2) {
            if (typedSet.add(map, i)) {
                keyType.appendTo(map, i, blockBuilder);
                valueType.appendTo(map, i + 1, blockBuilder);
            }
        }
    }
    // the first map
    map = maps[firstMapIndex];
    for (int i = 0; i < map.getPositionCount(); i += 2) {
        if (!typedSet.contains(map, i)) {
            keyType.appendTo(map, i, blockBuilder);
            valueType.appendTo(map, i + 1, blockBuilder);
        }
    }
    mapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
}