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

Example 1 with LazyAggregateCombineFn

use of org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn in project beam by apache.

the class SqlCreateFunction method execute.

@Override
public void execute(CalcitePrepare.Context context) {
    final Pair<CalciteSchema, String> pair = SqlDdlNodes.schema(context, true, functionName);
    SchemaPlus schema = pair.left.plus();
    String lastName = pair.right;
    if (!schema.getFunctions(lastName).isEmpty()) {
        throw SqlUtil.newContextException(functionName.getParserPosition(), RESOURCE.internal(String.format("Function %s is already defined.", lastName)));
    }
    JavaUdfLoader udfLoader = new JavaUdfLoader();
    // TODO(BEAM-12355) Support qualified function names.
    List<String> functionPath = ImmutableList.of(lastName);
    if (!(jarPath instanceof SqlCharStringLiteral)) {
        throw SqlUtil.newContextException(jarPath.getParserPosition(), RESOURCE.internal("Jar path is not instanceof SqlCharStringLiteral."));
    }
    String unquotedJarPath = ((SqlCharStringLiteral) jarPath).getNlsString().getValue();
    if (isAggregate) {
        // Try loading the aggregate function just to make sure it exists. LazyAggregateCombineFn will
        // need to fetch it again at runtime.
        udfLoader.loadAggregateFunction(functionPath, unquotedJarPath);
        LazyAggregateCombineFn<?, ?, ?> combineFn = new LazyAggregateCombineFn<>(functionPath, unquotedJarPath);
        schema.add(lastName, combineFn.getUdafImpl());
    } else {
        ScalarFn scalarFn = udfLoader.loadScalarFunction(functionPath, unquotedJarPath);
        Method method = ScalarFnReflector.getApplyMethod(scalarFn);
        Function function = ScalarFunctionImpl.create(method, unquotedJarPath);
        schema.add(lastName, function);
    }
}

Example 2 with LazyAggregateCombineFn

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