Search in sources :

Example 11 with ForLoop

use of io.airlift.bytecode.control.ForLoop in project hetu-core by openlookeng.

the class MapTransformKeyFunction method generateTransformKey.

private static MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, Type valueType, Type resultMapType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> keyJavaType = Primitives.wrap(keyType.getJavaType());
    Class<?> transformedKeyJavaType = Primitives.wrap(transformedKeyType.getJavaType());
    Class<?> valueJavaType = Primitives.wrap(valueType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("MapTransformKey"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    Parameter state = arg("state", Object.class);
    Parameter session = arg("session", ConnectorSession.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", BinaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "transform", type(Block.class), ImmutableList.of(state, session, 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 blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable typedSet = scope.declareVariable(TypedSet.class, "typeSet");
    Variable keyElement = scope.declareVariable(keyJavaType, "keyElement");
    Variable transformedKeyElement = scope.declareVariable(transformedKeyJavaType, "transformedKeyElement");
    Variable valueElement = scope.declareVariable(valueJavaType, "valueElement");
    // 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(blockBuilder.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class)));
    // create typed set
    body.append(typedSet.set(newInstance(TypedSet.class, constantType(binder, transformedKeyType), divide(positionCount, constantInt(2)), constantString(MAP_TRANSFORM_KEY_FUNCTION.getSignature().getNameSuffix()))));
    // throw null key exception block
    BytecodeNode throwNullKeyException = new BytecodeBlock().append(newInstance(PrestoException.class, getStatic(INVALID_FUNCTION_ARGUMENT.getDeclaringClass(), "INVALID_FUNCTION_ARGUMENT").cast(ErrorCodeSupplier.class), constantString("map key cannot be null"))).throwObject();
    SqlTypeBytecodeExpression keySqlType = constantType(binder, keyType);
    BytecodeNode loadKeyElement;
    if (!keyType.equals(UNKNOWN)) {
        loadKeyElement = new BytecodeBlock().append(keyElement.set(keySqlType.getValue(block, position).cast(keyJavaType)));
    } else {
        // make sure invokeExact will not take uninitialized keys during compile time
        // but if we reach this point during runtime, it is an exception
        // also close the block builder before throwing as we may be in a TRY() call
        // so that subsequent calls do not find it in an inconsistent state
        loadKeyElement = new BytecodeBlock().append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()).append(keyElement.set(constantNull(keyJavaType))).append(throwNullKeyException);
    }
    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 {
        // make sure invokeExact will not take uninitialized keys during compile time
        loadValueElement = new BytecodeBlock().append(valueElement.set(constantNull(valueJavaType)));
    }
    SqlTypeBytecodeExpression transformedKeySqlType = constantType(binder, transformedKeyType);
    BytecodeNode writeKeyElement;
    BytecodeNode throwDuplicatedKeyException;
    if (!transformedKeyType.equals(UNKNOWN)) {
        writeKeyElement = new BytecodeBlock().append(transformedKeyElement.set(function.invoke("apply", Object.class, keyElement.cast(Object.class), valueElement.cast(Object.class)).cast(transformedKeyJavaType))).append(new IfStatement().condition(equal(transformedKeyElement, constantNull(transformedKeyJavaType))).ifTrue(throwNullKeyException).ifFalse(new BytecodeBlock().append(constantType(binder, transformedKeyType).writeValue(blockBuilder, transformedKeyElement.cast(transformedKeyType.getJavaType()))).append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), blockBuilder))));
        // make sure getObjectValue takes a known key type
        throwDuplicatedKeyException = new BytecodeBlock().append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()).append(newInstance(PrestoException.class, getStatic(INVALID_FUNCTION_ARGUMENT.getDeclaringClass(), "INVALID_FUNCTION_ARGUMENT").cast(ErrorCodeSupplier.class), invokeStatic(String.class, "format", String.class, constantString("Duplicate keys (%s) are not allowed"), newArray(type(Object[].class), ImmutableList.of(transformedKeySqlType.invoke("getObjectValue", Object.class, session, blockBuilder.cast(Block.class), position)))))).throwObject();
    } else {
        // key cannot be unknown
        // if we reach this point during runtime, it is an exception
        writeKeyElement = throwNullKeyException;
        throwDuplicatedKeyException = throwNullKeyException;
    }
    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(writeKeyElement).append(new IfStatement().condition(typedSet.invoke("contains", boolean.class, blockBuilder.cast(Block.class), position)).ifTrue(throwDuplicatedKeyException).ifFalse(typedSet.invoke("add", void.class, blockBuilder.cast(Block.class), position)))));
    body.append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop());
    body.append(pageBuilder.invoke("declarePosition", void.class));
    body.append(constantType(binder, resultMapType).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(), MapTransformKeyFunction.class.getClassLoader());
    return methodHandle(generatedClass, "transform", Object.class, ConnectorSession.class, Block.class, BinaryFunctionInterface.class);
}
Also used : Variable(io.airlift.bytecode.Variable) VariableInstruction.incrementVariable(io.airlift.bytecode.instruction.VariableInstruction.incrementVariable) Signature.typeVariable(io.prestosql.spi.function.Signature.typeVariable) ErrorCodeSupplier(io.prestosql.spi.ErrorCodeSupplier) ForLoop(io.airlift.bytecode.control.ForLoop) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) PrestoException(io.prestosql.spi.PrestoException) ClassDefinition(io.airlift.bytecode.ClassDefinition) IfStatement(io.airlift.bytecode.control.IfStatement) Scope(io.airlift.bytecode.Scope) MethodDefinition(io.airlift.bytecode.MethodDefinition) CallSiteBinder(io.prestosql.sql.gen.CallSiteBinder) TypeSignatureParameter(io.prestosql.spi.type.TypeSignatureParameter) Parameter(io.airlift.bytecode.Parameter) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Block(io.prestosql.spi.block.Block) BytecodeNode(io.airlift.bytecode.BytecodeNode) SqlTypeBytecodeExpression(io.prestosql.sql.gen.SqlTypeBytecodeExpression) BlockBuilder(io.prestosql.spi.block.BlockBuilder)

