Search in sources :

Example 1 with SingleMapBlock

use of io.trino.spi.block.SingleMapBlock in project trino by trinodb.

the class TestSetDigest method testHashCounts.

@Test
public void testHashCounts() {
    SetDigest digest1 = new SetDigest();
    digest1.add(0);
    digest1.add(0);
    digest1.add(1);
    SetDigest digest2 = new SetDigest();
    digest2.add(0);
    digest2.add(0);
    digest2.add(2);
    digest2.add(2);
    MapType mapType = new MapType(BIGINT, SMALLINT, new TypeOperators());
    Block block = hashCounts(mapType, digest1.serialize());
    assertTrue(block instanceof SingleMapBlock);
    Set<Short> blockValues = new HashSet<>();
    for (int i = 1; i < block.getPositionCount(); i += 2) {
        blockValues.add(block.getShort(i, 0));
    }
    Set<Short> expected = ImmutableSet.of((short) 1, (short) 2);
    assertEquals(blockValues, expected);
    digest1.mergeWith(digest2);
    block = hashCounts(mapType, digest1.serialize());
    assertTrue(block instanceof SingleMapBlock);
    expected = ImmutableSet.of((short) 1, (short) 2, (short) 4);
    blockValues = new HashSet<>();
    for (int i = 1; i < block.getPositionCount(); i += 2) {
        blockValues.add(block.getShort(i, 0));
    }
    assertEquals(blockValues, expected);
}
Also used : SingleMapBlock(io.trino.spi.block.SingleMapBlock) Block(io.trino.spi.block.Block) SingleMapBlock(io.trino.spi.block.SingleMapBlock) MapType(io.trino.spi.type.MapType) TypeOperators(io.trino.spi.type.TypeOperators) HashSet(java.util.HashSet) Test(org.testng.annotations.Test)

Example 2 with SingleMapBlock

use of io.trino.spi.block.SingleMapBlock in project trino by trinodb.

the class MapElementAtFunction method elementAt.

@UsedByGeneratedCode
public static Object elementAt(Type valueType, Block map, boolean key) {
    SingleMapBlock mapBlock = (SingleMapBlock) map;
    int valuePosition = mapBlock.seekKeyExact(key);
    if (valuePosition == -1) {
        return null;
    }
    return readNativeValue(valueType, mapBlock, valuePosition);
}
Also used : SingleMapBlock(io.trino.spi.block.SingleMapBlock) UsedByGeneratedCode(io.trino.annotation.UsedByGeneratedCode)

Example 3 with SingleMapBlock

use of io.trino.spi.block.SingleMapBlock in project trino by trinodb.

the class MapSubscriptOperator method subscript.

@UsedByGeneratedCode
public static Object subscript(MissingKeyExceptionFactory missingKeyExceptionFactory, Type keyType, Type valueType, ConnectorSession session, Block map, boolean key) {
    SingleMapBlock mapBlock = (SingleMapBlock) map;
    int valuePosition = mapBlock.seekKeyExact(key);
    if (valuePosition == -1) {
        throw missingKeyExceptionFactory.create(session, key);
    }
    return readNativeValue(valueType, mapBlock, valuePosition);
}
Also used : SingleMapBlock(io.trino.spi.block.SingleMapBlock) UsedByGeneratedCode(io.trino.annotation.UsedByGeneratedCode)

Example 4 with SingleMapBlock

use of io.trino.spi.block.SingleMapBlock 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);
}
Also used : Type(io.trino.spi.type.Type) TypeSignature.functionType(io.trino.spi.type.TypeSignature.functionType) MapType(io.trino.spi.type.MapType) TypeSignature.mapType(io.trino.spi.type.TypeSignature.mapType) PageBuilder(io.trino.spi.PageBuilder) SingleMapBlock(io.trino.spi.block.SingleMapBlock) BlockBuilder(io.trino.spi.block.BlockBuilder)

Example 5 with SingleMapBlock

use of io.trino.spi.block.SingleMapBlock in project trino by trinodb.

the class MapElementAtFunction method elementAt.

@UsedByGeneratedCode
public static Object elementAt(Type valueType, Block map, double key) {
    SingleMapBlock mapBlock = (SingleMapBlock) map;
    int valuePosition = mapBlock.seekKeyExact(key);
    if (valuePosition == -1) {
        return null;
    }
    return readNativeValue(valueType, mapBlock, valuePosition);
}
Also used : SingleMapBlock(io.trino.spi.block.SingleMapBlock) UsedByGeneratedCode(io.trino.annotation.UsedByGeneratedCode)

Aggregations

SingleMapBlock (io.trino.spi.block.SingleMapBlock)16 UsedByGeneratedCode (io.trino.annotation.UsedByGeneratedCode)8 Block (io.trino.spi.block.Block)4 MapBlock (io.trino.spi.block.MapBlock)3 MapType (io.trino.spi.type.MapType)3 HashMap (java.util.HashMap)3 BlockBuilder (io.trino.spi.block.BlockBuilder)2 ImplementAvgFloatingPoint (io.trino.plugin.jdbc.aggregation.ImplementAvgFloatingPoint)1 PageBuilder (io.trino.spi.PageBuilder)1 MapBlockBuilder (io.trino.spi.block.MapBlockBuilder)1 Type (io.trino.spi.type.Type)1 TypeOperators (io.trino.spi.type.TypeOperators)1 TypeSignature.functionType (io.trino.spi.type.TypeSignature.functionType)1 TypeSignature.mapType (io.trino.spi.type.TypeSignature.mapType)1 HashSet (java.util.HashSet)1 Map (java.util.Map)1 Test (org.testng.annotations.Test)1