Examples with ExchangeNode io.prestosql.sql.planner.plan.ExchangeNode used on opensource projects

Example 6 with ExchangeNode

use of io.prestosql.sql.planner.plan.ExchangeNode in project hetu-core by openlookeng.

the class PushPartialAggregationThroughExchange method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    ExchangeNode exchangeNode = captures.get(EXCHANGE_NODE);
    // Do not push down aggregation since CTEScanNode requires an exchange node above it
    if (exchangeNode.getSources().size() == 1 && context.getLookup().resolve(exchangeNode.getSources().get(0)) instanceof CTEScanNode) {
        return Result.empty();
    }
    boolean decomposable = isDecomposable(aggregationNode, metadata);
    if (aggregationNode.getStep().equals(SINGLE) && aggregationNode.hasEmptyGroupingSet() && aggregationNode.hasNonEmptyGroupingSet() && exchangeNode.getType() == REPARTITION) {
        // single-step aggregation w/ empty grouping sets in a partitioned stage, so we need a partial that will produce
        // the default intermediates for the empty grouping set that will be routed to the appropriate final aggregation.
        // TODO: technically, AddExchanges generates a broken plan that this rule "fixes"
        checkState(decomposable, "Distributed aggregation with empty grouping set requires partial but functions are not decomposable");
        return Result.ofPlanNode(split(aggregationNode, context));
    }
    if (!decomposable || !preferPartialAggregation(context.getSession())) {
        return Result.empty();
    }
    // the cardinality of the stream (i.e., gather or repartition)
    if ((exchangeNode.getType() != GATHER && exchangeNode.getType() != REPARTITION) || exchangeNode.getPartitioningScheme().isReplicateNullsAndAny()) {
        return Result.empty();
    }
    if (exchangeNode.getType() == REPARTITION) {
        // if partitioning columns are not a subset of grouping keys,
        // we can't push this through
        List<Symbol> partitioningColumns = exchangeNode.getPartitioningScheme().getPartitioning().getArguments().stream().filter(Partitioning.ArgumentBinding::isVariable).map(Partitioning.ArgumentBinding::getColumn).collect(Collectors.toList());
        if (!aggregationNode.getGroupingKeys().containsAll(partitioningColumns)) {
            return Result.empty();
        }
    }
    // currently, we only support plans that don't use pre-computed hash functions
    if (aggregationNode.getHashSymbol().isPresent() || exchangeNode.getPartitioningScheme().getHashColumn().isPresent()) {
        return Result.empty();
    }
    switch(aggregationNode.getStep()) {
        case SINGLE:
            // Split it into a FINAL on top of a PARTIAL and
            return Result.ofPlanNode(split(aggregationNode, context));
        case PARTIAL:
            // Push it underneath each branch of the exchange
            return Result.ofPlanNode(pushPartial(aggregationNode, exchangeNode, context));
        default:
            return Result.empty();
    }
}

Example 7 with ExchangeNode

use of io.prestosql.sql.planner.plan.ExchangeNode in project hetu-core by openlookeng.

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.getOutputSymbols().size(); outputIndex++) {
            Symbol output = exchange.getOutputSymbols().get(outputIndex);
            Symbol input = exchange.getInputs().get(i).get(outputIndex);
            if (!output.equals(input)) {
                mappingsBuilder.put(output, input);
            }
        }
        SymbolMapper symbolMapper = mappingsBuilder.build();
        if (symbolMapper.getTypes() == null) {
            symbolMapper.setTypes(context.getSymbolAllocator().getTypes());
        }
        AggregationNode mappedPartial = symbolMapper.map(aggregation, source, context.getIdAllocator());
        Assignments.Builder assignments = Assignments.builder();
        for (Symbol output : aggregation.getOutputSymbols()) {
            Symbol input = symbolMapper.map(output);
            assignments.put(output, VariableReferenceSymbolConverter.toVariableReference(input, context.getSymbolAllocator().getTypes()));
        }
        partials.add(new ProjectNode(context.getIdAllocator().getNextId(), mappedPartial, assignments.build()));
    }
    for (PlanNode node : partials) {
        verify(aggregation.getOutputSymbols().equals(node.getOutputSymbols()));
    }
    // 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
    PartitioningScheme partitioning = new PartitioningScheme(exchange.getPartitioningScheme().getPartitioning(), aggregation.getOutputSymbols(), exchange.getPartitioningScheme().getHashColumn(), exchange.getPartitioningScheme().isReplicateNullsAndAny(), exchange.getPartitioningScheme().getBucketToPartition());
    return new ExchangeNode(context.getIdAllocator().getNextId(), exchange.getType(), exchange.getScope(), partitioning, partials, ImmutableList.copyOf(Collections.nCopies(partials.size(), aggregation.getOutputSymbols())), Optional.empty(), aggregation.getAggregationType());
}

