Search in sources :

Example 1 with Function

use of com.google.zetasql.Function 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()));
            }
        }
    }
}
Also used : ScalarFunctionImpl(org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) Method(java.lang.reflect.Method) FunctionSignature(com.google.zetasql.FunctionSignature) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) TableValuedFunction(com.google.zetasql.TableValuedFunction) Function(com.google.zetasql.Function) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter)

Example 2 with Function

use of com.google.zetasql.Function 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())));
}
Also used : TableValuedFunction(com.google.zetasql.TableValuedFunction) Function(com.google.zetasql.Function) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) Type(com.google.zetasql.Type) ZetaSQLType(com.google.zetasql.ZetaSQLType) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) ScalarFn(org.apache.beam.sdk.extensions.sql.udf.ScalarFn) Combine(org.apache.beam.sdk.transforms.Combine) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) LazyAggregateCombineFn(org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn) Method(java.lang.reflect.Method)

Aggregations

Function (com.google.zetasql.Function)2 FunctionArgumentType (com.google.zetasql.FunctionArgumentType)2 TableValuedFunction (com.google.zetasql.TableValuedFunction)2 Method (java.lang.reflect.Method)2 FunctionSignature (com.google.zetasql.FunctionSignature)1 Type (com.google.zetasql.Type)1 ZetaSQLType (com.google.zetasql.ZetaSQLType)1 LazyAggregateCombineFn (org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn)1 ScalarFunctionImpl (org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl)1 UdafImpl (org.apache.beam.sdk.extensions.sql.impl.UdafImpl)1 ScalarFn (org.apache.beam.sdk.extensions.sql.udf.ScalarFn)1 Combine (org.apache.beam.sdk.transforms.Combine)1 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)1 FunctionParameter (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter)1