Search in sources :

Example 1 with ExpressionRewriter

use of com.facebook.presto.sql.tree.ExpressionRewriter in project presto by prestodb.

the class CreateFunctionTask method createSqlInvokedFunction.

private SqlInvokedFunction createSqlInvokedFunction(CreateFunction statement, Metadata metadata, Analysis analysis) {
    QualifiedObjectName functionName = statement.isTemporary() ? QualifiedObjectName.valueOf(SESSION_NAMESPACE, statement.getFunctionName().getSuffix()) : qualifyObjectName(statement.getFunctionName());
    List<Parameter> parameters = statement.getParameters().stream().map(parameter -> new Parameter(parameter.getName().toString(), parseTypeSignature(parameter.getType()))).collect(toImmutableList());
    TypeSignature returnType = parseTypeSignature(statement.getReturnType());
    String description = statement.getComment().orElse("");
    RoutineCharacteristics routineCharacteristics = RoutineCharacteristics.builder().setLanguage(new RoutineCharacteristics.Language(statement.getCharacteristics().getLanguage().getLanguage())).setDeterminism(RoutineCharacteristics.Determinism.valueOf(statement.getCharacteristics().getDeterminism().name())).setNullCallClause(RoutineCharacteristics.NullCallClause.valueOf(statement.getCharacteristics().getNullCallClause().name())).build();
    RoutineBody body = statement.getBody();
    if (statement.getBody() instanceof Return) {
        Expression bodyExpression = ((Return) statement.getBody()).getExpression();
        Type bodyType = analysis.getType(bodyExpression);
        // Coerce expressions in body if necessary
        bodyExpression = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter<Void>() {

            @Override
            public Expression rewriteExpression(Expression expression, Void context, ExpressionTreeRewriter<Void> treeRewriter) {
                Expression rewritten = treeRewriter.defaultRewrite(expression, null);
                Type coercion = analysis.getCoercion(expression);
                if (coercion != null) {
                    return new Cast(rewritten, coercion.getTypeSignature().toString(), false, analysis.isTypeOnlyCoercion(expression));
                }
                return rewritten;
            }
        }, bodyExpression, null);
        if (!bodyType.equals(metadata.getType(returnType))) {
            // Casting is safe here, since we have verified at analysis time that the actual type of the body is coercible to declared return type.
            bodyExpression = new Cast(bodyExpression, statement.getReturnType());
        }
        body = new Return(bodyExpression);
    }
    return new SqlInvokedFunction(functionName, parameters, returnType, description, routineCharacteristics, formatSql(body, Optional.empty()), notVersioned());
}
Also used : ExpressionTreeRewriter(com.facebook.presto.sql.tree.ExpressionTreeRewriter) WarningCollector(com.facebook.presto.spi.WarningCollector) Parameter(com.facebook.presto.spi.function.Parameter) Analyzer(com.facebook.presto.sql.analyzer.Analyzer) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) SESSION_NAMESPACE(com.facebook.presto.metadata.SessionFunctionHandle.SESSION_NAMESPACE) SqlInvokedFunction(com.facebook.presto.spi.function.SqlInvokedFunction) PrestoException(com.facebook.presto.spi.PrestoException) CreateFunction(com.facebook.presto.sql.tree.CreateFunction) TypeSignature(com.facebook.presto.common.type.TypeSignature) Inject(javax.inject.Inject) ALREADY_EXISTS(com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) QualifiedObjectName(com.facebook.presto.common.QualifiedObjectName) Cast(com.facebook.presto.sql.tree.Cast) SqlFunctionHandle(com.facebook.presto.spi.function.SqlFunctionHandle) TransactionManager(com.facebook.presto.transaction.TransactionManager) FunctionAndTypeManager.qualifyObjectName(com.facebook.presto.metadata.FunctionAndTypeManager.qualifyObjectName) Type(com.facebook.presto.common.type.Type) RoutineCharacteristics(com.facebook.presto.spi.function.RoutineCharacteristics) Futures.immediateFuture(com.google.common.util.concurrent.Futures.immediateFuture) Session(com.facebook.presto.Session) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) RoutineBody(com.facebook.presto.sql.tree.RoutineBody) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) String.format(java.lang.String.format) SqlParser(com.facebook.presto.sql.parser.SqlParser) ExpressionRewriter(com.facebook.presto.sql.tree.ExpressionRewriter) List(java.util.List) SqlFunctionId(com.facebook.presto.spi.function.SqlFunctionId) TypeSignature.parseTypeSignature(com.facebook.presto.common.type.TypeSignature.parseTypeSignature) Expression(com.facebook.presto.sql.tree.Expression) NOT_SUPPORTED(com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED) Return(com.facebook.presto.sql.tree.Return) FunctionVersion.notVersioned(com.facebook.presto.spi.function.FunctionVersion.notVersioned) SqlFormatter.formatSql(com.facebook.presto.sql.SqlFormatter.formatSql) Analysis(com.facebook.presto.sql.analyzer.Analysis) Optional(java.util.Optional) VisibleForTesting(com.google.common.annotations.VisibleForTesting) Metadata(com.facebook.presto.metadata.Metadata) AccessControl(com.facebook.presto.security.AccessControl) Cast(com.facebook.presto.sql.tree.Cast) ExpressionTreeRewriter(com.facebook.presto.sql.tree.ExpressionTreeRewriter) Return(com.facebook.presto.sql.tree.Return) SqlInvokedFunction(com.facebook.presto.spi.function.SqlInvokedFunction) QualifiedObjectName(com.facebook.presto.common.QualifiedObjectName) TypeSignature(com.facebook.presto.common.type.TypeSignature) TypeSignature.parseTypeSignature(com.facebook.presto.common.type.TypeSignature.parseTypeSignature) Type(com.facebook.presto.common.type.Type) ExpressionRewriter(com.facebook.presto.sql.tree.ExpressionRewriter) RoutineCharacteristics(com.facebook.presto.spi.function.RoutineCharacteristics) Expression(com.facebook.presto.sql.tree.Expression) Parameter(com.facebook.presto.spi.function.Parameter) RoutineBody(com.facebook.presto.sql.tree.RoutineBody)

