Examples with SymbolMapper com.facebook.presto.sql.planner.optimizations.SymbolMapper used on opensource projects

Example 1 with SymbolMapper

use of com.facebook.presto.sql.planner.optimizations.SymbolMapper in project presto by prestodb.

the class PushPartialAggregationThroughExchange method pushPartial.

private PlanNode pushPartial(AggregationNode aggregation, ExchangeNode exchange, Context context) {
    List<PlanNode> partials = new ArrayList<>();
    for (int i = 0; i < exchange.getSources().size(); i++) {
        PlanNode source = exchange.getSources().get(i);
        SymbolMapper.Builder mappingsBuilder = SymbolMapper.builder();
        for (int outputIndex = 0; outputIndex < exchange.getOutputVariables().size(); outputIndex++) {
            VariableReferenceExpression output = exchange.getOutputVariables().get(outputIndex);
            VariableReferenceExpression input = exchange.getInputs().get(i).get(outputIndex);
            if (!output.equals(input)) {
                mappingsBuilder.put(output, input);
            }
        }
        SymbolMapper symbolMapper = mappingsBuilder.build();
        AggregationNode mappedPartial = symbolMapper.map(aggregation, source, context.getIdAllocator());
        Assignments.Builder assignments = Assignments.builder();
        for (VariableReferenceExpression output : aggregation.getOutputVariables()) {
            VariableReferenceExpression input = symbolMapper.map(output);
            assignments.put(output, input);
        }
        partials.add(new ProjectNode(exchange.getSourceLocation(), context.getIdAllocator().getNextId(), mappedPartial, assignments.build(), LOCAL));
    }
    for (PlanNode node : partials) {
        verify(aggregation.getOutputVariables().equals(node.getOutputVariables()));
    }
    // Since this exchange source is now guaranteed to have the same symbols as the inputs to the the partial
    // aggregation, we don't need to rewrite symbols in the partitioning function
    List<VariableReferenceExpression> aggregationOutputs = aggregation.getOutputVariables();
    PartitioningScheme partitioning = new PartitioningScheme(exchange.getPartitioningScheme().getPartitioning(), aggregationOutputs, exchange.getPartitioningScheme().getHashColumn(), exchange.getPartitioningScheme().isReplicateNullsAndAny(), exchange.getPartitioningScheme().getBucketToPartition());
    return new ExchangeNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), exchange.getType(), exchange.getScope(), partitioning, partials, ImmutableList.copyOf(Collections.nCopies(partials.size(), aggregationOutputs)), exchange.isEnsureSourceOrdering(), Optional.empty());
}

Example 2 with SymbolMapper

use of com.facebook.presto.sql.planner.optimizations.SymbolMapper in project presto by prestodb.

the class PushTopNThroughUnion method apply.

@Override
public Result apply(TopNNode topNNode, Captures captures, Context context) {
    UnionNode unionNode = captures.get(CHILD);
    ImmutableList.Builder<PlanNode> sources = ImmutableList.builder();
    for (PlanNode source : unionNode.getSources()) {
        SymbolMapper.Builder symbolMapper = SymbolMapper.builder();
        Set<VariableReferenceExpression> sourceOutputVariables = ImmutableSet.copyOf(source.getOutputVariables());
        for (VariableReferenceExpression unionOutput : unionNode.getOutputVariables()) {
            Set<VariableReferenceExpression> inputVariables = ImmutableSet.copyOf(unionNode.getVariableMapping().get(unionOutput));
            VariableReferenceExpression unionInput = getLast(intersection(inputVariables, sourceOutputVariables));
            symbolMapper.put(unionOutput, unionInput);
        }
        sources.add(symbolMapper.build().map(topNNode, source, context.getIdAllocator().getNextId()));
    }
    return Result.ofPlanNode(new UnionNode(unionNode.getSourceLocation(), unionNode.getId(), sources.build(), unionNode.getOutputVariables(), unionNode.getVariableMapping()));
}

Example 3 with SymbolMapper

use of com.facebook.presto.sql.planner.optimizations.SymbolMapper in project presto by prestodb.

the class PlanRemotePojections method dedupVariables.

