public StageExecutionPlan withBucketToPartition(Optional<int[]> bucketToPartition) { return new StageExecutionPlan(fragment.withBucketToPartition(bucketToPartition), splitSources, subStages); }
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); } }
new AtomicInteger(), locationFactory, plan.withBucketToPartition(Optional.of(new int[1])), nodeScheduler, remoteTaskFactory,
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().getStageExecutionStrategy().isAnyScanGroupedExecution()); 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().getStageExecutionStrategy().isAnyScanGroupedExecution(); if (groupedExecutionForStage) { connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle); if (plan.getSubStages().isEmpty()) {
stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes());
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); }
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; }
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()) {
stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes());
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); }
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; }
public StageExecutionPlan withBucketToPartition(Optional<int[]> bucketToPartition) { return new StageExecutionPlan(fragment.withBucketToPartition(bucketToPartition), splitSources, subStages); }
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); } }
@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,
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()); }
@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()); } }
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.getStageExecutionStrategy(), 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)),
StatsAndCosts.empty()); return new StageExecutionPlan( testFragment, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, TestingTransactionHandle.create(), splitSource)),