Example 2 with ExpressionRewriter

use of com.facebook.presto.sql.tree.ExpressionRewriter in project presto by prestodb.

the class SqlFunctionUtils method rewriteLambdaExpression.

private static Expression rewriteLambdaExpression(Expression sqlFunction, Map<String, VariableReferenceExpression> arguments, ExpressionAnalysis functionAnalysis, PlanVariableAllocator variableAllocator) {
    Map<NodeRef<Identifier>, LambdaArgumentDeclaration> lambdaArgumentReferences = functionAnalysis.getLambdaArgumentReferences();
    Map<NodeRef<Expression>, Type> expressionTypes = functionAnalysis.getExpressionTypes();
    // Rewrite reference to LambdaArgumentDeclaration
    Map<NodeRef<LambdaArgumentDeclaration>, VariableReferenceExpression> variables = expressionTypes.entrySet().stream().filter(entry -> entry.getKey().getNode() instanceof LambdaArgumentDeclaration).distinct().collect(toImmutableMap(entry -> NodeRef.of((LambdaArgumentDeclaration) entry.getKey().getNode()), entry -> variableAllocator.newVariable(((LambdaArgumentDeclaration) entry.getKey().getNode()).getName(), entry.getValue(), "lambda")));
    Expression rewritten = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter<Map<NodeRef<Identifier>, LambdaArgumentDeclaration>>() {

        @Override
        public Expression rewriteLambdaExpression(LambdaExpression node, Map<NodeRef<Identifier>, LambdaArgumentDeclaration> context, ExpressionTreeRewriter<Map<NodeRef<Identifier>, LambdaArgumentDeclaration>> treeRewriter) {
            return new LambdaExpression(node.getArguments().stream().map(argument -> new LambdaArgumentDeclaration(new Identifier(variables.get(NodeRef.of(argument)).getName()))).collect(toImmutableList()), treeRewriter.rewrite(node.getBody(), context));
        }

        @Override
        public Expression rewriteIdentifier(Identifier node, Map<NodeRef<Identifier>, LambdaArgumentDeclaration> context, ExpressionTreeRewriter<Map<NodeRef<Identifier>, LambdaArgumentDeclaration>> treeRewriter) {
            NodeRef<Identifier> ref = NodeRef.of(node);
            if (context.containsKey(ref)) {
                return createSymbolReference(variables.get(NodeRef.of(context.get(ref))));
            }
            return node;
        }
    }, sqlFunction, lambdaArgumentReferences);
    // Rewrite function input referenced in lambda
    rewritten = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter<Map<String, VariableReferenceExpression>>() {

        @Override
        public Expression rewriteIdentifier(Identifier node, Map<String, VariableReferenceExpression> context, ExpressionTreeRewriter<Map<String, VariableReferenceExpression>> treeRewriter) {
            if (context.containsKey(node.getValue())) {
                return createSymbolReference(context.get(node.getValue()));
            }
            return node;
        }
    }, rewritten, arguments);
    // Desugar lambda capture
    return LambdaCaptureDesugaringRewriter.rewrite(rewritten, variableAllocator);
}
Also used : ExpressionTreeRewriter(com.facebook.presto.sql.tree.ExpressionTreeRewriter) LambdaArgumentDeclaration(com.facebook.presto.sql.tree.LambdaArgumentDeclaration) RowExpressionRewriter(com.facebook.presto.expressions.RowExpressionRewriter) VariableReferenceExpression(com.facebook.presto.spi.relation.VariableReferenceExpression) ExpressionAnalyzer.analyzeSqlFunctionExpression(com.facebook.presto.sql.analyzer.ExpressionAnalyzer.analyzeSqlFunctionExpression) SqlInvokedFunction(com.facebook.presto.spi.function.SqlInvokedFunction) Identifier(com.facebook.presto.sql.tree.Identifier) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) Map(java.util.Map) LambdaCaptureDesugaringRewriter(com.facebook.presto.sql.planner.iterative.rule.LambdaCaptureDesugaringRewriter) Cast(com.facebook.presto.sql.tree.Cast) SQL(com.facebook.presto.spi.function.FunctionImplementationType.SQL) ENGLISH(java.util.Locale.ENGLISH) Type(com.facebook.presto.common.type.Type) RowExpression(com.facebook.presto.spi.relation.RowExpression) SymbolReference(com.facebook.presto.sql.tree.SymbolReference) RowExpressionTreeRewriter(com.facebook.presto.expressions.RowExpressionTreeRewriter) ImmutableMap(com.google.common.collect.ImmutableMap) LambdaExpression(com.facebook.presto.sql.tree.LambdaExpression) SqlFunctionProperties(com.facebook.presto.common.function.SqlFunctionProperties) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) String.format(java.lang.String.format) SqlParser(com.facebook.presto.sql.parser.SqlParser) ExpressionRewriter(com.facebook.presto.sql.tree.ExpressionRewriter) Preconditions.checkState(com.google.common.base.Preconditions.checkState) NodeRef(com.facebook.presto.sql.tree.NodeRef) FunctionMetadata(com.facebook.presto.spi.function.FunctionMetadata) List(java.util.List) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) SqlFunctionId(com.facebook.presto.spi.function.SqlFunctionId) Expression(com.facebook.presto.sql.tree.Expression) AS_DOUBLE(com.facebook.presto.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DOUBLE) AS_DECIMAL(com.facebook.presto.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DECIMAL) PlanVariableAllocator(com.facebook.presto.sql.planner.PlanVariableAllocator) Function.identity(java.util.function.Function.identity) Optional(java.util.Optional) ParsingOptions(com.facebook.presto.sql.parser.ParsingOptions) ExpressionAnalysis(com.facebook.presto.sql.analyzer.ExpressionAnalysis) Metadata(com.facebook.presto.metadata.Metadata) SqlInvokedScalarFunctionImplementation(com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation) ExpressionTreeUtils.createSymbolReference(com.facebook.presto.sql.analyzer.ExpressionTreeUtils.createSymbolReference) ExpressionTreeRewriter(com.facebook.presto.sql.tree.ExpressionTreeRewriter) RowExpressionTreeRewriter(com.facebook.presto.expressions.RowExpressionTreeRewriter) NodeRef(com.facebook.presto.sql.tree.NodeRef) Type(com.facebook.presto.common.type.Type) LambdaArgumentDeclaration(com.facebook.presto.sql.tree.LambdaArgumentDeclaration) Identifier(com.facebook.presto.sql.tree.Identifier) RowExpressionRewriter(com.facebook.presto.expressions.RowExpressionRewriter) ExpressionRewriter(com.facebook.presto.sql.tree.ExpressionRewriter) VariableReferenceExpression(com.facebook.presto.spi.relation.VariableReferenceExpression) ExpressionAnalyzer.analyzeSqlFunctionExpression(com.facebook.presto.sql.analyzer.ExpressionAnalyzer.analyzeSqlFunctionExpression) RowExpression(com.facebook.presto.spi.relation.RowExpression) LambdaExpression(com.facebook.presto.sql.tree.LambdaExpression) Expression(com.facebook.presto.sql.tree.Expression) VariableReferenceExpression(com.facebook.presto.spi.relation.VariableReferenceExpression) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) LambdaExpression(com.facebook.presto.sql.tree.LambdaExpression)

