MapTaskToNetworkFunction implements Function<MapTask, MutableNetwork<Node, Edge>> { @Override public MutableNetwork<Node, Edge> apply(MapTask mapTask) { List<ParallelInstruction> parallelInstructions = Apiary.listOrEmpty(mapTask.getInstructions()); MutableNetwork<Node, Edge> network = NetworkBuilder.directed() .allowsSelfLoops(false) .allowsParallelEdges(true) .expectedNodeCount(parallelInstructions.size() * 2) .build(); ParallelInstructionNode[] instructionNodes = new ParallelInstructionNode[parallelInstructions.size()]; InstructionOutputNode[][] outputNodes = new InstructionOutputNode[parallelInstructions.size()][]; for (int i = 0; i < parallelInstructions.size(); ++i) { List<InstructionOutput> outputs = Apiary.listOrEmpty(parallelInstructions.get(i).getOutputs()); outputNodes[i] = new InstructionOutputNode[outputs.size()]; JsonFactory factory = MoreObjects.firstNonNull(mapTask.getFactory(), Transport.getJsonFactory()); ParallelInstruction parallelInstruction = clone(factory, parallelInstructions.get(i)).setOutputs(null); ParallelInstructionNode instructionNode = ParallelInstructionNode.create(parallelInstruction, Nodes.ExecutionLocation.UNKNOWN); instructionNodes[i] = instructionNode; network.addNode(instructionNode); for (int j = 0; j < outputs.size(); ++j) { InstructionOutput instructionOutput = outputs.get(j); InstructionOutputNode outputNode = InstructionOutputNode.create( instructionOutput, "generatedPcollection" + this.idGenerator.getId()); network.addNode(outputNode); if (parallelInstruction.getParDo() != null) { network.addEdge( instructionNode, outputNode, MultiOutputInfoEdge.create( parallelInstruction.getParDo().getMultiOutputInfos().get(j))); } else { network.addEdge(instructionNode, outputNode, DefaultEdge.create()); } outputNodes[i][j] = outputNode; } } for (ParallelInstructionNode instructionNode : instructionNodes) { ParallelInstruction parallelInstruction = instructionNode.getParallelInstruction(); if (parallelInstruction.getFlatten() != null) { for (InstructionInput input : Apiary.listOrEmpty(parallelInstruction.getFlatten().getInputs())) { attachInput(input, network, instructionNode, outputNodes); } } else if (parallelInstruction.getParDo() != null) { attachInput( parallelInstruction.getParDo().getInput(), network, instructionNode, outputNodes); } else if (parallelInstruction.getPartialGroupByKey() != null) { attachInput( parallelInstruction.getPartialGroupByKey().getInput(), network, instructionNode, outputNodes); } else if (parallelInstruction.getRead() != null) { } else if (parallelInstruction.getWrite() != null) { attachInput( parallelInstruction.getWrite().getInput(), network, instructionNode, outputNodes); } else { throw new IllegalArgumentException( String.format( "Unknown type of instruction %s for map task %s", parallelInstruction, mapTask)); } } return network; } MapTaskToNetworkFunction(IdGenerator idGenerator); @Override MutableNetwork<Node, Edge> apply(MapTask mapTask);  }
@Test public void testEmptyMapTask() { Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(new MapTask()); assertTrue(network.isDirected()); assertTrue(network.allowsParallelEdges()); assertFalse(network.allowsSelfLoops()); assertThat(network.nodes(), emptyCollectionOf(Node.class)); }
@Test public void testRead() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(2, network.nodes().size()); assertEquals(1, network.edges().size()); ParallelInstructionNode readNode = get(network, read); InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode); assertEquals(readOutput, readOutputNode.getInstructionOutput()); }
@Test public void testParDo() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); MultiOutputInfo parDoMultiOutput = createMultiOutputInfo("output"); ParDoInstruction parDoInstruction = new ParDoInstruction(); parDoInstruction.setInput(createInstructionInput(0, 0)); parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput)); InstructionOutput parDoOutput = createInstructionOutput("ParDo.out"); ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput); parDo.setParDo(parDoInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read, parDo)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(4, network.nodes().size()); assertEquals(3, network.edges().size()); ParallelInstructionNode readNode = get(network, read); InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode); assertEquals(readOutput, readOutputNode.getInstructionOutput()); ParallelInstructionNode parDoNode = getOnlySuccessor(network, readOutputNode); InstructionOutputNode parDoOutputNode = getOnlySuccessor(network, parDoNode); assertEquals(parDoOutput, parDoOutputNode.getInstructionOutput()); assertEquals( parDoMultiOutput, ((MultiOutputInfoEdge) Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutputNode))) .getMultiOutputInfo()); }
@Test public void testFlatten() { InstructionOutput readOutputA = createInstructionOutput("ReadA.out"); ParallelInstruction readA = createParallelInstruction("ReadA", readOutputA); readA.setRead(new ReadInstruction()); InstructionOutput readOutputB = createInstructionOutput("ReadB.out"); ParallelInstruction readB = createParallelInstruction("ReadB", readOutputB); readB.setRead(new ReadInstruction()); FlattenInstruction flattenInstruction = new FlattenInstruction(); flattenInstruction.setInputs( ImmutableList.of( createInstructionInput(0, 0), createInstructionInput(1, 0))); InstructionOutput flattenOutput = createInstructionOutput("Flatten.out"); ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput); flatten.setFlatten(flattenInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(readA, readB, flatten)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(6, network.nodes().size()); assertEquals(5, network.edges().size()); ParallelInstructionNode readANode = get(network, readA); InstructionOutputNode readOutputANode = getOnlySuccessor(network, readANode); assertEquals(readOutputA, readOutputANode.getInstructionOutput()); ParallelInstructionNode readBNode = get(network, readB); InstructionOutputNode readOutputBNode = getOnlySuccessor(network, readBNode); assertEquals(readOutputB, readOutputBNode.getInstructionOutput()); assertEquals(network.successors(readOutputANode), network.successors(readOutputBNode)); ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputANode); InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode); assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput()); }
@Test public void testParallelEdgeFlatten() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); FlattenInstruction flattenInstruction = new FlattenInstruction(); flattenInstruction.setInputs( ImmutableList.of( createInstructionInput(0, 0), createInstructionInput(0, 0), createInstructionInput(0, 0))); InstructionOutput flattenOutput = createInstructionOutput("Flatten.out"); ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput); flatten.setFlatten(flattenInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read, flatten)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(4, network.nodes().size()); assertEquals(5, network.edges().size()); ParallelInstructionNode readNode = get(network, read); InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode); assertEquals(readOutput, readOutputNode.getInstructionOutput()); ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputNode); assertEquals(3, network.edgesConnecting(readOutputNode, flattenNode).size()); InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode); assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput()); }
@Test public void testWrite() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); WriteInstruction writeInstruction = new WriteInstruction(); writeInstruction.setInput(createInstructionInput(0, 0)); ParallelInstruction write = createParallelInstruction("Write"); write.setWrite(writeInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read, write)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(3, network.nodes().size()); assertEquals(2, network.edges().size()); ParallelInstructionNode readNode = get(network, read); InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode); assertEquals(readOutput, readOutputNode.getInstructionOutput()); ParallelInstructionNode writeNode = getOnlySuccessor(network, readOutputNode); assertNotNull(writeNode); }
@Test public void testPartialGroupByKey() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction(); pgbkInstruction.setInput(createInstructionInput(0, 0)); InstructionOutput pgbkOutput = createInstructionOutput("PGBK.out"); ParallelInstruction pgbk = createParallelInstruction("PGBK", pgbkOutput); pgbk.setPartialGroupByKey(pgbkInstruction); WriteInstruction writeInstruction = new WriteInstruction(); writeInstruction.setInput(createInstructionInput(1, 0)); ParallelInstruction write = createParallelInstruction("Write"); write.setWrite(writeInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read, pgbk, write)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(5, network.nodes().size()); assertEquals(4, network.edges().size()); ParallelInstructionNode readNode = get(network, read); InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode); assertEquals(readOutput, readOutputNode.getInstructionOutput()); ParallelInstructionNode pgbkNode = getOnlySuccessor(network, readOutputNode); InstructionOutputNode pgbkOutputNode = getOnlySuccessor(network, pgbkNode); assertEquals(pgbkOutput, pgbkOutputNode.getInstructionOutput()); ParallelInstructionNode writeNode = getOnlySuccessor(network, pgbkOutputNode); assertNotNull(write); }
@Test public void testMultipleOutput() { InstructionOutput readOutput = createInstructionOutput("Read.out"); ParallelInstruction read = createParallelInstruction("Read", readOutput); read.setRead(new ReadInstruction()); MultiOutputInfo parDoMultiOutput1 = createMultiOutputInfo("output1"); MultiOutputInfo parDoMultiOutput2 = createMultiOutputInfo("output2"); ParDoInstruction parDoInstruction = new ParDoInstruction(); parDoInstruction.setInput(createInstructionInput(0, 0)); parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput1, parDoMultiOutput2)); InstructionOutput parDoOutput1 = createInstructionOutput("ParDo.out1"); InstructionOutput parDoOutput2 = createInstructionOutput("ParDo.out2"); ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput1, parDoOutput2); parDo.setParDo(parDoInstruction); WriteInstruction writeAInstruction = new WriteInstruction(); writeAInstruction.setInput(createInstructionInput(1, 0)); ParallelInstruction writeA = createParallelInstruction("WriteA"); writeA.setWrite(writeAInstruction); WriteInstruction writeBInstruction = new WriteInstruction(); writeBInstruction.setInput(createInstructionInput(1, 1)); ParallelInstruction writeB = createParallelInstruction("WriteB"); writeB.setWrite(writeBInstruction); MapTask mapTask = new MapTask(); mapTask.setInstructions(ImmutableList.of(read, parDo, writeA, writeB)); mapTask.setFactory(Transport.getJsonFactory()); Network<Node, Edge> network = new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask); assertNetworkProperties(network); assertEquals(7, network.nodes().size()); assertEquals(6, network.edges().size()); ParallelInstructionNode parDoNode = get(network, parDo); ParallelInstructionNode writeANode = get(network, writeA); ParallelInstructionNode writeBNode = get(network, writeB); InstructionOutputNode parDoOutput1Node = getOnlyPredecessor(network, writeANode); assertEquals(parDoOutput1, parDoOutput1Node.getInstructionOutput()); InstructionOutputNode parDoOutput2Node = getOnlyPredecessor(network, writeBNode); assertEquals(parDoOutput2, parDoOutput2Node.getInstructionOutput()); assertThat( network.successors(parDoNode), Matchers.<Node>containsInAnyOrder(parDoOutput1Node, parDoOutput2Node)); assertEquals( parDoMultiOutput1, ((MultiOutputInfoEdge) Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput1Node))) .getMultiOutputInfo()); assertEquals( parDoMultiOutput2, ((MultiOutputInfoEdge) Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput2Node))) .getMultiOutputInfo()); }