public StageExecutionPlan withBucketToPartition(Optional<int[]> bucketToPartition) { return new StageExecutionPlan(fragment.withBucketToPartition(bucketToPartition), splitSources, subStages); }
stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), remoteTaskFactory, session, PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning(); if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) { Entry<PlanNodeId, SplitSource> entry = Iterables.getOnlyElement(plan.getSplitSources().entrySet()); PlanNodeId planNodeId = entry.getKey(); SplitSource splitSource = entry.getValue(); checkArgument(!plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution()); stageSchedulers.put(stageId, newSourcePartitionedSchedulerAsStageScheduler(stage, planNodeId, splitSource, placementPolicy, splitBatchSize)); bucketToPartition = Optional.of(new int[1]); Map<PlanNodeId, SplitSource> splitSources = plan.getSplitSources(); if (!splitSources.isEmpty()) { List<PlanNodeId> schedulingOrder = plan.getFragment().getPartitionedSources(); ConnectorId connectorId = partitioningHandle.getConnectorId().orElseThrow(IllegalStateException::new); List<ConnectorPartitionHandle> connectorPartitionHandles; boolean groupedExecutionForStage = plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution(); if (groupedExecutionForStage) { connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle); if (plan.getSubStages().isEmpty()) {
private static void closeSplitSources(StageExecutionPlan plan) { for (SplitSource source : plan.getSplitSources().values()) { try { source.close(); } catch (Throwable t) { log.warn(t, "Error closing split source"); } } for (StageExecutionPlan stage : plan.getSubStages()) { closeSplitSources(stage); } }
stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), remoteTaskFactory, session, PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning(); if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) { SplitSource splitSource = plan.getDataSource().get(); NodeSelector nodeSelector = nodeScheduler.createNodeSelector(splitSource.getDataSourceName()); SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stage::getAllTasks); NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning()); if (plan.getDataSource().isPresent()) { stageSchedulers.put(stageId, new FixedSourcePartitionedScheduler( stage, plan.getDataSource().get(), nodePartitionMap, splitBatchSize, for (StageExecutionPlan subStagePlan : plan.getSubStages()) { List<SqlStageExecution> subTree = createStages( Optional.of(stage), nextStageId, locationFactory, subStagePlan.withBucketToPartition(bucketToPartition), nodeScheduler,
stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes());
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), 0); SqlStageExecution stage = SqlStageExecution.createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new NoOpFailureDetector(), new SplitSchedulerStats()); stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
private static StageScheduler getSourcePartitionedScheduler( StageExecutionPlan plan, SqlStageExecution stage, InternalNodeManager nodeManager, NodeTaskMap nodeTaskMap, int splitBatchSize) { NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig() .setIncludeCoordinator(false) .setMaxSplitsPerNode(20) .setMaxPendingSplitsPerTask(0); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap); PlanNodeId sourceNode = Iterables.getOnlyElement(plan.getSplitSources().keySet()); SplitSource splitSource = Iterables.getOnlyElement(plan.getSplitSources().values()); SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(splitSource.getConnectorId()), stage::getAllTasks); return newSourcePartitionedSchedulerAsStageScheduler(stage, sourceNode, splitSource, placementPolicy, splitBatchSize); }
new AtomicInteger(), locationFactory, plan.withBucketToPartition(Optional.of(new int[1])), nodeScheduler, remoteTaskFactory,
private static SourcePartitionedScheduler getSourcePartitionedScheduler( StageExecutionPlan plan, SqlStageExecution stage, NodeManager nodeManager, NodeTaskMap nodeTaskMap, int splitBatchSize) { NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig() .setIncludeCoordinator(false) .setMaxSplitsPerNode(20) .setMaxPendingSplitsPerNodePerTask(0); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap); SplitSource splitSource = plan.getDataSource().get(); SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(splitSource.getDataSourceName()), stage::getAllTasks); return new SourcePartitionedScheduler(stage, splitSource, placementPolicy, splitBatchSize); }
stateMachine.setOutputFieldNames(outputStageExecutionPlan.getFieldNames());
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), "stage"); SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(executor), TEST_SESSION, nodeTaskMap, executor); stage.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
@Test public void testNoNodes() { try { NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); InMemoryNodeManager nodeManager = new InMemoryNodeManager(); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, new NodeSchedulerConfig().setIncludeCoordinator(false), nodeTaskMap); StageExecutionPlan plan = createPlan(createFixedSplitSource(20, TestingSplit::createRemoteSplit)); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); StageScheduler scheduler = newSourcePartitionedSchedulerAsStageScheduler( stage, Iterables.getOnlyElement(plan.getSplitSources().keySet()), Iterables.getOnlyElement(plan.getSplitSources().values()), new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(CONNECTOR_ID), stage::getAllTasks), 2); scheduler.schedule(); fail("expected PrestoException"); } catch (PrestoException e) { assertEquals(e.getErrorCode(), NO_NODES_AVAILABLE.toErrorCode()); } }
new AtomicInteger(), locationFactory, plan.withBucketToPartition(Optional.of(new int[1])), nodeScheduler, remoteTaskFactory,
@Test public void testNoNodes() throws Exception { try { NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); InMemoryNodeManager nodeManager = new InMemoryNodeManager(); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, new NodeSchedulerConfig().setIncludeCoordinator(false), nodeTaskMap); StageExecutionPlan plan = createPlan(createFixedSplitSource(20, TestingSplit::createRemoteSplit)); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); SourcePartitionedScheduler scheduler = new SourcePartitionedScheduler( stage, plan.getDataSource().get(), new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector("test"), stage::getAllTasks), 2); scheduler.schedule(); fail("expected PrestoException"); } catch (PrestoException e) { assertEquals(e.getErrorCode(), NO_NODES_AVAILABLE.toErrorCode()); } }
private StageExecutionPlan doPlan(SubPlan root, Session session, ImmutableList.Builder<SplitSource> allSplitSources) { PlanFragment currentFragment = root.getFragment(); // get splits for this fragment, this is lazy so split assignments aren't actually calculated here Map<PlanNodeId, SplitSource> splitSources = currentFragment.getRoot().accept(new Visitor(session, currentFragment.getStageExecutionDescriptor(), allSplitSources), null); // create child stages ImmutableList.Builder<StageExecutionPlan> dependencies = ImmutableList.builder(); for (SubPlan childPlan : root.getChildren()) { dependencies.add(doPlan(childPlan, session, allSplitSources)); } return new StageExecutionPlan( currentFragment, splitSources, dependencies.build()); }
StatsAndCosts.empty()); return new StageExecutionPlan( testFragment, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, TestingTransactionHandle.create(), splitSource)),
public StageExecutionPlan withBucketToPartition(Optional<int[]> bucketToPartition) { return new StageExecutionPlan(fragment.withBucketToPartition(bucketToPartition), dataSource, subStages); }
public StageExecutionPlan plan(SubPlan root, Session session) { PlanFragment currentFragment = root.getFragment(); // get splits for this fragment, this is lazy so split assignments aren't actually calculated here Visitor visitor = new Visitor(session); Optional<SplitSource> splits = currentFragment.getRoot().accept(visitor, null); // create child stages ImmutableList.Builder<StageExecutionPlan> dependencies = ImmutableList.builder(); for (SubPlan childPlan : root.getChildren()) { dependencies.add(plan(childPlan, session)); } return new StageExecutionPlan( currentFragment, splits, dependencies.build()); }
private static StageExecutionPlan createPlan(ConnectorSplitSource splitSource) { Symbol symbol = new Symbol("column"); // table scan with splitCount splits PlanNodeId tableScanNodeId = new PlanNodeId("plan_id"); PlanFragment testFragment = new PlanFragment( new PlanFragmentId("plan_id"), new JoinNode(new PlanNodeId("join_id"), INNER, new TableScanNode( tableScanNodeId, new TableHandle(CONNECTOR_ID, new TestingTableHandle()), ImmutableList.of(symbol), ImmutableMap.of(symbol, new TestingColumnHandle("column")), Optional.empty(), TupleDomain.all(), null), new RemoteSourceNode(new PlanNodeId("remote_id"), new PlanFragmentId("plan_fragment_id"), ImmutableList.of()), ImmutableList.of(), Optional.<Symbol>empty(), Optional.<Symbol>empty()), ImmutableMap.<Symbol, Type>of(symbol, VARCHAR), SOURCE_DISTRIBUTION, tableScanNodeId, new PartitionFunctionBinding(SINGLE_DISTRIBUTION, ImmutableList.of(symbol), ImmutableList.of())); return new StageExecutionPlan(testFragment, Optional.of(new ConnectorAwareSplitSource(CONNECTOR_ID, TestingTransactionHandle.create(CONNECTOR_ID), splitSource)), ImmutableList.of()); }