Search in sources :

Example 1 with UdafImpl

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

the class AggregateScanConverter method convertAggCall.

private AggregateCall convertAggCall(ResolvedComputedColumn computedColumn, int columnRefOff, int groupCount, RelNode input) {
    ResolvedAggregateFunctionCall aggregateFunctionCall = (ResolvedAggregateFunctionCall) computedColumn.getExpr();
    // Reject AVG(INT64)
    if (aggregateFunctionCall.getFunction().getName().equals("avg")) {
        FunctionSignature signature = aggregateFunctionCall.getSignature();
        if (signature.getFunctionArgumentList().get(0).getType().getKind().equals(TypeKind.TYPE_INT64)) {
            throw new UnsupportedOperationException(AVG_ILLEGAL_LONG_INPUT_TYPE);
        }
    }
    // Reject aggregation DISTINCT
    if (aggregateFunctionCall.getDistinct()) {
        throw new UnsupportedOperationException("Does not support " + aggregateFunctionCall.getFunction().getSqlName() + " DISTINCT. 'SELECT DISTINCT' syntax could be used to deduplicate before" + " aggregation.");
    }
    final SqlAggFunction sqlAggFunction;
    if (aggregateFunctionCall.getFunction().getGroup().equals(BeamZetaSqlCatalog.USER_DEFINED_JAVA_AGGREGATE_FUNCTIONS)) {
        // Create a new operator for user-defined functions.
        SqlReturnTypeInference typeInference = x -> ZetaSqlCalciteTranslationUtils.toCalciteType(aggregateFunctionCall.getFunction().getSignatureList().get(0).getResultType().getType(), // TODO(BEAM-9514) set nullable=true
        false, getCluster().getRexBuilder());
        UdafImpl<?, ?, ?> impl = new UdafImpl<>(getExpressionConverter().userFunctionDefinitions.javaAggregateFunctions().get(aggregateFunctionCall.getFunction().getNamePath()));
        sqlAggFunction = SqlOperators.createUdafOperator(aggregateFunctionCall.getFunction().getName(), typeInference, impl);
    } else {
        // Look up builtin functions in SqlOperatorMappingTable.
        sqlAggFunction = (SqlAggFunction) SqlOperatorMappingTable.create(aggregateFunctionCall);
        if (sqlAggFunction == null) {
            throw new UnsupportedOperationException("Does not support ZetaSQL aggregate function: " + aggregateFunctionCall.getFunction().getName());
        }
    }
    List<Integer> argList = new ArrayList<>();
    ResolvedAggregateFunctionCall expr = ((ResolvedAggregateFunctionCall) computedColumn.getExpr());
    List<ZetaSQLResolvedNodeKind.ResolvedNodeKind> resolvedNodeKinds = Arrays.asList(RESOLVED_CAST, RESOLVED_COLUMN_REF, RESOLVED_GET_STRUCT_FIELD);
    for (int i = 0; i < expr.getArgumentList().size(); i++) {
        // Throw an error if aggregate function's input isn't either a ColumnRef or a cast(ColumnRef).
        // TODO: is there a general way to handle aggregation calls conversion?
        ZetaSQLResolvedNodeKind.ResolvedNodeKind resolvedNodeKind = expr.getArgumentList().get(i).nodeKind();
        if (i == 0 && resolvedNodeKinds.contains(resolvedNodeKind)) {
            argList.add(columnRefOff);
        } else if (i > 0 && resolvedNodeKind == RESOLVED_LITERAL) {
            continue;
        } else {
            throw new UnsupportedOperationException("Aggregate function only accepts Column Reference or CAST(Column Reference) as the first argument and " + "Literals as subsequent arguments as its inputs");
        }
    }
    String aggName = getTrait().resolveAlias(computedColumn.getColumn());
    return AggregateCall.create(sqlAggFunction, false, false, false, argList, -1, null, RelCollations.EMPTY, groupCount, input, // When we pass null as the return type, Calcite infers it for us.
    null, aggName);
}
Also used : ZetaSQLResolvedNodeKind(com.google.zetasql.ZetaSQLResolvedNodeKind) IntStream(java.util.stream.IntStream) AggregateCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.AggregateCall) ZetaSqlCalciteTranslationUtils(org.apache.beam.sdk.extensions.sql.zetasql.ZetaSqlCalciteTranslationUtils) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) Arrays(java.util.Arrays) ResolvedNode(com.google.zetasql.resolvedast.ResolvedNode) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) ResolvedAggregateFunctionCall(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall) ResolvedAggregateScan(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan) RelCollations(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollations) ArrayList(java.util.ArrayList) ImmutableBitSet(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet) SqlAggFunction(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlAggFunction) RESOLVED_COLUMN_REF(com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_COLUMN_REF) LogicalProject(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.logical.LogicalProject) FunctionSignature(com.google.zetasql.FunctionSignature) RESOLVED_CAST(com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_CAST) RESOLVED_GET_STRUCT_FIELD(com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_GET_STRUCT_FIELD) RESOLVED_LITERAL(com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_LITERAL) Collectors(java.util.stream.Collectors) LogicalAggregate(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.logical.LogicalAggregate) List(java.util.List) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) ResolvedComputedColumn(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn) ResolvedExpr(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr) TypeKind(com.google.zetasql.ZetaSQLType.TypeKind) SqlReturnTypeInference(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlReturnTypeInference) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) BeamZetaSqlCatalog(org.apache.beam.sdk.extensions.sql.zetasql.BeamZetaSqlCatalog) Collections(java.util.Collections) ArrayList(java.util.ArrayList) SqlAggFunction(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlAggFunction) FunctionSignature(com.google.zetasql.FunctionSignature) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) ResolvedAggregateFunctionCall(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall) SqlReturnTypeInference(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlReturnTypeInference) ZetaSQLResolvedNodeKind(com.google.zetasql.ZetaSQLResolvedNodeKind) ZetaSQLResolvedNodeKind(com.google.zetasql.ZetaSQLResolvedNodeKind)

