Examples with StageExecutionPlan com.facebook.presto.sql.planner.StageExecutionPlan used on opensource projects

Example 1 with StageExecutionPlan

use of com.facebook.presto.sql.planner.StageExecutionPlan 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 2 with StageExecutionPlan

use of com.facebook.presto.sql.planner.StageExecutionPlan in project presto by prestodb.

the class TestSourcePartitionedScheduler method createPlan.

private static StageExecutionPlan createPlan(ConnectorSplitSource splitSource) {
    Symbol symbol = new Symbol("column");
    // table scan with splitCount splits
    PlanNodeId tableScanNodeId = new PlanNodeId("plan_id");
    TableScanNode tableScan = new TableScanNode(tableScanNodeId, new TableHandle(CONNECTOR_ID, new TestingTableHandle()), ImmutableList.of(symbol), ImmutableMap.of(symbol, new TestingColumnHandle("column")), Optional.empty(), TupleDomain.all(), null);
    RemoteSourceNode remote = new RemoteSourceNode(new PlanNodeId("remote_id"), new PlanFragmentId("plan_fragment_id"), ImmutableList.of());
    PlanFragment testFragment = new PlanFragment(new PlanFragmentId("plan_id"), new JoinNode(new PlanNodeId("join_id"), INNER, tableScan, remote, ImmutableList.of(), ImmutableList.<Symbol>builder().addAll(tableScan.getOutputSymbols()).addAll(remote.getOutputSymbols()).build(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(JoinNode.DistributionType.PARTITIONED)), ImmutableMap.of(symbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(tableScanNodeId), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)));
    return new StageExecutionPlan(testFragment, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, TestingTransactionHandle.create(), splitSource)), ImmutableList.of());
}