Example 12 with ForLoop

use of io.airlift.bytecode.control.ForLoop in project hetu-core by openlookeng.

the class ArrayTransformFunction method generateTransform.

private static Class<?> generateTransform(Type inputType, Type outputType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> inputJavaType = Primitives.wrap(inputType.getJavaType());
    Class<?> outputJavaType = Primitives.wrap(outputType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("ArrayTransform"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    // define createPageBuilder
    MethodDefinition createPageBuilderMethod = definition.declareMethod(a(PUBLIC, STATIC), "createPageBuilder", type(PageBuilder.class));
    createPageBuilderMethod.getBody().append(newInstance(PageBuilder.class, constantType(binder, new ArrayType(outputType)).invoke("getTypeParameters", List.class)).ret());
    // define transform method
    Parameter pageBuilder = arg("pageBuilder", PageBuilder.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", UnaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "transform", type(Block.class), ImmutableList.of(pageBuilder, 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 blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable inputElement = scope.declareVariable(inputJavaType, "inputElement");
    Variable outputElement = scope.declareVariable(outputJavaType, "outputElement");
    // invoke block.getPositionCount()
    body.append(positionCount.set(block.invoke("getPositionCount", int.class)));
    // reset page builder if it is full
    body.append(new IfStatement().condition(pageBuilder.invoke("isFull", boolean.class)).ifTrue(pageBuilder.invoke("reset", void.class)));
    // get block builder
    body.append(blockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0))));
    BytecodeNode loadInputElement;
    if (!inputType.equals(UNKNOWN)) {
        loadInputElement = new IfStatement().condition(block.invoke("isNull", boolean.class, position)).ifTrue(inputElement.set(constantNull(inputJavaType))).ifFalse(inputElement.set(constantType(binder, inputType).getValue(block, position).cast(inputJavaType)));
    } else {
        loadInputElement = new BytecodeBlock().append(inputElement.set(constantNull(inputJavaType)));
    }
    BytecodeNode writeOutputElement;
    if (!outputType.equals(UNKNOWN)) {
        writeOutputElement = new IfStatement().condition(equal(outputElement, constantNull(outputJavaType))).ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()).ifFalse(constantType(binder, outputType).writeValue(blockBuilder, outputElement.cast(outputType.getJavaType())));
    } else {
        writeOutputElement = new BytecodeBlock().append(blockBuilder.invoke("appendNull", BlockBuilder.class).pop());
    }
    body.append(new ForLoop().initialize(position.set(constantInt(0))).condition(lessThan(position, positionCount)).update(incrementVariable(position, (byte) 1)).body(new BytecodeBlock().append(loadInputElement).append(outputElement.set(function.invoke("apply", Object.class, inputElement.cast(Object.class)).cast(outputJavaType))).append(writeOutputElement)));
    body.append(pageBuilder.invoke("declarePositions", void.class, positionCount));
    body.append(blockBuilder.invoke("getRegion", Block.class, subtract(blockBuilder.invoke("getPositionCount", int.class), positionCount), positionCount).ret());
    return defineClass(definition, Object.class, binder.getBindings(), ArrayTransformFunction.class.getClassLoader());
}
Also used : Variable(io.airlift.bytecode.Variable) VariableInstruction.incrementVariable(io.airlift.bytecode.instruction.VariableInstruction.incrementVariable) Signature.typeVariable(io.prestosql.spi.function.Signature.typeVariable) ForLoop(io.airlift.bytecode.control.ForLoop) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) PageBuilder(io.prestosql.spi.PageBuilder) ClassDefinition(io.airlift.bytecode.ClassDefinition) ArrayType(io.prestosql.spi.type.ArrayType) IfStatement(io.airlift.bytecode.control.IfStatement) Scope(io.airlift.bytecode.Scope) MethodDefinition(io.airlift.bytecode.MethodDefinition) CallSiteBinder(io.prestosql.sql.gen.CallSiteBinder) Parameter(io.airlift.bytecode.Parameter) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Block(io.prestosql.spi.block.Block) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) BytecodeNode(io.airlift.bytecode.BytecodeNode) BlockBuilder(io.prestosql.spi.block.BlockBuilder)

