Search in sources :

Example 1 with ResolvedCreateFunctionStmt

use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateFunctionStmt 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);
    }
}
Also used : SqlOperator(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator) RexBuilder(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder) ArrayList(java.util.ArrayList) ResolvedExpr(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr) ResolvedCreateFunctionStmt(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateFunctionStmt) Type(com.google.zetasql.Type) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) RelRecordType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType) ResolvedLiteral(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

Type (com.google.zetasql.Type)1 ResolvedCreateFunctionStmt (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateFunctionStmt)1 ResolvedExpr (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr)1 ResolvedLiteral (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral)1 ArrayList (java.util.ArrayList)1 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)1 RelRecordType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType)1 RexBuilder (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder)1 RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)1 SqlOperator (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator)1