Search in sources :

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);
    }
}
Also used : Function(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function) ScalarFn(org.apache.beam.sdk.extensions.sql.udf.ScalarFn) CalciteSchema(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteSchema) SchemaPlus(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.SchemaPlus) LazyAggregateCombineFn(org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn) JavaUdfLoader(org.apache.beam.sdk.extensions.sql.impl.JavaUdfLoader) SqlCharStringLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlCharStringLiteral) Method(java.lang.reflect.Method)

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

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