Examples with StageId com.facebook.presto.execution.StageId used on opensource projects

Example 11 with StageId

use of com.facebook.presto.execution.StageId in project presto by prestodb.

the class SectionExecutionFactory method createStreamingLinkedStageExecutions.

/**
 * returns a List of StageExecutionAndSchedulers in a postorder representation of the tree
 */
private List<StageExecutionAndScheduler> createStreamingLinkedStageExecutions(Session session, ExchangeLocationsConsumer parent, StreamingSubPlan plan, Function<PartitioningHandle, NodePartitionMap> partitioningCache, TableWriteInfo tableWriteInfo, Optional<SqlStageExecution> parentStageExecution, boolean summarizeTaskInfo, RemoteTaskFactory remoteTaskFactory, SplitSourceFactory splitSourceFactory, int attemptId) {
    ImmutableList.Builder<StageExecutionAndScheduler> stageExecutionAndSchedulers = ImmutableList.builder();
    PlanFragmentId fragmentId = plan.getFragment().getId();
    StageId stageId = new StageId(session.getQueryId(), fragmentId.getId());
    SqlStageExecution stageExecution = createSqlStageExecution(new StageExecutionId(stageId, attemptId), plan.getFragment(), remoteTaskFactory, session, summarizeTaskInfo, nodeTaskMap, executor, failureDetector, schedulerStats, tableWriteInfo);
    PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning();
    List<RemoteSourceNode> remoteSourceNodes = plan.getFragment().getRemoteSourceNodes();
    Optional<int[]> bucketToPartition = getBucketToPartition(partitioningHandle, partitioningCache, plan.getFragment().getRoot(), remoteSourceNodes);
    // create child stages
    ImmutableSet.Builder<SqlStageExecution> childStagesBuilder = ImmutableSet.builder();
    for (StreamingSubPlan stagePlan : plan.getChildren()) {
        List<StageExecutionAndScheduler> subTree = createStreamingLinkedStageExecutions(session, stageExecution::addExchangeLocations, stagePlan.withBucketToPartition(bucketToPartition), partitioningCache, tableWriteInfo, Optional.of(stageExecution), summarizeTaskInfo, remoteTaskFactory, splitSourceFactory, attemptId);
        stageExecutionAndSchedulers.addAll(subTree);
        childStagesBuilder.add(getLast(subTree).getStageExecution());
    }
    Set<SqlStageExecution> childStageExecutions = childStagesBuilder.build();
    stageExecution.addStateChangeListener(newState -> {
        if (newState.isDone()) {
            childStageExecutions.forEach(SqlStageExecution::cancel);
        }
    });
    StageLinkage stageLinkage = new StageLinkage(fragmentId, parent, childStageExecutions);
    StageScheduler stageScheduler = createStageScheduler(splitSourceFactory, session, plan, partitioningCache, parentStageExecution, stageId, stageExecution, partitioningHandle, tableWriteInfo, childStageExecutions);
    stageExecutionAndSchedulers.add(new StageExecutionAndScheduler(stageExecution, stageLinkage, stageScheduler));
    return stageExecutionAndSchedulers.build();
}

Example 12 with StageId

use of com.facebook.presto.execution.StageId in project presto by prestodb.

the class LegacySqlQueryScheduler method buildStageInfo.

private StageInfo buildStageInfo(SubPlan subPlan, Map<StageId, StageExecutionInfo> stageExecutionInfos) {
    StageId stageId = getStageId(subPlan.getFragment().getId());
    StageExecutionInfo stageExecutionInfo = stageExecutionInfos.get(stageId);
    checkArgument(stageExecutionInfo != null, "No stageExecutionInfo for %s", stageId);
    return new StageInfo(stageId, locationFactory.createStageLocation(stageId), Optional.of(subPlan.getFragment()), stageExecutionInfo, ImmutableList.of(), subPlan.getChildren().stream().map(plan -> buildStageInfo(plan, stageExecutionInfos)).collect(toImmutableList()), runtimeOptimizedStages.contains(stageId));
}

Example 13 with StageId

use of com.facebook.presto.execution.StageId in project presto by prestodb.

the class SqlQueryScheduler method createStageExecutions.

private List<SectionExecution> createStageExecutions(List<StreamingPlanSection> sections) {
    ImmutableList.Builder<SectionExecution> result = ImmutableList.builder();
    for (StreamingPlanSection section : sections) {
        StageId sectionId = getStageId(section.getPlan().getFragment().getId());
        List<SectionExecution> attempts = sectionExecutions.computeIfAbsent(sectionId, (ignored) -> new CopyOnWriteArrayList<>());
        // sectionExecutions only get created when they are about to be scheduled, so there should
        // never be a non-failed SectionExecution for a section that's ready for execution
        verify(attempts.isEmpty() || getLast(attempts).isFailed(), "Non-failed sectionExecutions already exists");
        PlanFragment sectionRootFragment = section.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) -> {
            };
        }
        int attemptId = attempts.size();
        SectionExecution sectionExecution = sectionExecutionFactory.createSectionExecutions(session, section, locationsConsumer, bucketToPartition, outputBuffers, summarizeTaskInfo, remoteTaskFactory, splitSourceFactory, attemptId);
        addStateChangeListeners(sectionExecution);
        attempts.add(sectionExecution);
        result.add(sectionExecution);
    }
    return result.build();
}

Example 14 with StageId

use of com.facebook.presto.execution.StageId in project presto by prestodb.

the class TestSourcePartitionedScheduler method createSqlStageExecution.

private SqlStageExecution createSqlStageExecution(SubPlan tableScanPlan, NodeTaskMap nodeTaskMap) {
    StageId stageId = new StageId(new QueryId("query"), 0);
    SqlStageExecution stage = SqlStageExecution.createSqlStageExecution(new StageExecutionId(stageId, 0), tableScanPlan.getFragment(), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new TableWriteInfo(Optional.empty(), Optional.empty(), Optional.empty()));
    stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
    return stage;
}