Search in sources :

Example 21 with BytecodeNode

use of io.airlift.bytecode.BytecodeNode in project trino by trinodb.

the class BytecodeUtils method generateFullInvocation.

private static BytecodeNode generateFullInvocation(Scope scope, String functionName, FunctionNullability functionNullability, List<Boolean> argumentIsFunctionType, Function<InvocationConvention, FunctionInvoker> functionInvokerProvider, Function<MethodHandle, BytecodeNode> instanceFactory, List<Function<Optional<Class<?>>, BytecodeNode>> argumentCompilers, CallSiteBinder binder) {
    verify(argumentIsFunctionType.size() == argumentCompilers.size());
    List<InvocationArgumentConvention> argumentConventions = new ArrayList<>();
    List<BytecodeNode> arguments = new ArrayList<>();
    for (int i = 0; i < argumentIsFunctionType.size(); i++) {
        if (argumentIsFunctionType.get(i)) {
            argumentConventions.add(FUNCTION);
            arguments.add(null);
        } else {
            BytecodeNode argument = argumentCompilers.get(i).apply(Optional.empty());
            argumentConventions.add(getPreferredArgumentConvention(argument, argumentCompilers.size(), functionNullability.isArgumentNullable(i)));
            arguments.add(argument);
        }
    }
    InvocationConvention invocationConvention = new InvocationConvention(argumentConventions, functionNullability.isReturnNullable() ? NULLABLE_RETURN : FAIL_ON_NULL, true, true);
    FunctionInvoker functionInvoker = functionInvokerProvider.apply(invocationConvention);
    Binding binding = binder.bind(functionInvoker.getMethodHandle());
    LabelNode end = new LabelNode("end");
    BytecodeBlock block = new BytecodeBlock().setDescription("invoke " + functionName);
    Optional<BytecodeNode> instance = functionInvoker.getInstanceFactory().map(instanceFactory);
    // Index of current parameter in the MethodHandle
    int currentParameterIndex = 0;
    // Index of parameter (without @IsNull) in Trino function
    int realParameterIndex = 0;
    // Index of function argument types
    int lambdaArgumentIndex = 0;
    MethodType methodType = binding.getType();
    Class<?> returnType = methodType.returnType();
    Class<?> unboxedReturnType = Primitives.unwrap(returnType);
    List<Class<?>> stackTypes = new ArrayList<>();
    boolean instanceIsBound = false;
    while (currentParameterIndex < methodType.parameterArray().length) {
        Class<?> type = methodType.parameterArray()[currentParameterIndex];
        stackTypes.add(type);
        if (instance.isPresent() && !instanceIsBound) {
            checkState(type.equals(functionInvoker.getInstanceFactory().get().type().returnType()), "Mismatched type for instance parameter");
            block.append(instance.get());
            instanceIsBound = true;
        } else if (type == ConnectorSession.class) {
            block.append(scope.getVariable("session"));
        } else {
            switch(invocationConvention.getArgumentConvention(realParameterIndex)) {
                case NEVER_NULL:
                    block.append(arguments.get(realParameterIndex));
                    checkArgument(!Primitives.isWrapperType(type), "Non-nullable argument must not be primitive wrapper type");
                    block.append(ifWasNullPopAndGoto(scope, end, unboxedReturnType, Lists.reverse(stackTypes)));
                    break;
                case NULL_FLAG:
                    block.append(arguments.get(realParameterIndex));
                    block.append(scope.getVariable("wasNull"));
                    block.append(scope.getVariable("wasNull").set(constantFalse()));
                    stackTypes.add(boolean.class);
                    currentParameterIndex++;
                    break;
                case BOXED_NULLABLE:
                    block.append(arguments.get(realParameterIndex));
                    block.append(boxPrimitiveIfNecessary(scope, type));
                    block.append(scope.getVariable("wasNull").set(constantFalse()));
                    break;
                case BLOCK_POSITION:
                    InputReferenceNode inputReferenceNode = (InputReferenceNode) arguments.get(realParameterIndex);
                    block.append(inputReferenceNode.produceBlockAndPosition());
                    stackTypes.add(int.class);
                    if (!functionNullability.isArgumentNullable(realParameterIndex)) {
                        block.append(scope.getVariable("wasNull").set(inputReferenceNode.blockAndPositionIsNull()));
                        block.append(ifWasNullPopAndGoto(scope, end, unboxedReturnType, Lists.reverse(stackTypes)));
                    }
                    currentParameterIndex++;
                    break;
                case FUNCTION:
                    Class<?> lambdaInterface = functionInvoker.getLambdaInterfaces().get(lambdaArgumentIndex);
                    block.append(argumentCompilers.get(realParameterIndex).apply(Optional.of(lambdaInterface)));
                    lambdaArgumentIndex++;
                    break;
                default:
                    throw new UnsupportedOperationException(format("Unsupported argument conventsion type: %s", invocationConvention.getArgumentConvention(realParameterIndex)));
            }
            realParameterIndex++;
        }
        currentParameterIndex++;
    }
    block.append(invoke(binding, functionName));
    if (functionNullability.isReturnNullable()) {
        block.append(unboxPrimitiveIfNecessary(scope, returnType));
    }
    block.visitLabel(end);
    return block;
}
Also used : LabelNode(io.airlift.bytecode.instruction.LabelNode) MethodType(java.lang.invoke.MethodType) ArrayList(java.util.ArrayList) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) InputReferenceNode(io.trino.sql.gen.InputReferenceCompiler.InputReferenceNode) InvocationArgumentConvention(io.trino.spi.function.InvocationConvention.InvocationArgumentConvention) InvocationConvention(io.trino.spi.function.InvocationConvention) BytecodeNode(io.airlift.bytecode.BytecodeNode) FunctionInvoker(io.trino.metadata.FunctionInvoker) ConnectorSession(io.trino.spi.connector.ConnectorSession)