Aggregations

Type (com.facebook.presto.common.type.Type)2 Metadata (com.facebook.presto.metadata.Metadata)2 SqlFunctionId (com.facebook.presto.spi.function.SqlFunctionId)2 SqlInvokedFunction (com.facebook.presto.spi.function.SqlInvokedFunction)2 SqlParser (com.facebook.presto.sql.parser.SqlParser)2 Cast (com.facebook.presto.sql.tree.Cast)2 Expression (com.facebook.presto.sql.tree.Expression)2 ExpressionRewriter (com.facebook.presto.sql.tree.ExpressionRewriter)2 ExpressionTreeRewriter (com.facebook.presto.sql.tree.ExpressionTreeRewriter)2 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)2 String.format (java.lang.String.format)2 List (java.util.List)2 Map (java.util.Map)2 Optional (java.util.Optional)2 Session (com.facebook.presto.Session)1 QualifiedObjectName (com.facebook.presto.common.QualifiedObjectName)1 SqlFunctionProperties (com.facebook.presto.common.function.SqlFunctionProperties)1 TypeSignature (com.facebook.presto.common.type.TypeSignature)1 TypeSignature.parseTypeSignature (com.facebook.presto.common.type.TypeSignature.parseTypeSignature)1 RowExpressionRewriter (com.facebook.presto.expressions.RowExpressionRewriter)1