Examples with FunctionArgumentType com.google.zetasql.FunctionArgumentType used on opensource projects

Example 1 with FunctionArgumentType

use of com.google.zetasql.FunctionArgumentType in project beam by apache.

the class BeamZetaSqlCatalog method addUdfsFromSchema.

private void addUdfsFromSchema() {
    for (String functionName : calciteSchema.getFunctionNames()) {
        Collection<org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function> functions = calciteSchema.getFunctions(functionName);
        if (functions.size() != 1) {
            throw new IllegalArgumentException(String.format("Expected exactly 1 definition for function '%s', but found %d." + " Beam ZetaSQL supports only a single function definition per function name (BEAM-12073).", functionName, functions.size()));
        }
        for (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function function : functions) {
            List<String> path = Arrays.asList(functionName.split("\\."));
            if (function instanceof ScalarFunctionImpl) {
                ScalarFunctionImpl scalarFunction = (ScalarFunctionImpl) function;
                // for unsupported types.
                for (FunctionParameter parameter : scalarFunction.getParameters()) {
                    validateJavaUdfCalciteType(parameter.getType(typeFactory), functionName);
                }
                validateJavaUdfCalciteType(scalarFunction.getReturnType(typeFactory), functionName);
                Method method = scalarFunction.method;
                javaScalarUdfs.put(path, UserFunctionDefinitions.JavaScalarFunction.create(method, ""));
                FunctionArgumentType resultType = new FunctionArgumentType(ZetaSqlCalciteTranslationUtils.toZetaSqlType(scalarFunction.getReturnType(typeFactory)));
                FunctionSignature functionSignature = new FunctionSignature(resultType, getArgumentTypes(scalarFunction), 0L);
                zetaSqlCatalog.addFunction(new Function(path, USER_DEFINED_JAVA_SCALAR_FUNCTIONS, ZetaSQLFunctions.FunctionEnums.Mode.SCALAR, ImmutableList.of(functionSignature)));
            } else if (function instanceof UdafImpl) {
                UdafImpl<?, ?, ?> udaf = (UdafImpl) function;
                javaUdafs.put(path, udaf.getCombineFn());
                FunctionArgumentType resultType = new FunctionArgumentType(ZetaSqlCalciteTranslationUtils.toZetaSqlType(udaf.getReturnType(typeFactory)));
                FunctionSignature functionSignature = new FunctionSignature(resultType, getArgumentTypes(udaf), 0L);
                zetaSqlCatalog.addFunction(new Function(path, USER_DEFINED_JAVA_AGGREGATE_FUNCTIONS, ZetaSQLFunctions.FunctionEnums.Mode.AGGREGATE, ImmutableList.of(functionSignature)));
            } else {
                throw new IllegalArgumentException(String.format("Function %s has unrecognized implementation type %s.", functionName, function.getClass().getName()));
            }
        }
    }
}

Example 2 with FunctionArgumentType

use of com.google.zetasql.FunctionArgumentType in project beam by apache.

the class BeamZetaSqlCatalog method addFunction.

