Search in sources :

Example 1 with Type

use of com.google.zetasql.Type in project beam by apache.

the class ZetaSqlCalciteTranslationUtils method toRexNode.

// Value conversion: ZetaSQL => Calcite
public static RexNode toRexNode(Value value, RexBuilder rexBuilder) {
    Type type = value.getType();
    if (value.isNull()) {
        return rexBuilder.makeNullLiteral(toCalciteType(type, true, rexBuilder));
    }
    switch(type.getKind()) {
        case TYPE_INT64:
            return rexBuilder.makeExactLiteral(new BigDecimal(value.getInt64Value()), toCalciteType(type, false, rexBuilder));
        case TYPE_DOUBLE:
            // Cannot simply call makeApproxLiteral() because +inf, -inf, and NaN cannot be represented
            // as BigDecimal. So we create wrapper functions here for these three cases such that we can
            // later recognize it and customize its unparsing in BeamBigQuerySqlDialect.
            double val = value.getDoubleValue();
            String wrapperFun = null;
            if (val == Double.POSITIVE_INFINITY) {
                wrapperFun = BeamBigQuerySqlDialect.DOUBLE_POSITIVE_INF_WRAPPER;
            } else if (val == Double.NEGATIVE_INFINITY) {
                wrapperFun = BeamBigQuerySqlDialect.DOUBLE_NEGATIVE_INF_WRAPPER;
            } else if (Double.isNaN(val)) {
                wrapperFun = BeamBigQuerySqlDialect.DOUBLE_NAN_WRAPPER;
            }
            RelDataType returnType = toCalciteType(type, false, rexBuilder);
            if (wrapperFun == null) {
                return rexBuilder.makeApproxLiteral(new BigDecimal(val), returnType);
            } else if (BeamBigQuerySqlDialect.DOUBLE_NAN_WRAPPER.equals(wrapperFun)) {
                // true, which should be false.)
                return rexBuilder.makeCall(SqlOperators.createZetaSqlFunction(wrapperFun, returnType.getSqlTypeName()), rexBuilder.makeApproxLiteral(BigDecimal.valueOf(Math.random()), returnType));
            } else {
                return rexBuilder.makeCall(SqlOperators.createZetaSqlFunction(wrapperFun, returnType.getSqlTypeName()));
            }
        case TYPE_BOOL:
            return rexBuilder.makeLiteral(value.getBoolValue());
        case TYPE_STRING:
            // If not allow cast, rexBuilder() will only build a literal with CHAR type.
            return rexBuilder.makeLiteral(value.getStringValue(), toCalciteType(type, false, rexBuilder), true);
        case TYPE_BYTES:
            return rexBuilder.makeBinaryLiteral(new ByteString(value.getBytesValue().toByteArray()));
        case TYPE_NUMERIC:
            // its unparsing in BeamBigQuerySqlDialect.
            return rexBuilder.makeCall(SqlOperators.createZetaSqlFunction(BeamBigQuerySqlDialect.NUMERIC_LITERAL_WRAPPER, toCalciteType(type, false, rexBuilder).getSqlTypeName()), rexBuilder.makeExactLiteral(value.getNumericValue(), toCalciteType(type, false, rexBuilder)));
        case TYPE_DATE:
            return rexBuilder.makeDateLiteral(dateValueToDateString(value));
        case TYPE_TIME:
            return rexBuilder.makeTimeLiteral(timeValueToTimeString(value), rexBuilder.getTypeFactory().getTypeSystem().getMaxPrecision(SqlTypeName.TIME));
        case TYPE_DATETIME:
            return rexBuilder.makeTimestampWithLocalTimeZoneLiteral(datetimeValueToTimestampString(value), rexBuilder.getTypeFactory().getTypeSystem().getMaxPrecision(SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE));
        case TYPE_TIMESTAMP:
            return rexBuilder.makeTimestampLiteral(timestampValueToTimestampString(value), rexBuilder.getTypeFactory().getTypeSystem().getMaxPrecision(SqlTypeName.TIMESTAMP));
        case TYPE_ARRAY:
            return arrayValueToRexNode(value, rexBuilder);
        case TYPE_STRUCT:
            return structValueToRexNode(value, rexBuilder);
        case // internal only, used for DateTimestampPart
        TYPE_ENUM:
            return enumValueToRexNode(value, rexBuilder);
        default:
            throw new UnsupportedOperationException("Unknown ZetaSQL type: " + type.getKind().name());
    }
}
Also used : Type(com.google.zetasql.Type) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) StructType(com.google.zetasql.StructType) ByteString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.util.ByteString) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) ByteString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.util.ByteString) TimeString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.TimeString) TimestampString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.TimestampString) DateString(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.DateString) BigDecimal(java.math.BigDecimal)