Example 8 with ExchangeNode

use of io.prestosql.sql.planner.plan.ExchangeNode in project hetu-core by openlookeng.

the class PushRemoteExchangeThroughAssignUniqueId method apply.

@Override
public Result apply(ExchangeNode node, Captures captures, Context context) {
    checkArgument(!node.getOrderingScheme().isPresent(), "Merge exchange over AssignUniqueId not supported");
    AssignUniqueId assignUniqueId = captures.get(ASSIGN_UNIQUE_ID);
    PartitioningScheme partitioningScheme = node.getPartitioningScheme();
    if (partitioningScheme.getPartitioning().getColumns().contains(assignUniqueId.getIdColumn())) {
        // Hence, AssignUniqueId node has to stay below the exchange node.
        return Result.empty();
    }
    return Result.ofPlanNode(new AssignUniqueId(assignUniqueId.getId(), new ExchangeNode(node.getId(), node.getType(), node.getScope(), new PartitioningScheme(partitioningScheme.getPartitioning(), removeSymbol(partitioningScheme.getOutputLayout(), assignUniqueId.getIdColumn()), partitioningScheme.getHashColumn(), partitioningScheme.isReplicateNullsAndAny(), partitioningScheme.getBucketToPartition()), ImmutableList.of(assignUniqueId.getSource()), ImmutableList.of(removeSymbol(getOnlyElement(node.getInputs()), assignUniqueId.getIdColumn())), Optional.empty(), node.getAggregationType()), assignUniqueId.getIdColumn()));
}

Example 9 with ExchangeNode

use of io.prestosql.sql.planner.plan.ExchangeNode in project hetu-core by openlookeng.

the class AddIntermediateAggregations method addGatheringIntermediate.

private PlanNode addGatheringIntermediate(AggregationNode aggregation, PlanNodeIdAllocator idAllocator) {
    verify(aggregation.getGroupingKeys().isEmpty(), "Should be an un-grouped aggregation");
    ExchangeNode gatheringExchange = ExchangeNode.gatheringExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, aggregation);
    return new AggregationNode(idAllocator.getNextId(), gatheringExchange, outputsAsInputs(aggregation.getAggregations()), aggregation.getGroupingSets(), aggregation.getPreGroupedSymbols(), AggregationNode.Step.INTERMEDIATE, aggregation.getHashSymbol(), aggregation.getGroupIdSymbol(), aggregation.getAggregationType(), aggregation.getFinalizeSymbol());
}

Example 10 with ExchangeNode

use of io.prestosql.sql.planner.plan.ExchangeNode in project hetu-core by openlookeng.

the class AddIntermediateAggregations method recurseToPartial.

/**
 * Recurse through a series of preceding ExchangeNodes and ProjectNodes to find the preceding PARTIAL aggregation
 */
private Optional<PlanNode> recurseToPartial(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator) {
    if (node instanceof AggregationNode && ((AggregationNode) node).getStep() == AggregationNode.Step.PARTIAL) {
        return Optional.of(addGatheringIntermediate((AggregationNode) node, idAllocator));
    }
    if (!(node instanceof ExchangeNode) && !(node instanceof ProjectNode)) {
        return Optional.empty();
    }
    ImmutableList.Builder<PlanNode> builder = ImmutableList.builder();
    for (PlanNode source : node.getSources()) {
        Optional<PlanNode> planNode = recurseToPartial(lookup.resolve(source), lookup, idAllocator);
        if (!planNode.isPresent()) {
            return Optional.empty();
        }
        builder.add(planNode.get());
    }
    return Optional.of(node.replaceChildren(builder.build()));
}