Examples with ResolvedExpr com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr used on opensource projects

Example 1 with ResolvedExpr

use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr 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 2 with ResolvedExpr

use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr in project beam by apache.

the class AggregateScanConverter method convertAggregateScanInputScanToLogicalProject.

private LogicalProject convertAggregateScanInputScanToLogicalProject(ResolvedAggregateScan node, RelNode input) {
    // AggregateScan's input is the source of data (e.g. TableScan), which is different from the
    // design of CalciteSQL, in which the LogicalAggregate's input is a LogicalProject, whose input
    // is a LogicalTableScan. When AggregateScan's input is WithRefScan, the WithRefScan is
    // ebullient to a LogicalTableScan. So it's still required to build another LogicalProject as
    // the input of LogicalAggregate.
    List<RexNode> projects = new ArrayList<>();
    List<String> fieldNames = new ArrayList<>();
    // LogicalAggregate.
    for (ResolvedComputedColumn computedColumn : node.getGroupByList()) {
        projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(computedColumn.getExpr(), node.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of()));
        fieldNames.add(getTrait().resolveAlias(computedColumn.getColumn()));
    }
    // TODO: remove duplicate columns in projects.
    for (ResolvedComputedColumn resolvedComputedColumn : node.getAggregateList()) {
        // Should create Calcite's RexInputRef from ResolvedColumn from ResolvedComputedColumn.
        // TODO: handle aggregate function with more than one argument and handle OVER
        // TODO: is there is general way for column reference tracking and deduplication for
        // aggregation?
        ResolvedAggregateFunctionCall aggregateFunctionCall = ((ResolvedAggregateFunctionCall) resolvedComputedColumn.getExpr());
        if (aggregateFunctionCall.getArgumentList() != null && aggregateFunctionCall.getArgumentList().size() >= 1) {
            ResolvedExpr resolvedExpr = aggregateFunctionCall.getArgumentList().get(0);
            for (int i = 0; i < aggregateFunctionCall.getArgumentList().size(); i++) {
                if (i == 0) {
                    // TODO: assume aggregate function's input is either a ColumnRef or a cast(ColumnRef).
                    // TODO: user might use multiple CAST so we need to handle this rare case.
                    projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(resolvedExpr, node.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of()));
                } else {
                    projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(aggregateFunctionCall.getArgumentList().get(i)));
                }
                fieldNames.add(getTrait().resolveAlias(resolvedComputedColumn.getColumn()));
            }
        }
    }
    return LogicalProject.create(input, ImmutableList.of(), projects, fieldNames);
}