use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral 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);
}
}
use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral in project beam by apache.
the class ExpressionConverter method convertIntervalToRexIntervalLiteral.
private RexNode convertIntervalToRexIntervalLiteral(ResolvedLiteral resolvedLiteral) {
if (resolvedLiteral.getType().getKind() != TYPE_STRING) {
throw new ZetaSqlException(INTERVAL_FORMAT_MSG);
}
String valStr = resolvedLiteral.getValue().getStringValue();
List<String> stringList = Arrays.stream(valStr.split(" ")).filter(s -> !s.isEmpty()).collect(Collectors.toList());
if (stringList.size() != 3) {
throw new ZetaSqlException(INTERVAL_FORMAT_MSG);
}
if (!Ascii.toUpperCase(stringList.get(0)).equals("INTERVAL")) {
throw new ZetaSqlException(INTERVAL_FORMAT_MSG);
}
long intervalValue;
try {
intervalValue = Long.parseLong(stringList.get(1));
} catch (NumberFormatException e) {
throw new ZetaSqlException(Status.UNIMPLEMENTED.withDescription(INTERVAL_FORMAT_MSG).withCause(e).asRuntimeException());
}
String intervalDatepart = Ascii.toUpperCase(stringList.get(2));
return createCalciteIntervalRexLiteral(intervalValue, intervalDatepart);
}
use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral 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;
}
use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral in project beam by apache.
the class ExpressionConverter method convertRexNodeFromComputedColumnWithFieldList.
private RexNode convertRexNodeFromComputedColumnWithFieldList(ResolvedComputedColumn column, List<ResolvedColumn> columnList, List<RelDataTypeField> fieldList, int windowFieldIndex) {
if (column.getExpr().nodeKind() != RESOLVED_FUNCTION_CALL) {
return convertRexNodeFromResolvedExpr(column.getExpr(), columnList, fieldList, ImmutableMap.of());
}
ResolvedFunctionCall functionCall = (ResolvedFunctionCall) column.getExpr();
// TODO: is there any other illegal case?
if (functionCall.getFunction().getName().equals(FIXED_WINDOW) || functionCall.getFunction().getName().equals(SLIDING_WINDOW) || functionCall.getFunction().getName().equals(SESSION_WINDOW)) {
throw new ZetaSqlException(functionCall.getFunction().getName() + " shouldn't appear in SELECT exprlist.");
}
if (!functionCall.getFunction().getGroup().equals(PRE_DEFINED_WINDOW_FUNCTIONS)) {
// non-window function should still go through normal FunctionCall conversion process.
return convertRexNodeFromResolvedExpr(column.getExpr(), columnList, fieldList, ImmutableMap.of());
}
// ONLY window_start and window_end should arrive here.
// TODO: Have extra verification here to make sure window start/end functions have the same
// parameter with window function.
List<RexNode> operands = new ArrayList<>();
switch(functionCall.getFunction().getName()) {
case FIXED_WINDOW_START:
case SLIDING_WINDOW_START:
case SESSION_WINDOW_START:
// in Calcite.
case SESSION_WINDOW_END:
return rexBuilder().makeInputRef(fieldList.get(windowFieldIndex).getType(), windowFieldIndex);
case FIXED_WINDOW_END:
operands.add(rexBuilder().makeInputRef(fieldList.get(windowFieldIndex).getType(), windowFieldIndex));
// TODO: check window_end 's duration is the same as it's aggregate window.
operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) functionCall.getArgumentList().get(0)));
return rexBuilder().makeCall(SqlOperators.ZETASQL_TIMESTAMP_ADD, operands);
case SLIDING_WINDOW_END:
operands.add(rexBuilder().makeInputRef(fieldList.get(windowFieldIndex).getType(), windowFieldIndex));
operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) functionCall.getArgumentList().get(1)));
return rexBuilder().makeCall(SqlOperators.ZETASQL_TIMESTAMP_ADD, operands);
default:
throw new UnsupportedOperationException("Does not support window start/end: " + functionCall.getFunction().getName());
}
}
use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral 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());
}
Aggregations