void addFunction(ResolvedNodes.ResolvedCreateFunctionStmt createFunctionStmt) {
    String functionGroup = getFunctionGroup(createFunctionStmt);
    switch(functionGroup) {
        case USER_DEFINED_SQL_FUNCTIONS:
            sqlScalarUdfs.put(createFunctionStmt.getNamePath(), createFunctionStmt);
            break;
        case USER_DEFINED_JAVA_SCALAR_FUNCTIONS:
            String functionName = String.join(".", createFunctionStmt.getNamePath());
            for (FunctionArgumentType argumentType : createFunctionStmt.getSignature().getFunctionArgumentList()) {
                Type type = argumentType.getType();
                if (type == null) {
                    throw new UnsupportedOperationException("UDF templated argument types are not supported.");
                }
                validateJavaUdfZetaSqlType(type, functionName);
            }
            if (createFunctionStmt.getReturnType() == null) {
                throw new IllegalArgumentException("UDF return type must not be null.");
            }
            validateJavaUdfZetaSqlType(createFunctionStmt.getReturnType(), functionName);
            String jarPath = getJarPath(createFunctionStmt);
            ScalarFn scalarFn = javaUdfLoader.loadScalarFunction(createFunctionStmt.getNamePath(), jarPath);
            Method method = ScalarFnReflector.getApplyMethod(scalarFn);
            javaScalarUdfs.put(createFunctionStmt.getNamePath(), UserFunctionDefinitions.JavaScalarFunction.create(method, jarPath));
            break;
        case USER_DEFINED_JAVA_AGGREGATE_FUNCTIONS:
            jarPath = getJarPath(createFunctionStmt);
            // Try loading the aggregate function just to make sure it exists. LazyAggregateCombineFn
            // will need to fetch it again at runtime.
            javaUdfLoader.loadAggregateFunction(createFunctionStmt.getNamePath(), jarPath);
            Combine.CombineFn<?, ?, ?> combineFn = new LazyAggregateCombineFn<>(createFunctionStmt.getNamePath(), jarPath);
            javaUdafs.put(createFunctionStmt.getNamePath(), combineFn);
            break;
        default:
            throw new IllegalArgumentException(String.format("Encountered unrecognized function group %s.", functionGroup));
    }
    zetaSqlCatalog.addFunction(new Function(createFunctionStmt.getNamePath(), functionGroup, createFunctionStmt.getIsAggregate() ? ZetaSQLFunctions.FunctionEnums.Mode.AGGREGATE : ZetaSQLFunctions.FunctionEnums.Mode.SCALAR, ImmutableList.of(createFunctionStmt.getSignature())));
}

Example 3 with FunctionArgumentType

use of com.google.zetasql.FunctionArgumentType in project beam by apache.

the class BeamZetaSqlCatalog method addWindowTvfs.

@SuppressWarnings({ // customContext and volatility are in fact nullable, but they are missing the
"nullness" // annotation upstream. TODO Unsuppress when this is fixed in ZetaSQL.
})
private void addWindowTvfs() {
    FunctionArgumentType retType = new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_RELATION);
    FunctionArgumentType inputTableType = new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_RELATION);
    FunctionArgumentType descriptorType = new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_DESCRIPTOR, FunctionArgumentType.FunctionArgumentTypeOptions.builder().setDescriptorResolutionTableOffset(0).build(), 1);
    FunctionArgumentType stringType = new FunctionArgumentType(TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_STRING));
    // TUMBLE
    zetaSqlCatalog.addTableValuedFunction(new TableValuedFunction.ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF(ImmutableList.of(TVFStreamingUtils.FIXED_WINDOW_TVF), new FunctionSignature(retType, ImmutableList.of(inputTableType, descriptorType, stringType), -1), ImmutableList.of(TVFRelation.Column.create(TVFStreamingUtils.WINDOW_START, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP)), TVFRelation.Column.create(TVFStreamingUtils.WINDOW_END, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP))), null, null));
    // HOP
    zetaSqlCatalog.addTableValuedFunction(new TableValuedFunction.ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF(ImmutableList.of(TVFStreamingUtils.SLIDING_WINDOW_TVF), new FunctionSignature(retType, ImmutableList.of(inputTableType, descriptorType, stringType, stringType), -1), ImmutableList.of(TVFRelation.Column.create(TVFStreamingUtils.WINDOW_START, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP)), TVFRelation.Column.create(TVFStreamingUtils.WINDOW_END, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP))), null, null));
    // SESSION
    zetaSqlCatalog.addTableValuedFunction(new TableValuedFunction.ForwardInputSchemaToOutputSchemaWithAppendedColumnTVF(ImmutableList.of(TVFStreamingUtils.SESSION_WINDOW_TVF), new FunctionSignature(retType, ImmutableList.of(inputTableType, descriptorType, descriptorType, stringType), -1), ImmutableList.of(TVFRelation.Column.create(TVFStreamingUtils.WINDOW_START, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP)), TVFRelation.Column.create(TVFStreamingUtils.WINDOW_END, TypeFactory.createSimpleType(ZetaSQLType.TypeKind.TYPE_TIMESTAMP))), null, null));
}