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

Example 6 with SqlStageExecution

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

the class SqlQueryScheduler method createStages.

private List<SqlStageExecution> createStages(Optional<SqlStageExecution> parent, AtomicInteger nextStageId, LocationFactory locationFactory, StageExecutionPlan plan, NodeScheduler nodeScheduler, RemoteTaskFactory remoteTaskFactory, Session session, int splitBatchSize, Function<PartitioningHandle, NodePartitionMap> partitioningCache, ExecutorService executor, NodeTaskMap nodeTaskMap, ImmutableMap.Builder<StageId, StageScheduler> stageSchedulers, ImmutableMap.Builder<StageId, StageLinkage> stageLinkages) {
    ImmutableList.Builder<SqlStageExecution> stages = ImmutableList.builder();
    StageId stageId = new StageId(queryStateMachine.getQueryId(), nextStageId.getAndIncrement());
    SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), remoteTaskFactory, session, summarizeTaskInfo, nodeTaskMap, executor, schedulerStats);
    stages.add(stage);
    Optional<int[]> bucketToPartition;
    PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning();
    if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) {
        // nodes are selected dynamically based on the constraints of the splits and the system load
        Entry<PlanNodeId, SplitSource> entry = Iterables.getOnlyElement(plan.getSplitSources().entrySet());
        ConnectorId connectorId = entry.getValue().getConnectorId();
        if (isInternalSystemConnector(connectorId)) {
            connectorId = null;
        }
        NodeSelector nodeSelector = nodeScheduler.createNodeSelector(connectorId);
        SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stage::getAllTasks);
        stageSchedulers.put(stageId, new SourcePartitionedScheduler(stage, entry.getKey(), entry.getValue(), placementPolicy, splitBatchSize));
        bucketToPartition = Optional.of(new int[1]);
    } else {
        // nodes are pre determined by the nodePartitionMap
        NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning());
        Map<PlanNodeId, SplitSource> splitSources = plan.getSplitSources();
        if (!splitSources.isEmpty()) {
            stageSchedulers.put(stageId, new FixedSourcePartitionedScheduler(stage, splitSources, plan.getFragment().getPartitionedSources(), nodePartitionMap, splitBatchSize, nodeScheduler.createNodeSelector(null)));
            bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
        } else {
            Map<Integer, Node> partitionToNode = nodePartitionMap.getPartitionToNode();
            // todo this should asynchronously wait a standard timeout period before failing
            checkCondition(!partitionToNode.isEmpty(), NO_NODES_AVAILABLE, "No worker nodes available");
            stageSchedulers.put(stageId, new FixedCountScheduler(stage, partitionToNode));
            bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
        }
    }
    ImmutableSet.Builder<SqlStageExecution> childStagesBuilder = ImmutableSet.builder();
    for (StageExecutionPlan subStagePlan : plan.getSubStages()) {
        List<SqlStageExecution> subTree = createStages(Optional.of(stage), nextStageId, locationFactory, subStagePlan.withBucketToPartition(bucketToPartition), nodeScheduler, remoteTaskFactory, session, splitBatchSize, partitioningCache, executor, nodeTaskMap, stageSchedulers, stageLinkages);
        stages.addAll(subTree);
        SqlStageExecution childStage = subTree.get(0);
        childStagesBuilder.add(childStage);
    }
    Set<SqlStageExecution> childStages = childStagesBuilder.build();
    stage.addStateChangeListener(newState -> {
        if (newState.isDone()) {
            childStages.forEach(SqlStageExecution::cancel);
        }
    });
    stageLinkages.put(stageId, new StageLinkage(plan.getFragment().getId(), parent, childStages));
    return stages.build();
}

Example 7 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method testScheduleSlowSplitSource.

@Test
public void testScheduleSlowSplitSource() throws Exception {
    QueuedSplitSource queuedSplitSource = new QueuedSplitSource(TestingSplit::createRemoteSplit);
    StageExecutionPlan plan = createPlan(queuedSplitSource);
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    SourcePartitionedScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1);
    // schedule with no splits - will block
    ScheduleResult scheduleResult = scheduler.schedule();
    assertFalse(scheduleResult.isFinished());
    assertFalse(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 0);
    assertEquals(stage.getAllTasks().size(), 0);
    queuedSplitSource.addSplits(1);
    assertTrue(scheduleResult.getBlocked().isDone());
}

Example 8 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBlock.

@Test
public void testScheduleSplitsBlock() throws Exception {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(80, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    SourcePartitionedScheduler 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
        assertEquals(scheduleResult.isFinished(), i == 19);
        // 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 9 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method createSqlStageExecution.

private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) {
    StageId stageId = new StageId(new QueryId("query"), 0);
    SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(executor), TEST_SESSION, true, nodeTaskMap, executor, new SplitSchedulerStats());
    stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
    return stage;
}

Example 10 with SqlStageExecution

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

the class TestSourcePartitionedScheduler method testScheduleSplitsOneAtATime.

@Test
public void testScheduleSplitsOneAtATime() throws Exception {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    SourcePartitionedScheduler 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
        assertEquals(scheduleResult.isFinished(), i == 59);
        // 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();
}