Examples with Function org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function used on opensource projects

Example 6 with Function

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function in project beam by apache.

the class ExpressionConverter method convertResolvedFunctionCall.

private RexNode convertResolvedFunctionCall(ResolvedFunctionCall functionCall, @Nullable List<ResolvedColumn> columnList, @Nullable List<RelDataTypeField> fieldList, Map<String, RexNode> outerFunctionArguments) {
    final String funGroup = functionCall.getFunction().getGroup();
    final String funName = functionCall.getFunction().getName();
    SqlOperator op = SqlOperatorMappingTable.create(functionCall);
    List<RexNode> operands = new ArrayList<>();
    if (PRE_DEFINED_WINDOW_FUNCTIONS.equals(funGroup)) {
        switch(funName) {
            case FIXED_WINDOW:
            case SESSION_WINDOW:
                // TODO: check size and type of window function argument list.
                // Add ts column reference to operands.
                operands.add(convertRexNodeFromResolvedExpr(functionCall.getArgumentList().get(0), columnList, fieldList, outerFunctionArguments));
                // Add fixed window size or session window gap to operands.
                operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) functionCall.getArgumentList().get(1)));
                break;
            case SLIDING_WINDOW:
                // Add ts column reference to operands.
                operands.add(convertRexNodeFromResolvedExpr(functionCall.getArgumentList().get(0), columnList, fieldList, outerFunctionArguments));
                // add sliding window emit frequency to operands.
                operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) functionCall.getArgumentList().get(1)));
                // add sliding window size to operands.
                operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) functionCall.getArgumentList().get(2)));
                break;
            default:
                throw new UnsupportedOperationException("Unsupported function: " + funName + ". Only support TUMBLE, HOP, and SESSION now.");
        }
    } else if (ZETASQL_FUNCTION_GROUP_NAME.equals(funGroup)) {
        if (op == null) {
            Type returnType = functionCall.getSignature().getResultType().getType();
            if (returnType != null) {
                op = SqlOperators.createZetaSqlFunction(funName, ZetaSqlCalciteTranslationUtils.toCalciteType(returnType, false, rexBuilder()).getSqlTypeName());
            } else {
                throw new UnsupportedOperationException("Does not support ZetaSQL function: " + funName);
            }
        }
        for (ResolvedExpr expr : functionCall.getArgumentList()) {
            operands.add(convertRexNodeFromResolvedExpr(expr, columnList, fieldList, outerFunctionArguments));
        }
    } else if (USER_DEFINED_JAVA_SCALAR_FUNCTIONS.equals(funGroup)) {
        UserFunctionDefinitions.JavaScalarFunction javaScalarFunction = userFunctionDefinitions.javaScalarFunctions().get(functionCall.getFunction().getNamePath());
        ArrayList<RexNode> innerFunctionArguments = new ArrayList<>();
        for (int i = 0; i < functionCall.getArgumentList().size(); i++) {
            ResolvedExpr argExpr = functionCall.getArgumentList().get(i);
            RexNode argNode = convertRexNodeFromResolvedExpr(argExpr, columnList, fieldList, outerFunctionArguments);
            innerFunctionArguments.add(argNode);
        }
        return rexBuilder().makeCall(SqlOperators.createUdfOperator(functionCall.getFunction().getName(), javaScalarFunction.method(), USER_DEFINED_JAVA_SCALAR_FUNCTIONS, javaScalarFunction.jarPath()), innerFunctionArguments);
    } else if (USER_DEFINED_SQL_FUNCTIONS.equals(funGroup)) {
        ResolvedCreateFunctionStmt createFunctionStmt = userFunctionDefinitions.sqlScalarFunctions().get(functionCall.getFunction().getNamePath());
        ResolvedExpr functionExpression = createFunctionStmt.getFunctionExpression();
        ImmutableMap.Builder<String, RexNode> innerFunctionArguments = ImmutableMap.builder();
        for (int i = 0; i < functionCall.getArgumentList().size(); i++) {
            String argName = createFunctionStmt.getArgumentNameList().get(i);
            ResolvedExpr argExpr = functionCall.getArgumentList().get(i);
            RexNode argNode = convertRexNodeFromResolvedExpr(argExpr, columnList, fieldList, outerFunctionArguments);
            innerFunctionArguments.put(argName, argNode);
        }
        return this.convertRexNodeFromResolvedExpr(functionExpression, columnList, fieldList, innerFunctionArguments.build());
    } else {
        throw new UnsupportedOperationException("Does not support function group: " + funGroup);
    }
    SqlOperatorRewriter rewriter = SqlOperatorMappingTable.ZETASQL_FUNCTION_TO_CALCITE_SQL_OPERATOR_REWRITER.get(funName);
    if (rewriter != null) {
        return rewriter.apply(rexBuilder(), operands);
    } else {
        return rexBuilder().makeCall(op, operands);
    }
}