Example 13 with ForLoop

use of io.airlift.bytecode.control.ForLoop in project hetu-core by openlookeng.

the class JoinCompiler method generateConstructor.

private static void generateConstructor(ClassDefinition classDefinition, List<Integer> joinChannels, FieldDefinition sizeField, FieldDefinition instanceSizeField, List<FieldDefinition> channelFields, List<FieldDefinition> joinChannelFields, FieldDefinition hashChannelField) {
    Parameter channels = arg("channels", type(List.class, type(List.class, Block.class)));
    Parameter hashChannel = arg("hashChannel", type(OptionalInt.class));
    MethodDefinition constructorDefinition = classDefinition.declareConstructor(a(PUBLIC), channels, hashChannel);
    Variable thisVariable = constructorDefinition.getThis();
    Variable blockIndex = constructorDefinition.getScope().declareVariable(int.class, "blockIndex");
    BytecodeBlock constructor = constructorDefinition.getBody().comment("super();").append(thisVariable).invokeConstructor(Object.class);
    constructor.comment("this.size = INSTANCE_SIZE").append(thisVariable.setField(sizeField, getStatic(instanceSizeField)));
    constructor.comment("Set channel fields");
    for (int index = 0; index < channelFields.size(); index++) {
        BytecodeExpression channel = channels.invoke("get", Object.class, constantInt(index)).cast(type(List.class, Block.class));
        constructor.append(thisVariable.setField(channelFields.get(index), channel));
        BytecodeBlock loopBody = new BytecodeBlock();
        constructor.comment("for(blockIndex = 0; blockIndex < channel.size(); blockIndex++) { size += channel.get(i).getRetainedSizeInBytes() }").append(new ForLoop().initialize(blockIndex.set(constantInt(0))).condition(new BytecodeBlock().append(blockIndex).append(channel.invoke("size", int.class)).invokeStatic(CompilerOperations.class, "lessThan", boolean.class, int.class, int.class)).update(new BytecodeBlock().incrementVariable(blockIndex, (byte) 1)).body(loopBody));
        loopBody.append(thisVariable).append(thisVariable).getField(sizeField).append(channel.invoke("get", Object.class, blockIndex).cast(type(Block.class)).invoke("getRetainedSizeInBytes", long.class)).longAdd().putField(sizeField);
    }
    constructor.comment("Set join channel fields");
    for (int index = 0; index < joinChannelFields.size(); index++) {
        BytecodeExpression joinChannel = channels.invoke("get", Object.class, constantInt(joinChannels.get(index))).cast(type(List.class, Block.class));
        constructor.append(thisVariable.setField(joinChannelFields.get(index), joinChannel));
    }
    constructor.comment("Set hashChannel");
    constructor.append(new IfStatement().condition(hashChannel.invoke("isPresent", boolean.class)).ifTrue(thisVariable.setField(hashChannelField, channels.invoke("get", Object.class, hashChannel.invoke("getAsInt", int.class)))).ifFalse(thisVariable.setField(hashChannelField, constantNull(hashChannelField.getType()))));
    constructor.ret();
}
Also used : IfStatement(io.airlift.bytecode.control.IfStatement) Variable(io.airlift.bytecode.Variable) ForLoop(io.airlift.bytecode.control.ForLoop) MethodDefinition(io.airlift.bytecode.MethodDefinition) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Parameter(io.airlift.bytecode.Parameter) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Block(io.prestosql.spi.block.Block) LongArrayList(it.unimi.dsi.fastutil.longs.LongArrayList) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) List(java.util.List) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) OptionalInt(java.util.OptionalInt) BytecodeExpression(io.airlift.bytecode.expression.BytecodeExpression)