Example 2 with Type

use of com.google.zetasql.Type 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)

Example 3 with Type

use of com.google.zetasql.Type in project beam by apache.

the class BeamZetaSqlCatalog method addFunction.

void addFunction(ResolvedNodes.ResolvedCreateFunctionStmt createFunctionStmt) {
    String functionGroup = getFunctionGroup(createFunctionStmt);
    switch(functionGroup) {
        case USER_DEFINED_SQL_FUNCTIONS:
            sqlScalarUdfs.put(createFunctionStmt.getNamePath(), createFunctionStmt);
            break;
        case USER_DEFINED_JAVA_SCALAR_FUNCTIONS:
            String functionName = String.join(".", createFunctionStmt.getNamePath());
            for (FunctionArgumentType argumentType : createFunctionStmt.getSignature().getFunctionArgumentList()) {
                Type type = argumentType.getType();
                if (type == null) {
                    throw new UnsupportedOperationException("UDF templated argument types are not supported.");
                }
                validateJavaUdfZetaSqlType(type, functionName);
            }
            if (createFunctionStmt.getReturnType() == null) {
                throw new IllegalArgumentException("UDF return type must not be null.");
            }
            validateJavaUdfZetaSqlType(createFunctionStmt.getReturnType(), functionName);
            String jarPath = getJarPath(createFunctionStmt);
            ScalarFn scalarFn = javaUdfLoader.loadScalarFunction(createFunctionStmt.getNamePath(), jarPath);
            Method method = ScalarFnReflector.getApplyMethod(scalarFn);
            javaScalarUdfs.put(createFunctionStmt.getNamePath(), UserFunctionDefinitions.JavaScalarFunction.create(method, jarPath));
            break;
        case USER_DEFINED_JAVA_AGGREGATE_FUNCTIONS:
            jarPath = getJarPath(createFunctionStmt);
            // Try loading the aggregate function just to make sure it exists. LazyAggregateCombineFn
            // will need to fetch it again at runtime.
            javaUdfLoader.loadAggregateFunction(createFunctionStmt.getNamePath(), jarPath);
            Combine.CombineFn<?, ?, ?> combineFn = new LazyAggregateCombineFn<>(createFunctionStmt.getNamePath(), jarPath);
            javaUdafs.put(createFunctionStmt.getNamePath(), combineFn);
            break;
        default:
            throw new IllegalArgumentException(String.format("Encountered unrecognized function group %s.", functionGroup));
    }
    zetaSqlCatalog.addFunction(new Function(createFunctionStmt.getNamePath(), functionGroup, createFunctionStmt.getIsAggregate() ? ZetaSQLFunctions.FunctionEnums.Mode.AGGREGATE : ZetaSQLFunctions.FunctionEnums.Mode.SCALAR, ImmutableList.of(createFunctionStmt.getSignature())));
}
Also used : TableValuedFunction(com.google.zetasql.TableValuedFunction) Function(com.google.zetasql.Function) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) Type(com.google.zetasql.Type) ZetaSQLType(com.google.zetasql.ZetaSQLType) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) ScalarFn(org.apache.beam.sdk.extensions.sql.udf.ScalarFn) Combine(org.apache.beam.sdk.transforms.Combine) FunctionArgumentType(com.google.zetasql.FunctionArgumentType) LazyAggregateCombineFn(org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn) Method(java.lang.reflect.Method)

Aggregations

Type (com.google.zetasql.Type)3 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)3 Function (com.google.zetasql.Function)1 FunctionArgumentType (com.google.zetasql.FunctionArgumentType)1 StructType (com.google.zetasql.StructType)1 TableValuedFunction (com.google.zetasql.TableValuedFunction)1 ZetaSQLType (com.google.zetasql.ZetaSQLType)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 Method (java.lang.reflect.Method)1 BigDecimal (java.math.BigDecimal)1 ArrayList (java.util.ArrayList)1 LazyAggregateCombineFn (org.apache.beam.sdk.extensions.sql.impl.LazyAggregateCombineFn)1 ScalarFn (org.apache.beam.sdk.extensions.sql.udf.ScalarFn)1 Combine (org.apache.beam.sdk.transforms.Combine)1 ByteString (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.avatica.util.ByteString)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