Example 22 with BytecodeNode

use of io.airlift.bytecode.BytecodeNode 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);
}
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) ClassDefinition(io.airlift.bytecode.ClassDefinition) IfStatement(io.airlift.bytecode.control.IfStatement) SqlTypeBytecodeExpression.constantType(io.trino.sql.gen.SqlTypeBytecodeExpression.constantType) TypeSignature.mapType(io.trino.spi.type.TypeSignature.mapType) Type(io.trino.spi.type.Type) TypeSignature.functionType(io.trino.spi.type.TypeSignature.functionType) MapType(io.trino.spi.type.MapType) 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) BytecodeNode(io.airlift.bytecode.BytecodeNode) SqlTypeBytecodeExpression(io.trino.sql.gen.SqlTypeBytecodeExpression)

Example 23 with BytecodeNode

use of io.airlift.bytecode.BytecodeNode in project trino by trinodb.

the class MapTransformValuesFunction method generateTransform.

private static MethodHandle generateTransform(Type keyType, Type valueType, Type transformedValueType, Type resultMapType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> keyJavaType = Primitives.wrap(keyType.getJavaType());
    Class<?> valueJavaType = Primitives.wrap(valueType.getJavaType());
    Class<?> transformedValueJavaType = Primitives.wrap(transformedValueType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("MapTransformValue"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    // define transform method
    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), "transform", 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 blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable keyElement = scope.declareVariable(keyJavaType, "keyElement");
    Variable valueElement = scope.declareVariable(valueJavaType, "valueElement");
    Variable transformedValueElement = scope.declareVariable(transformedValueJavaType, "transformedValueElement");
    // 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)));
    // 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 {
        loadValueElement = new BytecodeBlock().append(valueElement.set(constantNull(valueJavaType)));
    }
    BytecodeNode writeTransformedValueElement;
    if (!transformedValueType.equals(UNKNOWN)) {
        writeTransformedValueElement = new IfStatement().condition(equal(transformedValueElement, constantNull(transformedValueJavaType))).ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()).ifFalse(constantType(binder, transformedValueType).writeValue(blockBuilder, transformedValueElement.cast(transformedValueType.getJavaType())));
    } else {
        writeTransformedValueElement = new BytecodeBlock().append(blockBuilder.invoke("appendNull", BlockBuilder.class).pop());
    }
    Variable transformationException = scope.declareVariable(Throwable.class, "transformationException");
    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(new TryCatch("Close builder before throwing to avoid subsequent calls finding it in an inconsistent state if we are in a TRY() call.", transformedValueElement.set(function.invoke("apply", Object.class, keyElement.cast(Object.class), valueElement.cast(Object.class)).cast(transformedValueJavaType)), ImmutableList.of(new TryCatch.CatchBlock(new BytecodeBlock().append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()).append(pageBuilder.invoke("declarePosition", void.class)).putVariable(transformationException).append(invokeStatic(Throwables.class, "throwIfUnchecked", void.class, transformationException)).append(newInstance(RuntimeException.class, transformationException)).throwObject(), ImmutableList.of(type(Throwable.class)))))).append(keySqlType.invoke("appendTo", void.class, block, position, blockBuilder)).append(writeTransformedValueElement)));
    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(), MapTransformValuesFunction.class.getClassLoader());
    return methodHandle(generatedClass, "transform", Object.class, Block.class, BinaryFunctionInterface.class);
}
Also used : TryCatch(io.airlift.bytecode.control.TryCatch) 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) Throwables(com.google.common.base.Throwables) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) ClassDefinition(io.airlift.bytecode.ClassDefinition) 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) BytecodeNode(io.airlift.bytecode.BytecodeNode) SqlTypeBytecodeExpression(io.trino.sql.gen.SqlTypeBytecodeExpression) BlockBuilder(io.trino.spi.block.BlockBuilder)

