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

Example 11 with MapType

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

the class AvroDecoderTestUtil method checkArrayValues.

public static void checkArrayValues(Block block, Type type, Object value) {
    assertNotNull(type, "Type is null");
    assertTrue(type instanceof ArrayType, "Unexpected type");
    assertNotNull(block, "Block is null");
    assertNotNull(value, "Value is null");
    List<?> list = (List<?>) value;
    assertEquals(block.getPositionCount(), list.size());
    Type elementType = ((ArrayType) type).getElementType();
    if (elementType instanceof ArrayType) {
        for (int index = 0; index < block.getPositionCount(); index++) {
            if (block.isNull(index)) {
                assertNull(list.get(index));
                continue;
            }
            Block arrayBlock = block.getObject(index, Block.class);
            checkArrayValues(arrayBlock, elementType, list.get(index));
        }
    } else if (elementType instanceof MapType) {
        for (int index = 0; index < block.getPositionCount(); index++) {
            if (block.isNull(index)) {
                assertNull(list.get(index));
                continue;
            }
            Block mapBlock = block.getObject(index, Block.class);
            checkMapValues(mapBlock, elementType, list.get(index));
        }
    } else if (elementType instanceof RowType) {
        for (int index = 0; index < block.getPositionCount(); index++) {
            if (block.isNull(index)) {
                assertNull(list.get(index));
                continue;
            }
            Block rowBlock = block.getObject(index, Block.class);
            checkRowValues(rowBlock, elementType, list.get(index));
        }
    } else {
        for (int index = 0; index < block.getPositionCount(); index++) {
            checkPrimitiveValue(getObjectValue(elementType, block, index), list.get(index));
        }
    }
}

Example 12 with MapType

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

the class TestDeltaLakeSchemaSupport method testSerializeSchemaAsJson.

@Test
public void testSerializeSchemaAsJson() throws Exception {
    DeltaLakeColumnHandle arrayColumn = new DeltaLakeColumnHandle("arr", new ArrayType(new ArrayType(INTEGER)), REGULAR);
    DeltaLakeColumnHandle structColumn = new DeltaLakeColumnHandle("str", RowType.from(ImmutableList.of(new RowType.Field(Optional.of("s1"), VarcharType.createUnboundedVarcharType()), new RowType.Field(Optional.of("s2"), RowType.from(ImmutableList.of(new RowType.Field(Optional.of("i1"), INTEGER), new RowType.Field(Optional.of("d2"), DecimalType.createDecimalType(38, 0))))))), REGULAR);
    TypeOperators typeOperators = new TypeOperators();
    DeltaLakeColumnHandle mapColumn = new DeltaLakeColumnHandle("m", new MapType(INTEGER, new MapType(INTEGER, INTEGER, typeOperators), typeOperators), REGULAR);
    URL expected = getResource("io/trino/plugin/deltalake/transactionlog/schema/nested_schema.json");
    ObjectMapper objectMapper = new ObjectMapper();
    String jsonEncoding = serializeSchemaAsJson(ImmutableList.of(arrayColumn, structColumn, mapColumn));
    assertEquals(objectMapper.readTree(jsonEncoding), objectMapper.readTree(expected));
}

Example 13 with MapType

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

the class MapFilterFunction method generateFilter.

