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

Example 1 with OutputBuffers

use of com.facebook.presto.execution.buffer.OutputBuffers 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);
    }
}

Example 2 with OutputBuffers

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

the class ScaledOutputBufferManager method addOutputBuffers.

@SuppressWarnings("ObjectEquality")
@Override
public void addOutputBuffers(List<OutputBufferId> newBuffers, boolean noMoreBuffers) {
    OutputBuffers newOutputBuffers;
    synchronized (this) {
        if (outputBuffers.isNoMoreBufferIds()) {
            // so ignore the new buffers
            return;
        }
        OutputBuffers originalOutputBuffers = outputBuffers;
        for (OutputBufferId newBuffer : newBuffers) {
            outputBuffers = outputBuffers.withBuffer(newBuffer, newBuffer.getId());
        }
        if (noMoreBuffers) {
            outputBuffers = outputBuffers.withNoMoreBufferIds();
        }
        // don't update if nothing changed
        if (outputBuffers == originalOutputBuffers) {
            return;
        }
        newOutputBuffers = this.outputBuffers;
    }
    outputBufferTarget.accept(newOutputBuffers);
}

Example 3 with OutputBuffers

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

the class SqlQueryExecution method planDistribution.

private void planDistribution(PlanRoot plan) {
    CloseableSplitSourceProvider splitSourceProvider = new CloseableSplitSourceProvider(splitManager::getSplits);
    // ensure split sources are closed
    stateMachine.addStateChangeListener(state -> {
        if (state.isDone()) {
            splitSourceProvider.close();
        }
    });
    // if query was canceled, skip creating scheduler
    if (stateMachine.isDone()) {
        return;
    }
    SubPlan outputStagePlan = plan.getRoot();
    // record output field
    stateMachine.setColumns(((OutputNode) outputStagePlan.getFragment().getRoot()).getColumnNames(), outputStagePlan.getFragment().getTypes());
    PartitioningHandle partitioningHandle = outputStagePlan.getFragment().getPartitioningScheme().getPartitioning().getHandle();
    OutputBuffers rootOutputBuffers;
    if (isSpoolingOutputBufferEnabled(getSession())) {
        rootOutputBuffers = createSpoolingOutputBuffers();
    } else {
        rootOutputBuffers = createInitialEmptyOutputBuffers(partitioningHandle).withBuffer(OUTPUT_BUFFER_ID, BROADCAST_PARTITION_ID).withNoMoreBufferIds();
    }
    SplitSourceFactory splitSourceFactory = new SplitSourceFactory(splitSourceProvider, stateMachine.getWarningCollector());
    // build the stage execution objects (this doesn't schedule execution)
    SqlQuerySchedulerInterface scheduler = isUseLegacyScheduler(getSession()) ? LegacySqlQueryScheduler.createSqlQueryScheduler(locationFactory, executionPolicy, queryExecutor, schedulerStats, sectionExecutionFactory, remoteTaskFactory, splitSourceFactory, stateMachine.getSession(), metadata.getFunctionAndTypeManager(), stateMachine, outputStagePlan, rootOutputBuffers, plan.isSummarizeTaskInfos(), runtimePlanOptimizers, stateMachine.getWarningCollector(), idAllocator, variableAllocator.get(), planChecker, metadata, sqlParser, partialResultQueryManager) : SqlQueryScheduler.createSqlQueryScheduler(locationFactory, executionPolicy, queryExecutor, schedulerStats, sectionExecutionFactory, remoteTaskFactory, splitSourceFactory, internalNodeManager, stateMachine.getSession(), stateMachine, outputStagePlan, plan.isSummarizeTaskInfos(), metadata.getFunctionAndTypeManager(), runtimePlanOptimizers, stateMachine.getWarningCollector(), idAllocator, variableAllocator.get(), planChecker, metadata, sqlParser, partialResultQueryManager);
    queryScheduler.set(scheduler);
    // directly since the callback may have already fired
    if (stateMachine.isDone()) {
        scheduler.abort();
        queryScheduler.set(null);
    }
}

Example 4 with OutputBuffers

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

the class SqlStageExecution method scheduleTask.

private synchronized RemoteTask scheduleTask(InternalNode node, TaskId taskId, Multimap<PlanNodeId, Split> sourceSplits) {
    checkArgument(!allTasks.contains(taskId), "A task with id %s already exists", taskId);
    ImmutableMultimap.Builder<PlanNodeId, Split> initialSplits = ImmutableMultimap.builder();
    initialSplits.putAll(sourceSplits);
    sourceTasks.forEach((planNodeId, task) -> {
        TaskStatus status = task.getTaskStatus();
        if (status.getState() != TaskState.FINISHED) {
            initialSplits.put(planNodeId, createRemoteSplitFor(taskId, task.getRemoteTaskLocation(), task.getTaskId()));
        }
    });
    OutputBuffers outputBuffers = this.outputBuffers.get();
    checkState(outputBuffers != null, "Initial output buffers must be set before a task can be scheduled");
    RemoteTask task = remoteTaskFactory.createRemoteTask(session, taskId, node, planFragment, initialSplits.build(), outputBuffers, nodeTaskMap.createTaskStatsTracker(node, taskId), summarizeTaskInfo, tableWriteInfo);
    completeSources.forEach(task::noMoreSplits);
    allTasks.add(taskId);
    tasks.computeIfAbsent(node, key -> newConcurrentHashSet()).add(task);
    nodeTaskMap.addTask(node, task);
    task.addStateChangeListener(new StageTaskListener(taskId));
    task.addFinalTaskInfoListener(this::updateFinalTaskInfo);
    if (!stateMachine.getState().isDone()) {
        task.start();
    } else {
        // stage finished while we were scheduling this task
        task.abort();
    }
    return task;
}

Example 5 with OutputBuffers

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

the class BroadcastOutputBufferManager method addOutputBuffers.

@Override
public void addOutputBuffers(List<OutputBufferId> newBuffers, boolean noMoreBuffers) {
    OutputBuffers newOutputBuffers;
    synchronized (this) {
        if (outputBuffers.isNoMoreBufferIds()) {
            // the new buffers
            return;
        }
        OutputBuffers originalOutputBuffers = outputBuffers;
        // Note: it does not matter which partition id the task is using, in broadcast all tasks read from the same partition
        for (OutputBufferId newBuffer : newBuffers) {
            outputBuffers = outputBuffers.withBuffer(newBuffer, BROADCAST_PARTITION_ID);
        }
        if (noMoreBuffers) {
            outputBuffers = outputBuffers.withNoMoreBufferIds();
        }
        // don't update if nothing changed
        if (outputBuffers == originalOutputBuffers) {
            return;
        }
        newOutputBuffers = this.outputBuffers;
    }
    outputBufferTarget.accept(newOutputBuffers);
}