Example 7 with Function

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function in project beam by apache.

the class RelConverter method createOneRow.

// This function creates a single dummy input row for queries that don't read from a table.
// For example: SELECT "hello"
// The code is copy-pasted from Calcite's LogicalValues.createOneRow() with a single line
// change: SqlTypeName.INTEGER replaced by SqlTypeName.BIGINT.
// Would like to call LogicalValues.createOneRow() directly, but it uses type SqlTypeName.INTEGER
// which corresponds to TypeKind.TYPE_INT32 in ZetaSQL, a type not supported in ZetaSQL
// PRODUCT_EXTERNAL mode. See
// https://github.com/google/zetasql/blob/c610a21ffdc110293c1c7bd255a2674ebc7ec7a8/java/com/google/zetasql/TypeFactory.java#L61
static LogicalValues createOneRow(RelOptCluster cluster) {
    final RelDataType rowType = cluster.getTypeFactory().builder().add("ZERO", SqlTypeName.BIGINT).nullable(false).build();
    final ImmutableList<ImmutableList<RexLiteral>> tuples = ImmutableList.of(ImmutableList.of(cluster.getRexBuilder().makeExactLiteral(BigDecimal.ZERO, rowType.getFieldList().get(0).getType())));
    return LogicalValues.create(cluster, rowType, tuples);
}

Example 8 with Function

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function in project beam by apache.

the class SqlInOperatorRewriter method apply.

@Override
public RexNode apply(RexBuilder rexBuilder, List<RexNode> operands) {
    Preconditions.checkArgument(operands.size() >= 2, "IN should have at least two arguments in function call.");
    final RexNode arg = operands.get(0);
    final List<RexNode> ranges = ImmutableList.copyOf(operands.subList(1, operands.size()));
    // ZetaSQL has weird behavior for NULL...
    for (RexNode node : ranges) {
        if (node instanceof RexLiteral && ((RexLiteral) node).isNull()) {
            throw new UnsupportedOperationException("IN NULL unsupported");
        }
    }
    return rexBuilder.makeIn(arg, ranges);
}

Example 9 with Function

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function in project beam by apache.

the class SqlOperators method createUdfOperator.

private static SqlUserDefinedFunction createUdfOperator(String name, Method method, final SqlSyntax syntax, String funGroup, String jarPath) {
    Function function = ZetaSqlScalarFunctionImpl.create(method, funGroup, jarPath);
    final RelDataTypeFactory typeFactory = createTypeFactory();
    List<RelDataType> argTypes = new ArrayList<>();
    List<SqlTypeFamily> typeFamilies = new ArrayList<>();
    for (FunctionParameter o : function.getParameters()) {
        final RelDataType type = o.getType(typeFactory);
        argTypes.add(type);
        typeFamilies.add(Util.first(type.getSqlTypeName().getFamily(), SqlTypeFamily.ANY));
    }
    final FamilyOperandTypeChecker typeChecker = OperandTypes.family(typeFamilies, i -> function.getParameters().get(i).isOptional());
    final List<RelDataType> paramTypes = toSql(typeFactory, argTypes);
    return new SqlUserDefinedFunction(new SqlIdentifier(name, SqlParserPos.ZERO), infer((ScalarFunction) function), InferTypes.explicit(argTypes), typeChecker, paramTypes, function) {

        @Override
        public SqlSyntax getSyntax() {
            return syntax;
        }
    };
}

Example 10 with Function

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Function 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);
}