Example 2 with UdafImpl

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

the class TypedCombineFnDelegateTest method testParameterExtractionFromCombineFn_CombineFnDelegate_WithListInsteadOfArray.

@Test
public void testParameterExtractionFromCombineFn_CombineFnDelegate_WithListInsteadOfArray() {
    Combine.BinaryCombineFn<List<List<String>>> max = Max.of((Comparator<List<List<String>>> & Serializable) (a, b) -> Integer.compare(a.get(0).get(0).length(), b.get(0).get(0).length()));
    UdafImpl<List<List<String>>, Combine.Holder<List<List<String>>>, List<List<String>>> udaf = new UdafImpl<>(new TypedCombineFnDelegate<List<List<String>>, Combine.Holder<List<List<String>>>, List<List<String>>>(max) {
    });
    RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
    List<FunctionParameter> parameters = udaf.getParameters();
    assertEquals(1, parameters.size());
    assertEquals(SqlTypeName.ARRAY, parameters.get(0).getType(typeFactory).getSqlTypeName());
}
Also used : RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter) Combine(org.apache.beam.sdk.transforms.Combine) Test(org.junit.Test) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) Serializable(java.io.Serializable) SqlTypeName(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName) List(java.util.List) Max(org.apache.beam.sdk.transforms.Max) Rule(org.junit.Rule) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) RelDataTypeSystem(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeSystem) ExpectedException(org.junit.rules.ExpectedException) Assert.assertEquals(org.junit.Assert.assertEquals) Serializable(java.io.Serializable) Combine(org.apache.beam.sdk.transforms.Combine) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) List(java.util.List) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter) Test(org.junit.Test)

Example 3 with UdafImpl

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

the class TypedCombineFnDelegateTest method testParameterExtractionFromCombineFn_CombineFnDelegate.

