Examples with UdfProvider org.apache.beam.sdk.extensions.sql.udf.UdfProvider used on opensource projects

Example 1 with UdfProvider

use of org.apache.beam.sdk.extensions.sql.udf.UdfProvider in project beam by apache.

the class JavaUdfLoader method loadJar.

private FunctionDefinitions loadJar(String jarPath) throws IOException {
    if (functionCache.containsKey(jarPath)) {
        LOG.debug("Using cached function definitions from {}", jarPath);
        return functionCache.get(jarPath);
    }
    ClassLoader classLoader = createClassLoader(jarPath);
    Map<List<String>, ScalarFn> scalarFunctions = new HashMap<>();
    Map<List<String>, AggregateFn> aggregateFunctions = new HashMap<>();
    Iterator<UdfProvider> providers = getUdfProviders(classLoader);
    int providersCount = 0;
    while (providers.hasNext()) {
        providersCount++;
        UdfProvider provider = providers.next();
        provider.userDefinedScalarFunctions().forEach((functionName, implementation) -> {
            List<String> functionPath = ImmutableList.copyOf(functionName.split("\\."));
            if (scalarFunctions.containsKey(functionPath)) {
                throw new IllegalArgumentException(String.format("Found multiple definitions of scalar function %s in %s.", functionName, jarPath));
            }
            scalarFunctions.put(functionPath, implementation);
        });
        provider.userDefinedAggregateFunctions().forEach((functionName, implementation) -> {
            List<String> functionPath = ImmutableList.copyOf(functionName.split("\\."));
            if (aggregateFunctions.containsKey(functionPath)) {
                throw new IllegalArgumentException(String.format("Found multiple definitions of aggregate function %s in %s.", functionName, jarPath));
            }
            aggregateFunctions.put(functionPath, implementation);
        });
    }
    if (providersCount == 0) {
        throw new ProviderNotFoundException(String.format("No %s implementation found in %s. Create a class implementing %s and annotate it with @AutoService(%s.class).", UdfProvider.class.getSimpleName(), jarPath, UdfProvider.class.getSimpleName(), UdfProvider.class.getSimpleName()));
    }
    LOG.info("Loaded {} implementations of {} from {} with {} scalar function(s).", providersCount, UdfProvider.class.getSimpleName(), jarPath, scalarFunctions.size());
    FunctionDefinitions userFunctionDefinitions = FunctionDefinitions.newBuilder().setScalarFunctions(ImmutableMap.copyOf(scalarFunctions)).setAggregateFunctions(ImmutableMap.copyOf(aggregateFunctions)).build();
    functionCache.put(jarPath, userFunctionDefinitions);
    return userFunctionDefinitions;
}