Examples with ProjectNode com.facebook.presto.spi.plan.ProjectNode used on opensource projects

Example 31 with ProjectNode

use of com.facebook.presto.spi.plan.ProjectNode in project presto by prestodb.

the class RewriteAggregationIfToFilter method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    ProjectNode sourceProject = captures.get(CHILD);
    Set<Aggregation> aggregationsToRewrite = aggregationNode.getAggregations().values().stream().filter(aggregation -> shouldRewriteAggregation(aggregation, sourceProject)).collect(toImmutableSet());
    if (aggregationsToRewrite.isEmpty()) {
        return Result.empty();
    }
    context.getSession().getRuntimeStats().addMetricValue(REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED, 1);
    // Get the corresponding assignments in the input project.
    // The aggregationReferences only has the aggregations to rewrite, thus the sourceAssignments only has IF/CAST(IF) expressions with NULL false results.
    // Multiple aggregations may reference the same input. We use a map to dedup them based on the VariableReferenceExpression, so that we only do the rewrite once per input
    // IF expression.
    // The order of sourceAssignments determines the order of generating the new variables for the IF conditions and results. We use a sorted map to get a deterministic
    // order based on the name of the VariableReferenceExpressions.
    Map<VariableReferenceExpression, RowExpression> sourceAssignments = aggregationsToRewrite.stream().map(aggregation -> (VariableReferenceExpression) aggregation.getArguments().get(0)).collect(toImmutableSortedMap(VariableReferenceExpression::compareTo, identity(), variable -> sourceProject.getAssignments().get(variable), (left, right) -> left));
    Assignments.Builder newAssignments = Assignments.builder();
    newAssignments.putAll(sourceProject.getAssignments());
    // Map from the aggregation reference to the IF condition reference which will be put in the mask.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToConditionReference = new HashMap<>();
    // Map from the aggregation reference to the IF result reference. This only contains the aggregates where the IF can be safely unwrapped.
    // E.g., SUM(IF(CARDINALITY(array) > 0, array[1])) will not be included in this map as array[1] can return errors if we unwrap the IF.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToIfResultReference = new HashMap<>();
    AggregationIfToFilterRewriteStrategy rewriteStrategy = getAggregationIfToFilterRewriteStrategy(context.getSession());
    for (Map.Entry<VariableReferenceExpression, RowExpression> entry : sourceAssignments.entrySet()) {
        VariableReferenceExpression outputVariable = entry.getKey();
        RowExpression rowExpression = entry.getValue();
        SpecialFormExpression ifExpression = (SpecialFormExpression) ((rowExpression instanceof CallExpression) ? ((CallExpression) rowExpression).getArguments().get(0) : rowExpression);
        RowExpression condition = ifExpression.getArguments().get(0);
        VariableReferenceExpression conditionReference = context.getVariableAllocator().newVariable(condition);
        newAssignments.put(conditionReference, condition);
        aggregationReferenceToConditionReference.put(outputVariable, conditionReference);
        if (canUnwrapIf(ifExpression, rewriteStrategy)) {
            RowExpression trueResult = ifExpression.getArguments().get(1);
            if (rowExpression instanceof CallExpression) {
                // Wrap the result with CAST().
                trueResult = new CallExpression(((CallExpression) rowExpression).getDisplayName(), ((CallExpression) rowExpression).getFunctionHandle(), rowExpression.getType(), ImmutableList.of(trueResult));
            }
            VariableReferenceExpression ifResultReference = context.getVariableAllocator().newVariable(trueResult);
            newAssignments.put(ifResultReference, trueResult);
            aggregationReferenceToIfResultReference.put(outputVariable, ifResultReference);
        }
    }
    // Build new aggregations.
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    // Stores the masks used to build the filter predicates. Use set to dedup the predicates.
    ImmutableSortedSet.Builder<VariableReferenceExpression> masks = ImmutableSortedSet.naturalOrder();
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        VariableReferenceExpression output = entry.getKey();
        Aggregation aggregation = entry.getValue();
        if (!aggregationsToRewrite.contains(aggregation)) {
            aggregations.put(output, aggregation);
            continue;
        }
        VariableReferenceExpression aggregationReference = (VariableReferenceExpression) aggregation.getArguments().get(0);
        CallExpression callExpression = aggregation.getCall();
        VariableReferenceExpression ifResultReference = aggregationReferenceToIfResultReference.get(aggregationReference);
        if (ifResultReference != null) {
            callExpression = new CallExpression(callExpression.getSourceLocation(), callExpression.getDisplayName(), callExpression.getFunctionHandle(), callExpression.getType(), ImmutableList.of(ifResultReference));
        }
        VariableReferenceExpression mask = aggregationReferenceToConditionReference.get(aggregationReference);
        aggregations.put(output, new Aggregation(callExpression, Optional.empty(), aggregation.getOrderBy(), aggregation.isDistinct(), Optional.of(aggregationReferenceToConditionReference.get(aggregationReference))));
        masks.add(mask);
    }
    RowExpression predicate = TRUE_CONSTANT;
    if (!aggregationNode.hasNonEmptyGroupingSet() && aggregationsToRewrite.size() == aggregationNode.getAggregations().size()) {
        // All aggregations are rewritten by this rule. We can add a filter with all the masks to make the query more efficient.
        predicate = or(masks.build());
    }
    return Result.ofPlanNode(new AggregationNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new FilterNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), sourceProject.getSource(), newAssignments.build()), predicate), aggregations.build(), aggregationNode.getGroupingSets(), aggregationNode.getPreGroupedVariables(), aggregationNode.getStep(), aggregationNode.getHashVariable(), aggregationNode.getGroupIdVariable()));
}

