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); } }
if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) { Entry<PlanNodeId, SplitSource> entry = Iterables.getOnlyElement(plan.getSplitSources().entrySet()); PlanNodeId planNodeId = entry.getKey(); SplitSource splitSource = entry.getValue(); Map<PlanNodeId, SplitSource> splitSources = plan.getSplitSources(); if (!splitSources.isEmpty()) {
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); }
@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()); } }