Search in sources :

Example 1 with RowExpressionOptimizer

use of io.prestosql.sql.relational.RowExpressionOptimizer in project hetu-core by openlookeng.

the class TranslateExpressions method createRewriter.

private static PlanRowExpressionRewriter createRewriter(Metadata metadata, SqlParser sqlParser) {
    return new PlanRowExpressionRewriter() {

        @Override
        public RowExpression rewrite(RowExpression expression, Rule.Context context) {
            // special treatment of the CallExpression in Aggregation
            if (expression instanceof CallExpression && ((CallExpression) expression).getArguments().stream().anyMatch(OriginalExpressionUtils::isExpression)) {
                return removeOriginalExpressionArguments((CallExpression) expression, context.getSession(), context.getSymbolAllocator(), context);
            }
            return removeOriginalExpression(expression, context, new HashMap<>());
        }

        private RowExpression removeOriginalExpressionArguments(CallExpression callExpression, Session session, PlanSymbolAllocator planSymbolAllocator, Rule.Context context) {
            Map<NodeRef<Expression>, Type> types = analyzeCallExpressionTypes(callExpression, session, planSymbolAllocator.getTypes());
            return new CallExpression(callExpression.getDisplayName(), callExpression.getFunctionHandle(), callExpression.getType(), callExpression.getArguments().stream().map(expression -> removeOriginalExpression(expression, session, types, context)).collect(toImmutableList()), Optional.empty());
        }

        private Map<NodeRef<Expression>, Type> analyzeCallExpressionTypes(CallExpression callExpression, Session session, TypeProvider typeProvider) {
            List<LambdaExpression> lambdaExpressions = callExpression.getArguments().stream().filter(OriginalExpressionUtils::isExpression).map(OriginalExpressionUtils::castToExpression).filter(LambdaExpression.class::isInstance).map(LambdaExpression.class::cast).collect(toImmutableList());
            ImmutableMap.Builder<NodeRef<Expression>, Type> builder = ImmutableMap.<NodeRef<Expression>, Type>builder();
            TypeAnalyzer typeAnalyzer = new TypeAnalyzer(sqlParser, metadata);
            if (!lambdaExpressions.isEmpty()) {
                List<FunctionType> functionTypes = metadata.getFunctionAndTypeManager().getFunctionMetadata(callExpression.getFunctionHandle()).getArgumentTypes().stream().filter(typeSignature -> typeSignature.getBase().equals(FunctionType.NAME)).map(metadata::getType).map(FunctionType.class::cast).collect(toImmutableList());
                InternalAggregationFunction internalAggregationFunction = metadata.getFunctionAndTypeManager().getAggregateFunctionImplementation(callExpression.getFunctionHandle());
                List<Class<?>> lambdaInterfaces = internalAggregationFunction.getLambdaInterfaces();
                verify(lambdaExpressions.size() == functionTypes.size());
                verify(lambdaExpressions.size() == lambdaInterfaces.size());
                for (int i = 0; i < lambdaExpressions.size(); i++) {
                    LambdaExpression lambdaExpression = lambdaExpressions.get(i);
                    FunctionType functionType = functionTypes.get(i);
                    // To compile lambda, LambdaDefinitionExpression needs to be generated from LambdaExpression,
                    // which requires the types of all sub-expressions.
                    // 
                    // In project and filter expression compilation, ExpressionAnalyzer.getExpressionTypesFromInput
                    // is used to generate the types of all sub-expressions. (see visitScanFilterAndProject and visitFilter)
                    // 
                    // This does not work here since the function call representation in final aggregation node
                    // is currently a hack: it takes intermediate type as input, and may not be a valid
                    // function call in Presto.
                    // 
                    // TODO: Once the final aggregation function call representation is fixed,
                    // the same mechanism in project and filter expression should be used here.
                    verify(lambdaExpression.getArguments().size() == functionType.getArgumentTypes().size());
                    Map<NodeRef<Expression>, Type> lambdaArgumentExpressionTypes = new HashMap<>();
                    Map<Symbol, Type> lambdaArgumentSymbolTypes = new HashMap<>();
                    for (int j = 0; j < lambdaExpression.getArguments().size(); j++) {
                        LambdaArgumentDeclaration argument = lambdaExpression.getArguments().get(j);
                        Type type = functionType.getArgumentTypes().get(j);
                        lambdaArgumentExpressionTypes.put(NodeRef.of(argument), type);
                        lambdaArgumentSymbolTypes.put(new Symbol(argument.getName().getValue()), type);
                    }
                    // the lambda expression itself
                    builder.put(NodeRef.of(lambdaExpression), functionType).putAll(lambdaArgumentExpressionTypes).putAll(typeAnalyzer.getTypes(session, TypeProvider.copyOf(lambdaArgumentSymbolTypes), lambdaExpression.getBody()));
                }
            }
            for (RowExpression argument : callExpression.getArguments()) {
                if (!isExpression(argument) || castToExpression(argument) instanceof LambdaExpression) {
                    continue;
                }
                builder.putAll(typeAnalyzer.getTypes(session, typeProvider, castToExpression(argument)));
            }
            return builder.build();
        }

        private RowExpression toRowExpression(Expression expression, Map<NodeRef<Expression>, Type> types, Map<Symbol, Integer> layout, Session session) {
            RowExpression rowExpression = SqlToRowExpressionTranslator.translate(expression, FunctionKind.SCALAR, types, layout, metadata.getFunctionAndTypeManager(), session, false);
            return new RowExpressionOptimizer(metadata).optimize(rowExpression, RowExpressionInterpreter.Level.SERIALIZABLE, session.toConnectorSession());
        }

        private RowExpression removeOriginalExpression(RowExpression expression, Rule.Context context, Map<Symbol, Integer> layout) {
            if (isExpression(expression)) {
                TypeAnalyzer typeAnalyzer = new TypeAnalyzer(sqlParser, metadata);
                return toRowExpression(castToExpression(expression), typeAnalyzer.getTypes(context.getSession(), context.getSymbolAllocator().getTypes(), castToExpression(expression)), layout, context.getSession());
            }
            return expression;
        }

        private RowExpression removeOriginalExpression(RowExpression rowExpression, Session session, Map<NodeRef<Expression>, Type> types, Rule.Context context) {
            if (isExpression(rowExpression)) {
                Expression expression = castToExpression(rowExpression);
                return toRowExpression(expression, types, new HashMap<>(), session);
            }
            return rowExpression;
        }
    };
}
Also used : FunctionKind(io.prestosql.spi.function.FunctionKind) TypeProvider(io.prestosql.sql.planner.TypeProvider) HashMap(java.util.HashMap) SqlParser(io.prestosql.sql.parser.SqlParser) TypeAnalyzer(io.prestosql.sql.planner.TypeAnalyzer) OriginalExpressionUtils.isExpression(io.prestosql.sql.relational.OriginalExpressionUtils.isExpression) InternalAggregationFunction(io.prestosql.operator.aggregation.InternalAggregationFunction) CallExpression(io.prestosql.spi.relation.CallExpression) Verify.verify(com.google.common.base.Verify.verify) Map(java.util.Map) Session(io.prestosql.Session) Type(io.prestosql.spi.type.Type) OriginalExpressionUtils.castToExpression(io.prestosql.sql.relational.OriginalExpressionUtils.castToExpression) SqlToRowExpressionTranslator(io.prestosql.sql.relational.SqlToRowExpressionTranslator) Symbol(io.prestosql.spi.plan.Symbol) ImmutableMap(com.google.common.collect.ImmutableMap) Rule(io.prestosql.sql.planner.iterative.Rule) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) FunctionType(io.prestosql.spi.type.FunctionType) RowExpressionOptimizer(io.prestosql.sql.relational.RowExpressionOptimizer) Metadata(io.prestosql.metadata.Metadata) LambdaArgumentDeclaration(io.prestosql.sql.tree.LambdaArgumentDeclaration) NodeRef(io.prestosql.sql.tree.NodeRef) PlanSymbolAllocator(io.prestosql.sql.planner.PlanSymbolAllocator) LambdaExpression(io.prestosql.sql.tree.LambdaExpression) RowExpressionInterpreter(io.prestosql.sql.planner.RowExpressionInterpreter) List(java.util.List) RowExpression(io.prestosql.spi.relation.RowExpression) Optional(java.util.Optional) Expression(io.prestosql.sql.tree.Expression) OriginalExpressionUtils(io.prestosql.sql.relational.OriginalExpressionUtils) HashMap(java.util.HashMap) Symbol(io.prestosql.spi.plan.Symbol) InternalAggregationFunction(io.prestosql.operator.aggregation.InternalAggregationFunction) PlanSymbolAllocator(io.prestosql.sql.planner.PlanSymbolAllocator) RowExpressionOptimizer(io.prestosql.sql.relational.RowExpressionOptimizer) NodeRef(io.prestosql.sql.tree.NodeRef) LambdaArgumentDeclaration(io.prestosql.sql.tree.LambdaArgumentDeclaration) CallExpression(io.prestosql.spi.relation.CallExpression) FunctionType(io.prestosql.spi.type.FunctionType) RowExpression(io.prestosql.spi.relation.RowExpression) TypeProvider(io.prestosql.sql.planner.TypeProvider) ImmutableMap(com.google.common.collect.ImmutableMap) TypeAnalyzer(io.prestosql.sql.planner.TypeAnalyzer) Type(io.prestosql.spi.type.Type) FunctionType(io.prestosql.spi.type.FunctionType) OriginalExpressionUtils.isExpression(io.prestosql.sql.relational.OriginalExpressionUtils.isExpression) CallExpression(io.prestosql.spi.relation.CallExpression) OriginalExpressionUtils.castToExpression(io.prestosql.sql.relational.OriginalExpressionUtils.castToExpression) LambdaExpression(io.prestosql.sql.tree.LambdaExpression) RowExpression(io.prestosql.spi.relation.RowExpression) Expression(io.prestosql.sql.tree.Expression) OriginalExpressionUtils(io.prestosql.sql.relational.OriginalExpressionUtils) LambdaExpression(io.prestosql.sql.tree.LambdaExpression) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) Session(io.prestosql.Session)

