Examples with FixedOutputSchemaTVF com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF used on opensource projects

Example 1 with FixedOutputSchemaTVF

use of com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF in project beam by apache.

the class TVFScanConverter method getInputs.

@Override
public List<ResolvedNode> getInputs(ResolvedTVFScan zetaNode) {
    List<ResolvedNode> inputs = new ArrayList();
    if (zetaNode.getTvf() != null && context.getUserDefinedTableValuedFunctions().containsKey(zetaNode.getTvf().getNamePath())) {
        inputs.add(context.getUserDefinedTableValuedFunctions().get(zetaNode.getTvf().getNamePath()));
    }
    for (ResolvedFunctionArgument argument : zetaNode.getArgumentList()) {
        if (argument.getScan() != null) {
            inputs.add(argument.getScan());
        }
    }
    // Extract ResolvedArguments for solving ResolvedArgumentRef in later conversion.
    if (zetaNode.getTvf() instanceof FixedOutputSchemaTVF) {
        FileDescriptorSetsBuilder temp = new FileDescriptorSetsBuilder();
        // TODO: migrate to public Java API to retrieve FunctionSignature.
        TableValuedFunctionProto tableValuedFunctionProto = zetaNode.getTvf().serialize(temp);
        for (int i = 0; i < tableValuedFunctionProto.getSignature().getArgumentList().size(); i++) {
            String argumentName = tableValuedFunctionProto.getSignature().getArgument(i).getOptions().getArgumentName();
            if (zetaNode.getArgumentList().get(i).nodeKind() == ResolvedNodeKind.RESOLVED_FUNCTION_ARGUMENT) {
                ResolvedFunctionArgument resolvedTVFArgument = zetaNode.getArgumentList().get(i);
                if (resolvedTVFArgument.getExpr().nodeKind() == ResolvedNodeKind.RESOLVED_LITERAL) {
                    ResolvedLiteral literal = (ResolvedLiteral) resolvedTVFArgument.getExpr();
                    context.addToFunctionArgumentRefMapping(argumentName, getExpressionConverter().convertResolvedLiteral(literal));
                }
            }
        }
    }
    return inputs;
}

Example 2 with FixedOutputSchemaTVF

use of com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF in project beam by apache.

the class ExpressionConverter method convertTableValuedFunction.

/**
 * Convert a TableValuedFunction in ZetaSQL to a RexCall in Calcite.
 */
public RexCall convertTableValuedFunction(RelNode input, TableValuedFunction tvf, List<ResolvedNodes.ResolvedFunctionArgument> argumentList, List<ResolvedColumn> inputTableColumns) {
    ResolvedColumn wmCol;
    // Handle builtin windowing TVF.
    switch(tvf.getName()) {
        case TVFStreamingUtils.FIXED_WINDOW_TVF:
            // TUMBLE tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.TUMBLE.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()));
        case TVFStreamingUtils.SLIDING_WINDOW_TVF:
            // HOP tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.HOP.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
        case TVFStreamingUtils.SESSION_WINDOW_TVF:
            // SESSION tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            // SESSION tvf's third argument is descriptor.
            List<ResolvedColumn> keyCol = extractSessionKeyColumnFromDescriptor(argumentList.get(2).getDescriptorArg());
            List<RexNode> operands = new ArrayList<>();
            operands.add(convertRelNodeToRexRangeRef(input));
            operands.add(convertResolvedColumnToRexInputRef(wmCol, inputTableColumns));
            operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
            operands.addAll(convertResolvedColumnsToRexInputRef(keyCol, inputTableColumns));
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.SESSION.name()), operands);
    }
    if (tvf instanceof FixedOutputSchemaTVF) {
        FixedOutputSchemaTVF fixedOutputSchemaTVF = (FixedOutputSchemaTVF) tvf;
        return (RexCall) rexBuilder().makeCall(new ZetaSqlUserDefinedSQLNativeTableValuedFunction(new SqlIdentifier(tvf.getName(), SqlParserPos.ZERO), opBinding -> {
            List<RelDataTypeField> relDataTypeFields = convertTVFRelationColumnsToRelDataTypeFields(fixedOutputSchemaTVF.getOutputSchema().getColumns());
            return new RelRecordType(relDataTypeFields);
        }, null, null, null, null));
    }
    throw new UnsupportedOperationException("Does not support table-valued function: " + tvf.getName());
}