public PhasedExecutionSchedule(Collection<SqlStageExecution> stages) { List<Set<PlanFragmentId>> phases = extractPhases(stages.stream().map(SqlStageExecution::getFragment).collect(toImmutableList())); Map<PlanFragmentId, SqlStageExecution> stagesByFragmentId = stages.stream().collect(toImmutableMap(stage -> stage.getFragment().getId(), identity())); // create a mutable list of mutable sets of stages, so we can remove completed stages schedulePhases = new ArrayList<>(); for (Set<PlanFragmentId> phase : phases) { schedulePhases.add(phase.stream() .map(stagesByFragmentId::get) .collect(Collectors.toCollection(HashSet::new))); } }
public PhasedExecutionSchedule(Collection<SqlStageExecution> stages) { List<Set<PlanFragmentId>> phases = extractPhases(stages.stream().map(SqlStageExecution::getFragment).collect(toImmutableList())); Map<PlanFragmentId, SqlStageExecution> stagesByFragmentId = stages.stream().collect(toImmutableMap(stage -> stage.getFragment().getId(), identity())); // create a mutable list of mutable sets of stages, so we can remove completed stages schedulePhases = new ArrayList<>(); for (Set<PlanFragmentId> phase : phases) { schedulePhases.add(phase.stream() .map(stagesByFragmentId::get) .collect(Collectors.toCollection(HashSet::new))); } }
private static void printFragmentNodes(StringBuilder output, PlanFragment fragment, PlanNodeIdGenerator idGenerator) { String clusterId = "cluster_" + fragment.getId(); output.append("subgraph ") .append(clusterId) .append(" {") .append('\n'); output.append(format("label = \"%s\"", fragment.getPartitioning())) .append('\n'); PlanNode plan = fragment.getRoot(); plan.accept(new NodePrinter(output, idGenerator), null); output.append("}") .append('\n'); }
@VisibleForTesting static List<PlanFragmentId> getPreferredScheduleOrder(Collection<PlanFragment> fragments) { // determine output fragment Set<PlanFragmentId> remoteSources = fragments.stream() .map(PlanFragment::getRemoteSourceNodes) .flatMap(Collection::stream) .map(RemoteSourceNode::getSourceFragmentIds) .flatMap(Collection::stream) .collect(toImmutableSet()); Set<PlanFragment> rootFragments = fragments.stream() .filter(fragment -> !remoteSources.contains(fragment.getId())) .collect(toImmutableSet()); checkArgument(rootFragments.size() == 1, "Expected one root fragment, but found: " + rootFragments); Visitor visitor = new Visitor(fragments); visitor.processFragment(getOnlyElement(rootFragments).getId()); return visitor.getSchedulerOrder(); }
private static void printFragmentNodes(StringBuilder output, PlanFragment fragment, PlanNodeIdGenerator idGenerator) { String clusterId = "cluster_" + fragment.getId(); output.append("subgraph ") .append(clusterId) .append(" {") .append('\n'); output.append(format("label = \"%s\"", fragment.getPartitioning())) .append('\n'); PlanNode plan = fragment.getRoot(); plan.accept(new NodePrinter(output, idGenerator), null); output.append("}") .append('\n'); }
@Test public void testJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment probeFragment = createTableScanPlanFragment("probe"); PlanFragment joinFragment = createJoinPlanFragment(INNER, "join", buildFragment, probeFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment, probeFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId()), ImmutableSet.of(buildFragment.getId()), ImmutableSet.of(probeFragment.getId()))); }
@Test public void testRightJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment probeFragment = createTableScanPlanFragment("probe"); PlanFragment joinFragment = createJoinPlanFragment(RIGHT, "join", buildFragment, probeFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment, probeFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId()), ImmutableSet.of(buildFragment.getId()), ImmutableSet.of(probeFragment.getId()))); }
@Test public void testBroadcastJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment joinFragment = createBroadcastJoinPlanFragment("join", buildFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId(), buildFragment.getId()))); }
@Test public void testRightJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment probeFragment = createTableScanPlanFragment("probe"); PlanFragment joinFragment = createJoinPlanFragment(RIGHT, "join", buildFragment, probeFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment, probeFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId()), ImmutableSet.of(buildFragment.getId()), ImmutableSet.of(probeFragment.getId()))); }
@Test public void testJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment probeFragment = createTableScanPlanFragment("probe"); PlanFragment joinFragment = createJoinPlanFragment(INNER, "join", buildFragment, probeFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment, probeFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId()), ImmutableSet.of(buildFragment.getId()), ImmutableSet.of(probeFragment.getId()))); }
@Test public void testUnion() { PlanFragment aFragment = createTableScanPlanFragment("a"); PlanFragment bFragment = createTableScanPlanFragment("b"); PlanFragment cFragment = createTableScanPlanFragment("c"); PlanFragment unionFragment = createUnionPlanFragment("union", aFragment, bFragment, cFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(aFragment, bFragment, cFragment, unionFragment)); assertEquals(phases, ImmutableList.of( ImmutableSet.of(unionFragment.getId()), ImmutableSet.of(aFragment.getId()), ImmutableSet.of(bFragment.getId()), ImmutableSet.of(cFragment.getId()))); }
@Test public void testExchange() { PlanFragment aFragment = createTableScanPlanFragment("a"); PlanFragment bFragment = createTableScanPlanFragment("b"); PlanFragment cFragment = createTableScanPlanFragment("c"); PlanFragment exchangeFragment = createExchangePlanFragment("exchange", aFragment, bFragment, cFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(aFragment, bFragment, cFragment, exchangeFragment)); assertEquals(phases, ImmutableList.of( ImmutableSet.of(exchangeFragment.getId()), ImmutableSet.of(aFragment.getId()), ImmutableSet.of(bFragment.getId()), ImmutableSet.of(cFragment.getId()))); }
@Test public void testBroadcastJoin() { PlanFragment buildFragment = createTableScanPlanFragment("build"); PlanFragment joinFragment = createBroadcastJoinPlanFragment("join", buildFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(joinFragment, buildFragment)); assertEquals(phases, ImmutableList.of(ImmutableSet.of(joinFragment.getId(), buildFragment.getId()))); }
@Test public void testExchange() { PlanFragment aFragment = createTableScanPlanFragment("a"); PlanFragment bFragment = createTableScanPlanFragment("b"); PlanFragment cFragment = createTableScanPlanFragment("c"); PlanFragment exchangeFragment = createExchangePlanFragment("exchange", aFragment, bFragment, cFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(aFragment, bFragment, cFragment, exchangeFragment)); assertEquals(phases, ImmutableList.of( ImmutableSet.of(exchangeFragment.getId()), ImmutableSet.of(aFragment.getId()), ImmutableSet.of(bFragment.getId()), ImmutableSet.of(cFragment.getId()))); }
@Test public void testUnion() { PlanFragment aFragment = createTableScanPlanFragment("a"); PlanFragment bFragment = createTableScanPlanFragment("b"); PlanFragment cFragment = createTableScanPlanFragment("c"); PlanFragment unionFragment = createUnionPlanFragment("union", aFragment, bFragment, cFragment); List<Set<PlanFragmentId>> phases = PhasedExecutionSchedule.extractPhases(ImmutableList.of(aFragment, bFragment, cFragment, unionFragment)); assertEquals(phases, ImmutableList.of( ImmutableSet.of(unionFragment.getId()), ImmutableSet.of(aFragment.getId()), ImmutableSet.of(bFragment.getId()), ImmutableSet.of(cFragment.getId()))); }
private static PlanFragment createJoinPlanFragment(JoinNode.Type joinType, String name, PlanFragment buildFragment, PlanFragment probeFragment) { RemoteSourceNode probe = new RemoteSourceNode(new PlanNodeId("probe_id"), probeFragment.getId(), ImmutableList.of(), Optional.empty(), REPARTITION); RemoteSourceNode build = new RemoteSourceNode(new PlanNodeId("build_id"), buildFragment.getId(), ImmutableList.of(), Optional.empty(), REPARTITION); PlanNode planNode = new JoinNode( new PlanNodeId(name + "_id"), joinType, probe, build, ImmutableList.of(), ImmutableList.<Symbol>builder() .addAll(probe.getOutputSymbols()) .addAll(build.getOutputSymbols()) .build(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); return createFragment(planNode); }
private static PlanFragment createUnionPlanFragment(String name, PlanFragment... fragments) { PlanNode planNode = new UnionNode( new PlanNodeId(name + "_id"), Stream.of(fragments) .map(fragment -> new RemoteSourceNode(new PlanNodeId(fragment.getId().toString()), fragment.getId(), fragment.getPartitioningScheme().getOutputLayout(), Optional.empty(), REPARTITION)) .collect(toImmutableList()), ImmutableListMultimap.of(), ImmutableList.of()); return createFragment(planNode); }
private static PlanFragment createUnionPlanFragment(String name, PlanFragment... fragments) { PlanNode planNode = new UnionNode( new PlanNodeId(name + "_id"), Stream.of(fragments) .map(fragment -> new RemoteSourceNode(new PlanNodeId(fragment.getId().toString()), fragment.getId(), fragment.getPartitioningScheme().getOutputLayout(), Optional.empty(), REPARTITION)) .collect(toImmutableList()), ImmutableListMultimap.of(), ImmutableList.of()); return createFragment(planNode); }