Examples with ScalarFunctionImpl org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl used on opensource projects

Example 1 with ScalarFunctionImpl

use of org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl 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()));
            }
        }
    }
}