Example 2 with RowExpressionOptimizer

use of io.prestosql.sql.relational.RowExpressionOptimizer in project hetu-core by openlookeng.

the class TestingRowExpressionTranslator method translateAndOptimize.

public RowExpression translateAndOptimize(Expression expression, Map<NodeRef<Expression>, Type> types) {
    RowExpression rowExpression = SqlToRowExpressionTranslator.translate(expression, SCALAR, types, ImmutableMap.of(), metadata.getFunctionAndTypeManager(), TEST_SESSION, false);
    RowExpressionOptimizer optimizer = new RowExpressionOptimizer(metadata);
    return optimizer.optimize(rowExpression, OPTIMIZED, TEST_SESSION.toConnectorSession());
}
Also used : RowExpression(io.prestosql.spi.relation.RowExpression) RowExpressionOptimizer(io.prestosql.sql.relational.RowExpressionOptimizer)

Aggregations

RowExpression (io.prestosql.spi.relation.RowExpression)2 RowExpressionOptimizer (io.prestosql.sql.relational.RowExpressionOptimizer)2 Verify.verify (com.google.common.base.Verify.verify)1 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 Session (io.prestosql.Session)1 Metadata (io.prestosql.metadata.Metadata)1 InternalAggregationFunction (io.prestosql.operator.aggregation.InternalAggregationFunction)1 FunctionKind (io.prestosql.spi.function.FunctionKind)1 Symbol (io.prestosql.spi.plan.Symbol)1 CallExpression (io.prestosql.spi.relation.CallExpression)1 FunctionType (io.prestosql.spi.type.FunctionType)1 Type (io.prestosql.spi.type.Type)1 SqlParser (io.prestosql.sql.parser.SqlParser)1 PlanSymbolAllocator (io.prestosql.sql.planner.PlanSymbolAllocator)1 RowExpressionInterpreter (io.prestosql.sql.planner.RowExpressionInterpreter)1 TypeAnalyzer (io.prestosql.sql.planner.TypeAnalyzer)1 TypeProvider (io.prestosql.sql.planner.TypeProvider)1 Rule (io.prestosql.sql.planner.iterative.Rule)1 OriginalExpressionUtils (io.prestosql.sql.relational.OriginalExpressionUtils)1