Example 14 with ForLoop

use of io.airlift.bytecode.control.ForLoop in project trino by trinodb.

the class ArrayTransformFunction method generateTransform.

private static Class<?> generateTransform(Type inputType, Type outputType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> inputJavaType = Primitives.wrap(inputType.getJavaType());
    Class<?> outputJavaType = Primitives.wrap(outputType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("ArrayTransform"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    // define createPageBuilder
    MethodDefinition createPageBuilderMethod = definition.declareMethod(a(PUBLIC, STATIC), "createPageBuilder", type(PageBuilder.class));
    createPageBuilderMethod.getBody().append(newInstance(PageBuilder.class, constantType(binder, new ArrayType(outputType)).invoke("getTypeParameters", List.class)).ret());
    // define transform method
    Parameter pageBuilder = arg("pageBuilder", PageBuilder.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", UnaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "transform", type(Block.class), ImmutableList.of(pageBuilder, 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 blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable inputElement = scope.declareVariable(inputJavaType, "inputElement");
    Variable outputElement = scope.declareVariable(outputJavaType, "outputElement");
    // invoke block.getPositionCount()
    body.append(positionCount.set(block.invoke("getPositionCount", int.class)));
    // reset page builder if it is full
    body.append(new IfStatement().condition(pageBuilder.invoke("isFull", boolean.class)).ifTrue(pageBuilder.invoke("reset", void.class)));
    // get block builder
    body.append(blockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0))));
    BytecodeNode loadInputElement;
    if (!inputType.equals(UNKNOWN)) {
        loadInputElement = new IfStatement().condition(block.invoke("isNull", boolean.class, position)).ifTrue(inputElement.set(constantNull(inputJavaType))).ifFalse(inputElement.set(constantType(binder, inputType).getValue(block, position).cast(inputJavaType)));
    } else {
        loadInputElement = new BytecodeBlock().append(inputElement.set(constantNull(inputJavaType)));
    }
    BytecodeNode writeOutputElement;
    if (!outputType.equals(UNKNOWN)) {
        writeOutputElement = new IfStatement().condition(equal(outputElement, constantNull(outputJavaType))).ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()).ifFalse(constantType(binder, outputType).writeValue(blockBuilder, outputElement.cast(outputType.getJavaType())));
    } else {
        writeOutputElement = new BytecodeBlock().append(blockBuilder.invoke("appendNull", BlockBuilder.class).pop());
    }
    body.append(new ForLoop().initialize(position.set(constantInt(0))).condition(lessThan(position, positionCount)).update(incrementVariable(position, (byte) 1)).body(new BytecodeBlock().append(loadInputElement).append(outputElement.set(function.invoke("apply", Object.class, inputElement.cast(Object.class)).cast(outputJavaType))).append(writeOutputElement)));
    body.append(pageBuilder.invoke("declarePositions", void.class, positionCount));
    body.append(blockBuilder.invoke("getRegion", Block.class, subtract(blockBuilder.invoke("getPositionCount", int.class), positionCount), positionCount).ret());
    return defineClass(definition, Object.class, binder.getBindings(), ArrayTransformFunction.class.getClassLoader());
}
Also used : Signature.typeVariable(io.trino.metadata.Signature.typeVariable) Variable(io.airlift.bytecode.Variable) VariableInstruction.incrementVariable(io.airlift.bytecode.instruction.VariableInstruction.incrementVariable) ForLoop(io.airlift.bytecode.control.ForLoop) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) PageBuilder(io.trino.spi.PageBuilder) ClassDefinition(io.airlift.bytecode.ClassDefinition) ArrayType(io.trino.spi.type.ArrayType) IfStatement(io.airlift.bytecode.control.IfStatement) Scope(io.airlift.bytecode.Scope) MethodDefinition(io.airlift.bytecode.MethodDefinition) CallSiteBinder(io.trino.sql.gen.CallSiteBinder) Parameter(io.airlift.bytecode.Parameter) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Block(io.trino.spi.block.Block) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) BytecodeNode(io.airlift.bytecode.BytecodeNode) BlockBuilder(io.trino.spi.block.BlockBuilder)

