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()));
}
}
}
}
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())));
}
Aggregations