Examples with Lookup com.facebook.presto.sql.planner.iterative.Lookup used on opensource projects

Example 1 with Lookup

use of com.facebook.presto.sql.planner.iterative.Lookup in project presto by prestodb.

the class SingleMarkDistinctToGroupBy method apply.

@Override
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator) {
    if (!(node instanceof AggregationNode)) {
        return Optional.empty();
    }
    AggregationNode parent = (AggregationNode) node;
    PlanNode source = lookup.resolve(parent.getSource());
    if (!(source instanceof MarkDistinctNode)) {
        return Optional.empty();
    }
    MarkDistinctNode child = (MarkDistinctNode) source;
    boolean hasFilters = parent.getAggregations().values().stream().map(FunctionCall::getFilter).anyMatch(Optional::isPresent);
    if (hasFilters) {
        return Optional.empty();
    }
    // optimize if and only if
    // all aggregation functions have a single common distinct mask symbol
    // AND all aggregation functions have mask
    Set<Symbol> masks = ImmutableSet.copyOf(parent.getMasks().values());
    if (masks.size() != 1 || parent.getMasks().size() != parent.getAggregations().size()) {
        return Optional.empty();
    }
    Symbol mask = Iterables.getOnlyElement(masks);
    if (!child.getMarkerSymbol().equals(mask)) {
        return Optional.empty();
    }
    return Optional.of(new AggregationNode(idAllocator.getNextId(), new AggregationNode(idAllocator.getNextId(), child.getSource(), Collections.emptyMap(), ImmutableList.of(child.getDistinctSymbols()), SINGLE, child.getHashSymbol(), Optional.empty()), // remove DISTINCT flag from function calls
    parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> removeDistinct(e.getValue()))), parent.getGroupingSets(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol()));
}

Example 2 with Lookup

use of com.facebook.presto.sql.planner.iterative.Lookup in project presto by prestodb.

the class RuleAssert method applyRule.

private RuleApplication applyRule() {
    PlanVariableAllocator variableAllocator = new PlanVariableAllocator(types.allVariables());
    Memo memo = new Memo(idAllocator, plan);
    Lookup lookup = Lookup.from(planNode -> Stream.of(memo.resolve(planNode)));
    PlanNode memoRoot = memo.getNode(memo.getRootGroup());
    return inTransaction(session -> applyRule(rule, memoRoot, ruleContext(statsCalculator, costCalculator, variableAllocator, memo, lookup, session)));
}

Example 3 with Lookup

use of com.facebook.presto.sql.planner.iterative.Lookup in project presto by prestodb.

the class PushAggregationThroughOuterJoin method createAggregationOverNull.