Example 32 with ProjectNode

use of com.facebook.presto.spi.plan.ProjectNode in project presto by prestodb.

the class ImplementFilteredAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Assignments.Builder newAssignments = Assignments.builder();
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    ImmutableList.Builder<Expression> maskSymbols = ImmutableList.builder();
    boolean aggregateWithoutFilterPresent = false;
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregation.getAggregations().entrySet()) {
        VariableReferenceExpression output = entry.getKey();
        // strip the filters
        Optional<VariableReferenceExpression> mask = entry.getValue().getMask();
        if (entry.getValue().getFilter().isPresent()) {
            // TODO remove cast once assignment can be RowExpression
            Expression filter = OriginalExpressionUtils.castToExpression(entry.getValue().getFilter().get());
            VariableReferenceExpression variable = context.getVariableAllocator().newVariable(filter, BOOLEAN);
            verify(!mask.isPresent(), "Expected aggregation without mask symbols, see Rule pattern");
            newAssignments.put(variable, castToRowExpression(filter));
            mask = Optional.of(variable);
            maskSymbols.add(createSymbolReference(variable));
        } else {
            aggregateWithoutFilterPresent = true;
        }
        aggregations.put(output, new Aggregation(entry.getValue().getCall(), Optional.empty(), entry.getValue().getOrderBy(), entry.getValue().isDistinct(), mask));
    }
    Expression predicate = TRUE_LITERAL;
    if (!aggregation.hasNonEmptyGroupingSet() && !aggregateWithoutFilterPresent) {
        predicate = combineDisjunctsWithDefault(maskSymbols.build(), TRUE_LITERAL);
    }
    // identity projection for all existing inputs
    newAssignments.putAll(identitiesAsSymbolReferences(aggregation.getSource().getOutputVariables()));
    return Result.ofPlanNode(new AggregationNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), new FilterNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), aggregation.getSource(), newAssignments.build()), castToRowExpression(predicate)), aggregations.build(), aggregation.getGroupingSets(), ImmutableList.of(), aggregation.getStep(), aggregation.getHashVariable(), aggregation.getGroupIdVariable()));
}

Example 33 with ProjectNode

use of com.facebook.presto.spi.plan.ProjectNode in project presto by prestodb.

the class InlineProjections method apply.

@Override
public Result apply(ProjectNode parent, Captures captures, Context context) {
    ProjectNode child = captures.get(CHILD);
    // Do not inline remote projections, or if parent and child has different locality
    if (parent.getLocality().equals(REMOTE) || child.getLocality().equals(REMOTE) || !parent.getLocality().equals(child.getLocality())) {
        return Result.empty();
    }
    Sets.SetView<VariableReferenceExpression> targets = extractInliningTargets(parent, child, context);
    if (targets.isEmpty()) {
        return Result.empty();
    }
    // inline the expressions
    Assignments assignments = child.getAssignments().filter(targets::contains);
    Map<VariableReferenceExpression, RowExpression> parentAssignments = parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> inlineReferences(entry.getValue(), assignments, context.getVariableAllocator().getTypes())));
    // Synthesize identity assignments for the inputs of expressions that were inlined
    // to place in the child projection.
    // If all assignments end up becoming identity assignments, they'll get pruned by
    // other rules
    Set<VariableReferenceExpression> inputs = child.getAssignments().entrySet().stream().filter(entry -> targets.contains(entry.getKey())).map(Map.Entry::getValue).flatMap(expression -> extractInputs(expression, context.getVariableAllocator().getTypes()).stream()).collect(toSet());
    Builder childAssignments = Assignments.builder();
    for (Map.Entry<VariableReferenceExpression, RowExpression> assignment : child.getAssignments().entrySet()) {
        if (!targets.contains(assignment.getKey())) {
            childAssignments.put(assignment);
        }
    }
    boolean allTranslated = child.getAssignments().entrySet().stream().map(Map.Entry::getValue).noneMatch(OriginalExpressionUtils::isExpression);
    for (VariableReferenceExpression input : inputs) {
        if (allTranslated) {
            childAssignments.put(input, input);
        } else {
            childAssignments.put(identityAsSymbolReference(input));
        }
    }
    return Result.ofPlanNode(new ProjectNode(parent.getSourceLocation(), parent.getId(), new ProjectNode(parent.getSourceLocation(), child.getId(), child.getSource(), childAssignments.build(), child.getLocality()), Assignments.copyOf(parentAssignments), parent.getLocality()));
}

