Examples with BROADCAST_PARTITION_ID com.facebook.presto.execution.buffer.OutputBuffers.BROADCAST_PARTITION_ID used on opensource projects

Example 1 with BROADCAST_PARTITION_ID

use of com.facebook.presto.execution.buffer.OutputBuffers.BROADCAST_PARTITION_ID in project presto by prestodb.

the class LegacySqlQueryScheduler method updateStageExecutions.

/**
 * Utility function that rebuild a StreamingPlanSection, re-create stageExecutionAndScheduler for each of its stage, and finally update the stageExecutions map.
 */
private void updateStageExecutions(StreamingPlanSection section, Map<PlanFragment, PlanFragment> oldToNewFragment) {
    StreamingPlanSection newSection = new StreamingPlanSection(rewriteStreamingSubPlan(section.getPlan(), oldToNewFragment), section.getChildren());
    PlanFragment sectionRootFragment = newSection.getPlan().getFragment();
    Optional<int[]> bucketToPartition;
    OutputBuffers outputBuffers;
    ExchangeLocationsConsumer locationsConsumer;
    if (isRootFragment(sectionRootFragment)) {
        bucketToPartition = Optional.of(new int[1]);
        outputBuffers = createInitialEmptyOutputBuffers(sectionRootFragment.getPartitioningScheme().getPartitioning().getHandle()).withBuffer(new OutputBufferId(0), BROADCAST_PARTITION_ID).withNoMoreBufferIds();
        OutputBufferId rootBufferId = getOnlyElement(outputBuffers.getBuffers().keySet());
        locationsConsumer = (fragmentId, tasks, noMoreExchangeLocations) -> updateQueryOutputLocations(queryStateMachine, rootBufferId, tasks, noMoreExchangeLocations);
    } else {
        bucketToPartition = Optional.empty();
        outputBuffers = createDiscardingOutputBuffers();
        locationsConsumer = (fragmentId, tasks, noMoreExchangeLocations) -> {
        };
    }
    SectionExecution sectionExecution = sectionExecutionFactory.createSectionExecutions(session, newSection, locationsConsumer, bucketToPartition, outputBuffers, summarizeTaskInfo, remoteTaskFactory, splitSourceFactory, 0);
    addStateChangeListeners(sectionExecution);
    Map<StageId, StageExecutionAndScheduler> updatedStageExecutions = sectionExecution.getSectionStages().stream().collect(toImmutableMap(execution -> execution.getStageExecution().getStageExecutionId().getStageId(), identity()));
    synchronized (this) {
        stageExecutions.putAll(updatedStageExecutions);
    }
}