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

Example 16 with SqlStageExecution

use of com.facebook.presto.execution.SqlStageExecution 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 17 with SqlStageExecution

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

the class LegacySqlQueryScheduler method initialize.

// this is a separate method to ensure that the `this` reference is not leaked during construction
private void initialize() {
    SqlStageExecution rootStage = stageExecutions.get(rootStageId).getStageExecution();
    rootStage.addStateChangeListener(state -> {
        if (state == FINISHED) {
            queryStateMachine.transitionToFinishing();
        } else if (state == CANCELED) {
            // output stage was canceled
            queryStateMachine.transitionToCanceled();
        }
    });
    for (StageExecutionAndScheduler stageExecutionInfo : stageExecutions.values()) {
        SqlStageExecution stageExecution = stageExecutionInfo.getStageExecution();
        stageExecution.addStateChangeListener(state -> {
            if (queryStateMachine.isDone()) {
                return;
            }
            if (state == FAILED) {
                queryStateMachine.transitionToFailed(stageExecution.getStageExecutionInfo().getFailureCause().get().toException());
            } else if (state == ABORTED) {
                // this should never happen, since abort can only be triggered in query clean up after the query is finished
                queryStateMachine.transitionToFailed(new PrestoException(GENERIC_INTERNAL_ERROR, "Query stage was aborted"));
            } else if (state == FINISHED) {
                // checks if there's any new sections available for execution and starts the scheduling if any
                startScheduling();
            } else if (queryStateMachine.getQueryState() == QueryState.STARTING) {
                // if the stage has at least one task, we are running
                if (stageExecution.hasTasks()) {
                    queryStateMachine.transitionToRunning();
                }
            }
        });
        stageExecution.addFinalStageInfoListener(status -> queryStateMachine.updateQueryInfo(Optional.of(getStageInfo())));
    }
    // when query is done or any time a stage completes, attempt to transition query to "final query info ready"
    queryStateMachine.addStateChangeListener(newState -> {
        if (newState.isDone()) {
            queryStateMachine.updateQueryInfo(Optional.of(getStageInfo()));
        }
    });
}

Example 18 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBlock.

@Test
public void testScheduleSplitsBlock() {
    SubPlan plan = createPlan();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(80, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 1);
    // schedule first 60 splits, which will cause the scheduler to block
    for (int i = 0; i <= 60; i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        assertFalse(scheduleResult.isFinished());
        // blocks at 20 per node
        assertEquals(scheduleResult.getBlocked().isDone(), i != 60);
        // first three splits create new tasks
        assertEquals(scheduleResult.getNewTasks().size(), i < 3 ? 1 : 0);
        assertEquals(stage.getAllTasks().size(), i < 3 ? i + 1 : 3);
        assertPartitionedSplitCount(stage, min(i + 1, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    // todo rewrite MockRemoteTask to fire a tate transition when splits are cleared, and then validate blocked future completes
    // drop the 20 splits from one node
    ((MockRemoteTask) stage.getAllTasks().get(0)).clearSplits();
    // schedule remaining 20 splits
    for (int i = 0; i < 20; i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // finishes when last split is fetched
        if (i == 19) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // does not block again
        assertTrue(scheduleResult.getBlocked().isDone());
        // no additional tasks will be created
        assertEquals(scheduleResult.getNewTasks().size(), 0);
        assertEquals(stage.getAllTasks().size(), 3);
        // we dropped 20 splits so start at 40 and count to 60
        assertPartitionedSplitCount(stage, min(i + 41, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    stage.abort();
}

Example 19 with SqlStageExecution

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

Example 20 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBatched.

@Test
public void testScheduleSplitsBatched() {
    SubPlan plan = createPlan();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(60, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 7);
    for (int i = 0; i <= (60 / 7); i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // finishes when last split is fetched
        if (i == (60 / 7)) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // never blocks
        assertTrue(scheduleResult.getBlocked().isDone());
        // first three splits create new tasks
        assertEquals(scheduleResult.getNewTasks().size(), i == 0 ? 3 : 0);
        assertEquals(stage.getAllTasks().size(), 3);
        assertPartitionedSplitCount(stage, min((i + 1) * 7, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    stage.abort();
}