Example 34 with ProjectNode

use of com.facebook.presto.spi.plan.ProjectNode in project presto by prestodb.

the class InlineProjections method extractInliningTargets.

private Sets.SetView<VariableReferenceExpression> extractInliningTargets(ProjectNode parent, ProjectNode child, Context context) {
    // candidates for inlining are
    // 1. references to simple constants
    // 2. references to complex expressions that
    // a. are not inputs to try() expressions
    // b. appear only once across all expressions
    // c. are not identity projections
    // which come from the child, as opposed to an enclosing scope.
    Set<VariableReferenceExpression> childOutputSet = ImmutableSet.copyOf(child.getOutputVariables());
    TypeProvider types = context.getVariableAllocator().getTypes();
    Map<VariableReferenceExpression, Long> dependencies = parent.getAssignments().getExpressions().stream().flatMap(expression -> extractInputs(expression, context.getVariableAllocator().getTypes()).stream()).filter(childOutputSet::contains).collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
    // find references to simple constants
    Set<VariableReferenceExpression> constants = dependencies.keySet().stream().filter(input -> isConstant(child.getAssignments().get(input))).collect(toSet());
    // exclude any complex inputs to TRY expressions. Inlining them would potentially
    // change the semantics of those expressions
    Set<VariableReferenceExpression> tryArguments = parent.getAssignments().getExpressions().stream().flatMap(expression -> extractTryArguments(expression, types).stream()).collect(toSet());
    Set<VariableReferenceExpression> singletons = dependencies.entrySet().stream().filter(// reference appears just once across all expressions in parent project node
    entry -> entry.getValue() == 1).filter(// they are not inputs to TRY. Otherwise, inlining might change semantics
    entry -> !tryArguments.contains(entry.getKey())).filter(// skip identities, otherwise, this rule will keep firing forever
    entry -> !isIdentity(child.getAssignments(), entry.getKey())).map(Map.Entry::getKey).collect(toSet());
    return Sets.union(singletons, constants);
}

Example 35 with ProjectNode

use of com.facebook.presto.spi.plan.ProjectNode in project presto by prestodb.

the class PushProjectionThroughUnion method apply.

@Override
public Result apply(ProjectNode parent, Captures captures, Context context) {
    UnionNode source = captures.get(CHILD);
    // OutputLayout of the resultant Union, will be same as the layout of the Project
    List<VariableReferenceExpression> outputLayout = parent.getOutputVariables();
    // Mapping from the output symbol to ordered list of symbols from each of the sources
    ImmutableListMultimap.Builder<VariableReferenceExpression, VariableReferenceExpression> mappings = ImmutableListMultimap.builder();
    // sources for the resultant UnionNode
    ImmutableList.Builder<PlanNode> outputSources = ImmutableList.builder();
    for (int i = 0; i < source.getSources().size(); i++) {
        // assignments for the new ProjectNode
        Assignments.Builder assignments = Assignments.builder();
        // mapping from current ProjectNode to new ProjectNode, used to identify the output layout
        Map<VariableReferenceExpression, VariableReferenceExpression> projectVariableMapping = new HashMap<>();
        // Translate the assignments in the ProjectNode using symbols of the source of the UnionNode
        for (Map.Entry<VariableReferenceExpression, RowExpression> entry : parent.getAssignments().entrySet()) {
            RowExpression translatedExpression;
            VariableReferenceExpression variable;
            translatedExpression = RowExpressionVariableInliner.inlineVariables(source.sourceVariableMap(i), entry.getValue());
            variable = context.getVariableAllocator().newVariable(translatedExpression);
            assignments.put(variable, translatedExpression);
            projectVariableMapping.put(entry.getKey(), variable);
        }
        outputSources.add(new ProjectNode(source.getSourceLocation(), context.getIdAllocator().getNextId(), source.getSources().get(i), assignments.build(), parent.getLocality()));
        outputLayout.forEach(variable -> mappings.put(variable, projectVariableMapping.get(variable)));
    }
    ListMultimap<VariableReferenceExpression, VariableReferenceExpression> outputsToInputs = mappings.build();
    return Result.ofPlanNode(new UnionNode(source.getSourceLocation(), parent.getId(), outputSources.build(), ImmutableList.copyOf(outputsToInputs.keySet()), fromListMultimap(outputsToInputs)));
}