Example 24 with BytecodeNode

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

the class AbstractMinMaxBy method generateInputMethod.

private void generateInputMethod(ClassDefinition definition, CallSiteBinder binder, MethodHandle compareMethod, Type keyType, Type valueType, Class<?> stateClass) {
    Parameter state = arg("state", stateClass);
    Parameter value = arg("value", Block.class);
    Parameter key = arg("key", Block.class);
    Parameter position = arg("position", int.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "input", type(void.class), state, value, key, position);
    SqlTypeBytecodeExpression keySqlType = constantType(binder, keyType);
    BytecodeBlock ifBlock = new BytecodeBlock().append(state.invoke("setFirst", void.class, keySqlType.getValue(key, position))).append(state.invoke("setFirstNull", void.class, constantBoolean(false))).append(state.invoke("setSecondNull", void.class, value.invoke("isNull", boolean.class, position)));
    BytecodeNode setValueNode;
    if (valueType.getJavaType().isPrimitive()) {
        SqlTypeBytecodeExpression valueSqlType = constantType(binder, valueType);
        setValueNode = state.invoke("setSecond", void.class, valueSqlType.getValue(value, position));
    } else {
        // Do not get value directly given it creates object overhead.
        // Such objects would live long enough in Block or SliceBigArray to cause GC pressure.
        setValueNode = new BytecodeBlock().append(state.invoke("setSecondBlock", void.class, value)).append(state.invoke("setSecondPosition", void.class, position));
    }
    ifBlock.append(new IfStatement().condition(value.invoke("isNull", boolean.class, position)).ifFalse(setValueNode));
    method.getBody().append(new IfStatement().condition(or(state.invoke("isFirstNull", boolean.class), and(not(key.invoke("isNull", boolean.class, position)), loadConstant(binder, compareMethod, MethodHandle.class).invoke("invokeExact", boolean.class, keySqlType.getValue(key, position), state.invoke("getFirst", keyType.getJavaType()))))).ifTrue(ifBlock)).ret();
}
Also used : IfStatement(io.airlift.bytecode.control.IfStatement) MethodDefinition(io.airlift.bytecode.MethodDefinition) BytecodeBlock(io.airlift.bytecode.BytecodeBlock) Signature.orderableTypeParameter(io.prestosql.spi.function.Signature.orderableTypeParameter) Parameter(io.airlift.bytecode.Parameter) BytecodeNode(io.airlift.bytecode.BytecodeNode) SqlTypeBytecodeExpression(io.prestosql.sql.gen.SqlTypeBytecodeExpression) MethodHandle(java.lang.invoke.MethodHandle)

Example 25 with BytecodeNode

use of io.airlift.bytecode.BytecodeNode 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)

Aggregations

BytecodeNode (io.airlift.bytecode.BytecodeNode)42 Variable (io.airlift.bytecode.Variable)37 BytecodeBlock (io.airlift.bytecode.BytecodeBlock)29 IfStatement (io.airlift.bytecode.control.IfStatement)26 MethodDefinition (io.airlift.bytecode.MethodDefinition)22 Parameter (io.airlift.bytecode.Parameter)20 Scope (io.airlift.bytecode.Scope)19 LabelNode (io.airlift.bytecode.instruction.LabelNode)15 ForLoop (io.airlift.bytecode.control.ForLoop)10 BytecodeExpression (io.airlift.bytecode.expression.BytecodeExpression)9 ClassDefinition (io.airlift.bytecode.ClassDefinition)8 VariableInstruction.incrementVariable (io.airlift.bytecode.instruction.VariableInstruction.incrementVariable)8 RowExpression (io.prestosql.spi.relation.RowExpression)6 ArrayList (java.util.ArrayList)6 ImmutableList (com.google.common.collect.ImmutableList)5 Block (io.prestosql.spi.block.Block)5 BuiltInScalarFunctionImplementation (io.prestosql.spi.function.BuiltInScalarFunctionImplementation)5 Block (io.trino.spi.block.Block)5 FunctionHandle (io.prestosql.spi.function.FunctionHandle)4 Signature.typeVariable (io.prestosql.spi.function.Signature.typeVariable)4