private static List<ProjectionContext> dedupVariables(List<ProjectionContext> projectionContexts) {
    ImmutableList.Builder<ProjectionContext> deduppedProjectionContexts = ImmutableList.builder();
    Set<VariableReferenceExpression> originalVariable = projectionContexts.get(projectionContexts.size() - 1).getProjections().keySet();
    SymbolMapper mapper = null;
    for (int i = 0; i < projectionContexts.size(); i++) {
        Map<VariableReferenceExpression, RowExpression> projections = projectionContexts.get(i).getProjections();
        // Apply mapping from previous projection
        if (mapper != null) {
            ImmutableMap.Builder<VariableReferenceExpression, RowExpression> newProjections = ImmutableMap.builder();
            for (Map.Entry<VariableReferenceExpression, RowExpression> entry : projections.entrySet()) {
                newProjections.put(entry.getKey(), mapper.map(entry.getValue()));
            }
            projections = newProjections.build();
        }
        // Dedup
        ImmutableMultimap.Builder<RowExpression, VariableReferenceExpression> reverseProjectionsBuilder = ImmutableMultimap.builder();
        projections.forEach((key, value) -> reverseProjectionsBuilder.put(value, key));
        ImmutableMultimap<RowExpression, VariableReferenceExpression> reverseProjections = reverseProjectionsBuilder.build();
        if (reverseProjections.keySet().size() == projectionContexts.get(i).getProjections().size() && reverseProjections.keySet().stream().noneMatch(VariableReferenceExpression.class::isInstance)) {
            // No duplication
            deduppedProjectionContexts.add(new ProjectionContext(projections, projectionContexts.get(i).isRemote()));
            mapper = null;
        } else {
            SymbolMapper.Builder mapperBuilder = SymbolMapper.builder();
            ImmutableMap.Builder<VariableReferenceExpression, RowExpression> dedupedProjectionsBuilder = ImmutableMap.builder();
            for (RowExpression key : reverseProjections.keySet()) {
                List<VariableReferenceExpression> values = ImmutableList.copyOf(reverseProjections.get(key));
                if (key instanceof VariableReferenceExpression) {
                    values.forEach(variable -> mapperBuilder.put(variable, (VariableReferenceExpression) key));
                    dedupedProjectionsBuilder.put((VariableReferenceExpression) key, key);
                } else if (values.size() > 1) {
                    // Consolidate to one variable, prefer variables from original plan
                    List<VariableReferenceExpression> fromOriginal = originalVariable.stream().filter(values::contains).collect(toImmutableList());
                    VariableReferenceExpression variable = fromOriginal.isEmpty() ? values.get(0) : getOnlyElement(fromOriginal);
                    for (int j = 0; j < values.size(); j++) {
                        if (!values.get(j).equals(variable)) {
                            mapperBuilder.put(values.get(j), variable);
                        }
                    }
                    dedupedProjectionsBuilder.put(variable, key);
                } else {
                    checkState(values.size() == 1, "Expect only 1 value");
                    dedupedProjectionsBuilder.put(values.get(0), key);
                }
            }
            deduppedProjectionContexts.add(new ProjectionContext(dedupedProjectionsBuilder.build(), projectionContexts.get(i).isRemote()));
            mapper = mapperBuilder.build();
        }
    }
    return deduppedProjectionContexts.build();
}

Example 4 with SymbolMapper

use of com.facebook.presto.sql.planner.optimizations.SymbolMapper in project presto by prestodb.

the class PushTableWriteThroughUnion method rewriteSource.

private static TableWriterNode rewriteSource(TableWriterNode writerNode, UnionNode unionNode, int source, List<Map<VariableReferenceExpression, VariableReferenceExpression>> sourceMappings, Context context) {
    Map<VariableReferenceExpression, VariableReferenceExpression> inputMappings = getInputVariableMapping(unionNode, source);
    ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> mappings = ImmutableMap.builder();
    mappings.putAll(inputMappings);
    ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> outputMappings = ImmutableMap.builder();
    for (VariableReferenceExpression outputVariable : writerNode.getOutputVariables()) {
        if (inputMappings.containsKey(outputVariable)) {
            outputMappings.put(outputVariable, inputMappings.get(outputVariable));
        } else {
            VariableReferenceExpression newVariable = context.getVariableAllocator().newVariable(outputVariable);
            outputMappings.put(outputVariable, newVariable);
            mappings.put(outputVariable, newVariable);
        }
    }
    sourceMappings.add(outputMappings.build());
    SymbolMapper symbolMapper = new SymbolMapper(mappings.build(), WarningCollector.NOOP);
    return symbolMapper.map(writerNode, unionNode.getSources().get(source), context.getIdAllocator().getNextId());
}