StagePartitioner implements Function<IRDAG, Map<IRVertex, Integer>> { @Override public Map<IRVertex, Integer> apply(final IRDAG irDAG) { final Map<IRVertex, Integer> vertexToStageIdMap = new HashMap<>(); irDAG.topologicalDo(irVertex -> { if (vertexToStageIdMap.get(irVertex) == null) { vertexToStageIdMap.put(irVertex, nextStageIndex.getValue()); nextStageIndex.increment(); } final int stageId = vertexToStageIdMap.get(irVertex); for (final IREdge edge : irDAG.getOutgoingEdgesOf(irVertex)) { final IRVertex connectedIRVertex = edge.getDst(); if (vertexToStageIdMap.containsKey(connectedIRVertex)) { continue; } if (testMergeability(edge, irDAG)) { vertexToStageIdMap.put(connectedIRVertex, stageId); } else { vertexToStageIdMap.put(connectedIRVertex, nextStageIndex.getValue()); nextStageIndex.increment(); } } }); return vertexToStageIdMap; }  void addIgnoredPropertyKey(final Class<? extends VertexExecutionProperty> ignoredPropertyKey); @Override Map<IRVertex, Integer> apply(final IRDAG irDAG); Set<VertexExecutionProperty> getStageProperties(final IRVertex vertex);  }
@Test public void testLinear() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(5, 0, Collections.emptyList()); final IRVertex v1 = newVertex(5, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertEquals(0, (int) partitioning.get(v0)); assertEquals(0, (int) partitioning.get(v1)); }
@Test public void testSplitByParallelism() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(5, 0, Collections.emptyList()); final IRVertex v1 = newVertex(1, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertNotEquals(partitioning.get(v0), partitioning.get(v1)); }
@Test public void testSplitByScheduleGroup() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(1, 0, Collections.emptyList()); final IRVertex v1 = newVertex(1, 1, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertNotEquals(partitioning.get(v0), partitioning.get(v1)); }
@Test public void testSplitByShuffle() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(1, 0, Collections.emptyList()); final IRVertex v1 = newVertex(1, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.SHUFFLE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertNotEquals(partitioning.get(v0), partitioning.get(v1)); }
@Test public void testSplitByOtherProperty() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(1, 0, Arrays.asList(ResourcePriorityProperty.of(ResourcePriorityProperty.RESERVED))); final IRVertex v1 = newVertex(1, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertNotEquals(partitioning.get(v0), partitioning.get(v1)); }
@Test public void testNotSplitByIgnoredProperty() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(1, 0, Arrays.asList(IgnoreSchedulingTempDataReceiverProperty.of())); final IRVertex v1 = newVertex(1, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v1)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertEquals(0, (int) partitioning.get(v0)); assertEquals(0, (int) partitioning.get(v1)); }
@Test public void testJoin() { final DAGBuilder<IRVertex, IREdge> dagBuilder = new DAGBuilder<>(); final IRVertex v0 = newVertex(5, 0, Collections.emptyList()); final IRVertex v1 = newVertex(5, 0, Collections.emptyList()); final IRVertex v2 = newVertex(5, 0, Collections.emptyList()); dagBuilder.addVertex(v0); dagBuilder.addVertex(v1); dagBuilder.addVertex(v2); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v0, v2)); dagBuilder.connectVertices(new IREdge(CommunicationPatternProperty.Value.ONE_TO_ONE, v1, v2)); final Map<IRVertex, Integer> partitioning = stagePartitioner.apply(new IRDAG(dagBuilder.buildWithoutSourceSinkCheck())); assertNotEquals(partitioning.get(v0), partitioning.get(v1)); assertNotEquals(partitioning.get(v1), partitioning.get(v2)); assertNotEquals(partitioning.get(v2), partitioning.get(v0)); }