private Optional<MappedAggregationInfo> createAggregationOverNull(AggregationNode referenceAggregation, PlanVariableAllocator variableAllocator, PlanNodeIdAllocator idAllocator, Lookup lookup) {
    // Create a values node that consists of a single row of nulls.
    // Map the output symbols from the referenceAggregation's source
    // to symbol references for the new values node.
    ImmutableList.Builder<VariableReferenceExpression> nullVariables = ImmutableList.builder();
    ImmutableList.Builder<RowExpression> nullLiterals = ImmutableList.builder();
    ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> sourcesVariableMappingBuilder = ImmutableMap.builder();
    for (VariableReferenceExpression sourceVariable : referenceAggregation.getSource().getOutputVariables()) {
        RowExpression nullLiteral = constantNull(sourceVariable.getSourceLocation(), sourceVariable.getType());
        nullLiterals.add(nullLiteral);
        VariableReferenceExpression nullVariable = variableAllocator.newVariable(nullLiteral);
        nullVariables.add(nullVariable);
        // TODO The type should be from sourceVariable.getType
        sourcesVariableMappingBuilder.put(sourceVariable, nullVariable);
    }
    ValuesNode nullRow = new ValuesNode(referenceAggregation.getSourceLocation(), idAllocator.getNextId(), nullVariables.build(), ImmutableList.of(nullLiterals.build()));
    Map<VariableReferenceExpression, VariableReferenceExpression> sourcesVariableMapping = sourcesVariableMappingBuilder.build();
    // For each aggregation function in the reference node, create a corresponding aggregation function
    // that points to the nullRow. Map the symbols from the aggregations in referenceAggregation to the
    // symbols in these new aggregations.
    ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> aggregationsVariableMappingBuilder = ImmutableMap.builder();
    ImmutableMap.Builder<VariableReferenceExpression, AggregationNode.Aggregation> aggregationsOverNullBuilder = ImmutableMap.builder();
    for (Map.Entry<VariableReferenceExpression, AggregationNode.Aggregation> entry : referenceAggregation.getAggregations().entrySet()) {
        VariableReferenceExpression aggregationVariable = entry.getKey();
        AggregationNode.Aggregation aggregation = entry.getValue();
        if (!isUsingVariables(aggregation, sourcesVariableMapping.keySet())) {
            return Optional.empty();
        }
        AggregationNode.Aggregation overNullAggregation = new AggregationNode.Aggregation(new CallExpression(aggregation.getCall().getSourceLocation(), aggregation.getCall().getDisplayName(), aggregation.getCall().getFunctionHandle(), aggregation.getCall().getType(), aggregation.getArguments().stream().map(argument -> inlineVariables(sourcesVariableMapping, argument)).collect(toImmutableList())), aggregation.getFilter().map(filter -> inlineVariables(sourcesVariableMapping, filter)), aggregation.getOrderBy().map(orderBy -> inlineOrderByVariables(sourcesVariableMapping, orderBy)), aggregation.isDistinct(), aggregation.getMask().map(x -> new VariableReferenceExpression(sourcesVariableMapping.get(x).getSourceLocation(), sourcesVariableMapping.get(x).getName(), x.getType())));
        QualifiedObjectName functionName = functionAndTypeManager.getFunctionMetadata(overNullAggregation.getFunctionHandle()).getName();
        VariableReferenceExpression overNull = variableAllocator.newVariable(aggregation.getCall().getSourceLocation(), functionName.getObjectName(), aggregationVariable.getType());
        aggregationsOverNullBuilder.put(overNull, overNullAggregation);
        aggregationsVariableMappingBuilder.put(aggregationVariable, overNull);
    }
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationsSymbolMapping = aggregationsVariableMappingBuilder.build();
    // create an aggregation node whose source is the null row.
    AggregationNode aggregationOverNullRow = new AggregationNode(referenceAggregation.getSourceLocation(), idAllocator.getNextId(), nullRow, aggregationsOverNullBuilder.build(), globalAggregation(), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), Optional.empty());
    return Optional.of(new MappedAggregationInfo(aggregationOverNullRow, aggregationsSymbolMapping));
}

Example 4 with Lookup

use of com.facebook.presto.sql.planner.iterative.Lookup in project presto by prestodb.

the class AddIntermediateAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Lookup lookup = context.getLookup();
    PlanNodeIdAllocator idAllocator = context.getIdAllocator();
    Session session = context.getSession();
    TypeProvider types = context.getVariableAllocator().getTypes();
    Optional<PlanNode> rewrittenSource = recurseToPartial(lookup.resolve(aggregation.getSource()), lookup, idAllocator, types);
    if (!rewrittenSource.isPresent()) {
        return Result.empty();
    }
    PlanNode source = rewrittenSource.get();
    if (getTaskConcurrency(session) > 1) {
        Map<VariableReferenceExpression, Aggregation> variableToAggregations = inputsAsOutputs(aggregation.getAggregations(), types);
        if (variableToAggregations.isEmpty()) {
            return Result.empty();
        }
        source = roundRobinExchange(idAllocator.getNextId(), LOCAL, source);
        source = new AggregationNode(aggregation.getSourceLocation(), idAllocator.getNextId(), source, variableToAggregations, aggregation.getGroupingSets(), aggregation.getPreGroupedVariables(), INTERMEDIATE, aggregation.getHashVariable(), aggregation.getGroupIdVariable());
        source = gatheringExchange(idAllocator.getNextId(), LOCAL, source);
    }
    return Result.ofPlanNode(aggregation.replaceChildren(ImmutableList.of(source)));
}