private static MethodHandle generateFilter(MapType mapType) {
    CallSiteBinder binder = new CallSiteBinder();
    Type keyType = mapType.getKeyType();
    Type valueType = mapType.getValueType();
    Class<?> keyJavaType = Primitives.wrap(keyType.getJavaType());
    Class<?> valueJavaType = Primitives.wrap(valueType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("MapFilter"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    Parameter state = arg("state", Object.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", BinaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "filter", type(Block.class), ImmutableList.of(state, block, function));
    BytecodeBlock body = method.getBody();
    Scope scope = method.getScope();
    Variable positionCount = scope.declareVariable(int.class, "positionCount");
    Variable position = scope.declareVariable(int.class, "position");
    Variable pageBuilder = scope.declareVariable(PageBuilder.class, "pageBuilder");
    Variable mapBlockBuilder = scope.declareVariable(BlockBuilder.class, "mapBlockBuilder");
    Variable singleMapBlockWriter = scope.declareVariable(BlockBuilder.class, "singleMapBlockWriter");
    Variable keyElement = scope.declareVariable(keyJavaType, "keyElement");
    Variable valueElement = scope.declareVariable(valueJavaType, "valueElement");
    Variable keep = scope.declareVariable(Boolean.class, "keep");
    // invoke block.getPositionCount()
    body.append(positionCount.set(block.invoke("getPositionCount", int.class)));
    // prepare the single map block builder
    body.append(pageBuilder.set(state.cast(PageBuilder.class)));
    body.append(new IfStatement().condition(pageBuilder.invoke("isFull", boolean.class)).ifTrue(pageBuilder.invoke("reset", void.class)));
    body.append(mapBlockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0))));
    body.append(singleMapBlockWriter.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class)));
    SqlTypeBytecodeExpression keySqlType = constantType(binder, keyType);
    BytecodeNode loadKeyElement;
    if (!keyType.equals(UNKNOWN)) {
        // key element must be non-null
        loadKeyElement = new BytecodeBlock().append(keyElement.set(keySqlType.getValue(block, position).cast(keyJavaType)));
    } else {
        loadKeyElement = new BytecodeBlock().append(keyElement.set(constantNull(keyJavaType)));
    }
    SqlTypeBytecodeExpression valueSqlType = constantType(binder, valueType);
    BytecodeNode loadValueElement;
    if (!valueType.equals(UNKNOWN)) {
        loadValueElement = new IfStatement().condition(block.invoke("isNull", boolean.class, add(position, constantInt(1)))).ifTrue(valueElement.set(constantNull(valueJavaType))).ifFalse(valueElement.set(valueSqlType.getValue(block, add(position, constantInt(1))).cast(valueJavaType)));
    } else {
        loadValueElement = new BytecodeBlock().append(valueElement.set(constantNull(valueJavaType)));
    }
    body.append(new ForLoop().initialize(position.set(constantInt(0))).condition(lessThan(position, positionCount)).update(incrementVariable(position, (byte) 2)).body(new BytecodeBlock().append(loadKeyElement).append(loadValueElement).append(keep.set(function.invoke("apply", Object.class, keyElement.cast(Object.class), valueElement.cast(Object.class)).cast(Boolean.class))).append(new IfStatement("if (keep != null && keep) ...").condition(and(notEqual(keep, constantNull(Boolean.class)), keep.cast(boolean.class))).ifTrue(new BytecodeBlock().append(keySqlType.invoke("appendTo", void.class, block, position, singleMapBlockWriter)).append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), singleMapBlockWriter))))));
    body.append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop());
    body.append(pageBuilder.invoke("declarePosition", void.class));
    body.append(constantType(binder, mapType).invoke("getObject", Object.class, mapBlockBuilder.cast(Block.class), subtract(mapBlockBuilder.invoke("getPositionCount", int.class), constantInt(1))).ret());
    Class<?> generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapFilterFunction.class.getClassLoader());
    return methodHandle(generatedClass, "filter", Object.class, Block.class, BinaryFunctionInterface.class);
}

Example 14 with MapType

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

the class MapFromEntriesFunction method mapFromEntries.

@TypeParameter("K")
@TypeParameter("V")
@SqlType("map(K,V)")
@SqlNullable
public Block mapFromEntries(@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, @TypeParameter("map(K,V)") MapType mapType, ConnectorSession session, @SqlType("array(row(K,V))") Block mapEntries) {
    Type keyType = mapType.getKeyType();
    Type valueType = mapType.getValueType();
    RowType mapEntryType = RowType.anonymous(ImmutableList.of(keyType, valueType));
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    int entryCount = mapEntries.getPositionCount();
    BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder resultBuilder = mapBlockBuilder.beginBlockEntry();
    TypedSet uniqueKeys = createEqualityTypedSet(keyType, keyEqual, keyHashCode, entryCount, "map_from_entries");
    for (int i = 0; i < entryCount; i++) {
        if (mapEntries.isNull(i)) {
            mapBlockBuilder.closeEntry();
            pageBuilder.declarePosition();
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null");
        }
        Block mapEntryBlock = mapEntryType.getObject(mapEntries, i);
        if (mapEntryBlock.isNull(0)) {
            mapBlockBuilder.closeEntry();
            pageBuilder.declarePosition();
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }
        if (!uniqueKeys.add(mapEntryBlock, 0)) {
            mapBlockBuilder.closeEntry();
            pageBuilder.declarePosition();
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("Duplicate keys (%s) are not allowed", keyType.getObjectValue(session, mapEntryBlock, 0)));
        }
        keyType.appendTo(mapEntryBlock, 0, resultBuilder);
        valueType.appendTo(mapEntryBlock, 1, resultBuilder);
    }
    mapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
}

Example 15 with MapType

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

the class MapToJsonCast method specialize.

@Override
public ScalarFunctionImplementation specialize(BoundSignature boundSignature) {
    MapType mapType = (MapType) boundSignature.getArgumentType(0);
    Type keyType = mapType.getKeyType();
    Type valueType = mapType.getValueType();
    checkCondition(canCastToJson(mapType), INVALID_CAST_ARGUMENT, "Cannot cast %s to JSON", mapType);
    ObjectKeyProvider provider = ObjectKeyProvider.createObjectKeyProvider(keyType);
    JsonGeneratorWriter writer = JsonGeneratorWriter.createJsonGeneratorWriter(valueType, legacyRowToJson);
    MethodHandle methodHandle = METHOD_HANDLE.bindTo(provider).bindTo(writer);
    return new ChoicesScalarFunctionImplementation(boundSignature, FAIL_ON_NULL, ImmutableList.of(NEVER_NULL), methodHandle);
}