@Test
public void testParameterExtractionFromCombineFn_CombineFnDelegate() {
    Combine.BinaryCombineFn<String> max = Max.of((Comparator<String> & Serializable) (a, b) -> Integer.compare(a.length(), b.length()));
    UdafImpl<String, Combine.Holder<String>, String> udaf = new UdafImpl<>(new TypedCombineFnDelegate<String, Combine.Holder<String>, String>(max) {
    });
    RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
    List<FunctionParameter> parameters = udaf.getParameters();
    assertEquals(1, parameters.size());
    assertEquals(SqlTypeName.VARCHAR, parameters.get(0).getType(typeFactory).getSqlTypeName());
}
Also used : RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter) Combine(org.apache.beam.sdk.transforms.Combine) Test(org.junit.Test) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) Serializable(java.io.Serializable) SqlTypeName(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName) List(java.util.List) Max(org.apache.beam.sdk.transforms.Max) Rule(org.junit.Rule) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) RelDataTypeSystem(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeSystem) ExpectedException(org.junit.rules.ExpectedException) Assert.assertEquals(org.junit.Assert.assertEquals) Serializable(java.io.Serializable) Combine(org.apache.beam.sdk.transforms.Combine) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter) Test(org.junit.Test)

Example 4 with UdafImpl

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

the class TypedCombineFnDelegateTest method testParameterExtractionFromCombineFn_CombineFnDelegate_WithGenericArray.

@Test
public void testParameterExtractionFromCombineFn_CombineFnDelegate_WithGenericArray() {
    Combine.BinaryCombineFn<List<String>[]> max = Max.of((Comparator<List<String>[]> & Serializable) (a, b) -> Integer.compare(a[0].get(0).length(), b[0].get(0).length()));
    UdafImpl<List<String>[], Combine.Holder<List<String>[]>, List<String>[]> udaf = new UdafImpl<>(new TypedCombineFnDelegate<List<String>[], Combine.Holder<List<String>[]>, List<String>[]>(max) {
    });
    exceptions.expect(IllegalArgumentException.class);
    RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
    udaf.getParameters().get(0).getType(typeFactory);
}
Also used : RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter) Combine(org.apache.beam.sdk.transforms.Combine) Test(org.junit.Test) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) Serializable(java.io.Serializable) SqlTypeName(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName) List(java.util.List) Max(org.apache.beam.sdk.transforms.Max) Rule(org.junit.Rule) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) RelDataTypeSystem(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeSystem) ExpectedException(org.junit.rules.ExpectedException) Assert.assertEquals(org.junit.Assert.assertEquals) Serializable(java.io.Serializable) Combine(org.apache.beam.sdk.transforms.Combine) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) Comparator(java.util.Comparator) JavaTypeFactoryImpl(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl) RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) List(java.util.List) Test(org.junit.Test)

Example 5 with UdafImpl

use of org.apache.beam.sdk.extensions.sql.impl.UdafImpl 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()));
            }
        }
    }
}
Also used : ScalarFunctionImpl(org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) Method(java.lang.reflect.Method) FunctionSignature(com.google.zetasql.FunctionSignature) UdafImpl(org.apache.beam.sdk.extensions.sql.impl.UdafImpl) TableValuedFunction(com.google.zetasql.TableValuedFunction) Function(com.google.zetasql.Function) FunctionParameter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter)

Aggregations

UdafImpl (org.apache.beam.sdk.extensions.sql.impl.UdafImpl)5 List (java.util.List)4 FunctionParameter (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.FunctionParameter)4 Serializable (java.io.Serializable)3 Comparator (java.util.Comparator)3 Combine (org.apache.beam.sdk.transforms.Combine)3 Max (org.apache.beam.sdk.transforms.Max)3 JavaTypeFactoryImpl (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.JavaTypeFactoryImpl)3 RelDataTypeFactory (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory)3 RelDataTypeSystem (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeSystem)3 SqlTypeName (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName)3 Assert.assertEquals (org.junit.Assert.assertEquals)3 Rule (org.junit.Rule)3 Test (org.junit.Test)3 ExpectedException (org.junit.rules.ExpectedException)3 FunctionSignature (com.google.zetasql.FunctionSignature)2 Function (com.google.zetasql.Function)1 FunctionArgumentType (com.google.zetasql.FunctionArgumentType)1 TableValuedFunction (com.google.zetasql.TableValuedFunction)1 ZetaSQLResolvedNodeKind (com.google.zetasql.ZetaSQLResolvedNodeKind)1