Example 15 with ForLoop

use of io.airlift.bytecode.control.ForLoop in project trino by trinodb.

the class MapTransformKeysFunction method generateTransformKey.

private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, Type valueType, Type resultMapType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> keyJavaType = Primitives.wrap(keyType.getJavaType());
    Class<?> transformedKeyJavaType = Primitives.wrap(transformedKeyType.getJavaType());
    Class<?> valueJavaType = Primitives.wrap(valueType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("MapTransformKey"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    Parameter state = arg("state", Object.class);
    Parameter session = arg("session", ConnectorSession.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", BinaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "transform", type(Block.class), ImmutableList.of(state, session, 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 blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable typedSet = scope.declareVariable(TypedSet.class, "typeSet");
    Variable keyElement = scope.declareVariable(keyJavaType, "keyElement");
    Variable transformedKeyElement = scope.declareVariable(transformedKeyJavaType, "transformedKeyElement");
    Variable valueElement = scope.declareVariable(valueJavaType, "valueElement");
    // 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(blockBuilder.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class)));
    // create typed set
    body.append(typedSet.set(invokeStatic(TypedSet.class, "createEqualityTypedSet", TypedSet.class, constantType(binder, transformedKeyType), loadConstant(binder, blockTypeOperators.getEqualOperator(transformedKeyType), BlockPositionEqual.class), loadConstant(binder, blockTypeOperators.getHashCodeOperator(transformedKeyType), BlockPositionHashCode.class), divide(positionCount, constantInt(2)), constantString(NAME))));
    // throw null key exception block
    BytecodeNode throwNullKeyException = new BytecodeBlock().append(newInstance(TrinoException.class, getStatic(INVALID_FUNCTION_ARGUMENT.getDeclaringClass(), "INVALID_FUNCTION_ARGUMENT").cast(ErrorCodeSupplier.class), constantString("map key cannot be null"))).throwObject();
    SqlTypeBytecodeExpression keySqlType = constantType(binder, keyType);
    BytecodeNode loadKeyElement;
    if (!keyType.equals(UNKNOWN)) {
        loadKeyElement = new BytecodeBlock().append(keyElement.set(keySqlType.getValue(block, position).cast(keyJavaType)));
    } else {
        // make sure invokeExact will not take uninitialized keys during compile time
        // but if we reach this point during runtime, it is an exception
        // also close the block builder before throwing as we may be in a TRY() call
        // so that subsequent calls do not find it in an inconsistent state
        loadKeyElement = new BytecodeBlock().append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()).append(keyElement.set(constantNull(keyJavaType))).append(throwNullKeyException);
    }
    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 {
        // make sure invokeExact will not take uninitialized keys during compile time
        loadValueElement = new BytecodeBlock().append(valueElement.set(constantNull(valueJavaType)));
    }
    SqlTypeBytecodeExpression transformedKeySqlType = constantType(binder, transformedKeyType);
    BytecodeNode writeKeyElement;
    BytecodeNode throwDuplicatedKeyException;
    if (!transformedKeyType.equals(UNKNOWN)) {
        writeKeyElement = new BytecodeBlock().append(transformedKeyElement.set(function.invoke("apply", Object.class, keyElement.cast(Object.class), valueElement.cast(Object.class)).cast(transformedKeyJavaType))).append(new IfStatement().condition(equal(transformedKeyElement, constantNull(transformedKeyJavaType))).ifTrue(throwNullKeyException).ifFalse(new BytecodeBlock().append(constantType(binder, transformedKeyType).writeValue(blockBuilder, transformedKeyElement.cast(transformedKeyType.getJavaType()))).append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), blockBuilder))));
        // make sure getObjectValue takes a known key type
        throwDuplicatedKeyException = new BytecodeBlock().append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()).append(newInstance(TrinoException.class, getStatic(INVALID_FUNCTION_ARGUMENT.getDeclaringClass(), "INVALID_FUNCTION_ARGUMENT").cast(ErrorCodeSupplier.class), invokeStatic(String.class, "format", String.class, constantString("Duplicate keys (%s) are not allowed"), newArray(type(Object[].class), ImmutableList.of(transformedKeySqlType.invoke("getObjectValue", Object.class, session, blockBuilder.cast(Block.class), position)))))).throwObject();
    } else {
        // key cannot be unknown
        // if we reach this point during runtime, it is an exception
        writeKeyElement = throwNullKeyException;
        throwDuplicatedKeyException = throwNullKeyException;
    }
    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(writeKeyElement).append(new IfStatement().condition(typedSet.invoke("add", boolean.class, blockBuilder.cast(Block.class), position)).ifFalse(throwDuplicatedKeyException))));
    body.append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop());
    body.append(pageBuilder.invoke("declarePosition", void.class));
    body.append(constantType(binder, resultMapType).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(), MapTransformKeysFunction.class.getClassLoader());
    return methodHandle(generatedClass, "transform", Object.class, ConnectorSession.class, Block.class, BinaryFunctionInterface.class);
}
Also used : Signature.typeVariable(io.trino.metadata.Signature.typeVariable) Variable(io.airlift.bytecode.Variable) VariableInstruction.incrementVariable(io.airlift.bytecode.instruction.VariableInstruction.incrementVariable) ErrorCodeSupplier(io.trino.spi.ErrorCodeSupplier) ForLoop(io.airlift.bytecode.control.ForLoop) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) ClassDefinition(io.airlift.bytecode.ClassDefinition) BlockPositionEqual(io.trino.type.BlockTypeOperators.BlockPositionEqual) IfStatement(io.airlift.bytecode.control.IfStatement) Scope(io.airlift.bytecode.Scope) MethodDefinition(io.airlift.bytecode.MethodDefinition) CallSiteBinder(io.trino.sql.gen.CallSiteBinder) Parameter(io.airlift.bytecode.Parameter) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Block(io.trino.spi.block.Block) TrinoException(io.trino.spi.TrinoException) BytecodeNode(io.airlift.bytecode.BytecodeNode) BlockPositionHashCode(io.trino.type.BlockTypeOperators.BlockPositionHashCode) SqlTypeBytecodeExpression(io.trino.sql.gen.SqlTypeBytecodeExpression) BlockBuilder(io.trino.spi.block.BlockBuilder)

