Examples with StageExecutionPlan io.prestosql.sql.planner.StageExecutionPlan used on opensource projects

Example 1 with StageExecutionPlan

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

the class SqlQueryExecution method createResumeScheduler.

private SqlQueryScheduler createResumeScheduler(PlanRoot plan, OutputBuffers rootOutputBuffers) {
    String resumeMessage = "Query encountered failures. Recovering using the distributed-snapshot feature.";
    warningCollector.add(new PrestoWarning(StandardWarningCode.SNAPSHOT_RECOVERY, resumeMessage));
    // Check if there is a snapshot we can restore to, or restart from beginning,
    // and update marker split sources so they know where to resume from.
    // This MUST be done BEFORE creating the new scheduler, because it resets the snapshotManager internal states.
    OptionalLong snapshotId = snapshotManager.getResumeSnapshotId();
    MarkerAnnouncer announcer = splitManager.getMarkerAnnouncer(stateMachine.getSession());
    announcer.resumeSnapshot(snapshotId.orElse(0));
    // Clear any temporary content that's not part of the snapshot
    resetOutputData(plan, snapshotId);
    // Create a new scheduler, to schedule new stages and tasks
    DistributedExecutionPlanner distributedExecutionPlanner = new DistributedExecutionPlanner(splitManager, metadata);
    StageExecutionPlan executionPlan = distributedExecutionPlanner.plan(plan.getRoot(), stateMachine.getSession(), RESUME, snapshotId.isPresent() ? snapshotId.getAsLong() : null, announcer.currentSnapshotId());
    // build the stage execution objects (this doesn't schedule execution)
    return createSqlQueryScheduler(stateMachine, locationFactory, executionPlan, nodePartitioningManager, nodeScheduler, remoteTaskFactory, stateMachine.getSession(), plan.isSummarizeTaskInfos(), scheduleSplitBatchSize, queryExecutor, schedulerExecutor, failureDetector, rootOutputBuffers, nodeTaskMap, executionPolicy, schedulerStats, dynamicFilterService, heuristicIndexerManager, snapshotManager, // Require same number of tasks to be scheduled, but do not require it if starting from beginning
    snapshotId.isPresent() ? queryScheduler.get().getStageTaskCounts() : null);
}

Example 2 with StageExecutionPlan

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

the class SqlQueryExecution method planDistribution.

private void planDistribution(PlanRoot plan) {
    // time distribution planning
    stateMachine.beginDistributedPlanning();
    // plan the execution on the active nodes
    DistributedExecutionPlanner distributedPlanner = new DistributedExecutionPlanner(splitManager, metadata);
    StageExecutionPlan outputStageExecutionPlan;
    Session session = stateMachine.getSession();
    if (SystemSessionProperties.isSnapshotEnabled(session)) {
        // Snapshot: need to plan different when snapshot is enabled.
        // See the "plan" method for difference between the different modes.
        MarkerAnnouncer announcer = splitManager.getMarkerAnnouncer(session);
        announcer.setSnapshotManager(snapshotManager);
        outputStageExecutionPlan = distributedPlanner.plan(plan.getRoot(), session, SNAPSHOT, null, announcer.currentSnapshotId());
    } else {
        outputStageExecutionPlan = distributedPlanner.plan(plan.getRoot(), session, NORMAL, null, 0);
    }
    stateMachine.endDistributedPlanning();
    // ensure split sources are closed
    stateMachine.addStateChangeListener(state -> {
        if (state.isDone()) {
            closeSplitSources(outputStageExecutionPlan);
        }
    });
    // if query was canceled, skip creating scheduler
    if (stateMachine.isDone()) {
        return;
    }
    // record output field
    stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes());
    PartitioningHandle partitioningHandle = plan.getRoot().getFragment().getPartitioningScheme().getPartitioning().getHandle();
    OutputBuffers rootOutputBuffers = createInitialEmptyOutputBuffers(partitioningHandle).withBuffer(OUTPUT_BUFFER_ID, BROADCAST_PARTITION_ID).withNoMoreBufferIds();
    // build the stage execution objects (this doesn't schedule execution)
    SqlQueryScheduler scheduler = createSqlQueryScheduler(stateMachine, locationFactory, outputStageExecutionPlan, nodePartitioningManager, nodeScheduler, remoteTaskFactory, stateMachine.getSession(), plan.isSummarizeTaskInfos(), scheduleSplitBatchSize, queryExecutor, schedulerExecutor, failureDetector, rootOutputBuffers, nodeTaskMap, executionPolicy, schedulerStats, dynamicFilterService, heuristicIndexerManager, snapshotManager, null);
    queryScheduler.set(scheduler);
    // directly since the callback may have already fired
    if (stateMachine.isDone()) {
        scheduler.abort();
        queryScheduler.set(null);
    }
}

Example 3 with StageExecutionPlan

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBlock.

@Test
public void testScheduleSplitsBlock() {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(80, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(plan, 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()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 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()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 20);
    }
    stage.abort();
}

Example 4 with StageExecutionPlan

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

the class TestSourcePartitionedScheduler method testScheduleSplitsOneAtATime.

@Test
public void testScheduleSplitsOneAtATime() {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    try (StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1)) {
        for (int i = 0; i < 60; i++) {
            ScheduleResult scheduleResult = scheduler.schedule();
            // only finishes when last split is fetched
            if (i == 59) {
                assertEffectivelyFinished(scheduleResult, scheduler);
            } else {
                assertFalse(scheduleResult.isFinished());
            }
            // never blocks
            assertTrue(scheduleResult.getBlocked().isDone());
            // 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()) {
            assertEquals(remoteTask.getPartitionedSplitCount(), 20);
        }
        stage.abort();
    }
}

Example 5 with StageExecutionPlan

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

the class TestSourcePartitionedScheduler method testScheduleNoSplits.

@Test
public void testScheduleNoSplits() {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(0, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1);
    ScheduleResult scheduleResult = scheduler.schedule();
    assertEquals(scheduleResult.getNewTasks().size(), 1);
    assertEffectivelyFinished(scheduleResult, scheduler);
    stage.abort();
}