Examples with CodegenScalarFunction com.facebook.presto.spi.function.CodegenScalarFunction used on opensource projects

Example 1 with CodegenScalarFunction

use of com.facebook.presto.spi.function.CodegenScalarFunction in project presto by prestodb.

the class CodegenScalarFromAnnotationsParser method createSqlScalarFunction.

private static SqlScalarFunction createSqlScalarFunction(Method method) {
    CodegenScalarFunction codegenScalarFunction = method.getAnnotation(CodegenScalarFunction.class);
    Signature signature = new Signature(QualifiedObjectName.valueOf(DEFAULT_NAMESPACE, codegenScalarFunction.value()), FunctionKind.SCALAR, Arrays.stream(method.getAnnotationsByType(TypeParameter.class)).map(t -> withVariadicBound(t.value(), t.boundedBy().isEmpty() ? null : t.boundedBy())).collect(toImmutableList()), ImmutableList.of(), parseTypeSignature(method.getAnnotation(SqlType.class).value()), Arrays.stream(method.getParameters()).map(p -> parseTypeSignature(p.getAnnotation(SqlType.class).value())).collect(toImmutableList()), false);
    return new SqlScalarFunction(signature) {

        @Override
        public BuiltInScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, FunctionAndTypeManager functionAndTypeManager) {
            Signature boundSignature = applyBoundVariables(signature, boundVariables, arity);
            MethodHandle handle;
            try {
                handle = (MethodHandle) method.invoke(null, boundSignature.getArgumentTypes().stream().map(t -> functionAndTypeManager.getType(t)).toArray());
            } catch (Exception e) {
                throw new PrestoException(FUNCTION_IMPLEMENTATION_ERROR, format("Method %s does not return valid MethodHandle", method), e);
            }
            return new BuiltInScalarFunctionImplementation(method.getAnnotation(SqlNullable.class) != null, getArgumentProperties(method), handle, Optional.empty());
        }

        @Override
        public SqlFunctionVisibility getVisibility() {
            return codegenScalarFunction.visibility();
        }

        @Override
        public boolean isDeterministic() {
            return codegenScalarFunction.deterministic();
        }

        @Override
        public String getDescription() {
            Description description = method.getAnnotation(Description.class);
            return description == null ? "" : description.value();
        }

        @Override
        public boolean isCalledOnNullInput() {
            return codegenScalarFunction.calledOnNullInput();
        }
    };
}

Example 2 with CodegenScalarFunction

use of com.facebook.presto.spi.function.CodegenScalarFunction in project presto by prestodb.

the class MapFilterFunction method mapFilter.

@CodegenScalarFunction(value = "map_filter", deterministic = false)
@Description("return map containing entries that match the given predicate")
@TypeParameter("K")
@TypeParameter("V")
@SqlType("map(K,V)")
public static MethodHandle mapFilter(@SqlType("map(K,V)") Type inputType, @SqlType("function(K,V,boolean)") Type filter) {
    MapType mapType = (MapType) inputType;
    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 block = arg("block", Block.class);
    Parameter function = arg("function", BinaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "filter", type(Block.class), ImmutableList.of(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 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(mapBlockBuilder.set(constantType(binder, mapType).invoke("createBlockBuilder", BlockBuilder.class, constantNull(BlockBuilderStatus.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(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", Block.class, BinaryFunctionInterface.class);
}