Aggregations

Variable (io.airlift.bytecode.Variable)20 ForLoop (io.airlift.bytecode.control.ForLoop)20 IfStatement (io.airlift.bytecode.control.IfStatement)20 BytecodeBlock (io.airlift.bytecode.BytecodeBlock)18 MethodDefinition (io.airlift.bytecode.MethodDefinition)16 Parameter (io.airlift.bytecode.Parameter)14 Scope (io.airlift.bytecode.Scope)14 BytecodeNode (io.airlift.bytecode.BytecodeNode)10 ClassDefinition (io.airlift.bytecode.ClassDefinition)8 VariableInstruction.incrementVariable (io.airlift.bytecode.instruction.VariableInstruction.incrementVariable)8 Block (io.prestosql.spi.block.Block)7 Block (io.trino.spi.block.Block)7 ImmutableList (com.google.common.collect.ImmutableList)5 List (java.util.List)5 BytecodeExpression (io.airlift.bytecode.expression.BytecodeExpression)4 Signature.typeVariable (io.prestosql.spi.function.Signature.typeVariable)4 CallSiteBinder (io.prestosql.sql.gen.CallSiteBinder)4 Signature.typeVariable (io.trino.metadata.Signature.typeVariable)4 CallSiteBinder (io.trino.sql.gen.CallSiteBinder)4 ArrayList (java.util.ArrayList)4