use of com.google.api.services.dataflow.model.ParallelInstruction in project beam by apache.
the class MapTaskToNetworkFunctionTest method testMultipleOutput.
@Test
public void testMultipleOutput() {
// /---> WriteA
// Read ---> ParDo
// \---> WriteB
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();
// Read.out
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();
// ParDo.out1
writeAInstruction.setInput(createInstructionInput(1, 0));
ParallelInstruction writeA = createParallelInstruction("WriteA");
writeA.setWrite(writeAInstruction);
WriteInstruction writeBInstruction = new WriteInstruction();
// ParDo.out2
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());
}
use of com.google.api.services.dataflow.model.ParallelInstruction in project beam by apache.
the class MapTaskToNetworkFunctionTest method testRead.
@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());
}
use of com.google.api.services.dataflow.model.ParallelInstruction in project beam by apache.
the class MapTaskToNetworkFunctionTest method testParallelEdgeFlatten.
@Test
public void testParallelEdgeFlatten() {
// /---\
// Read --> Read.out --> Flatten
// \---/
InstructionOutput readOutput = createInstructionOutput("Read.out");
ParallelInstruction read = createParallelInstruction("Read", readOutput);
read.setRead(new ReadInstruction());
FlattenInstruction flattenInstruction = new FlattenInstruction();
flattenInstruction.setInputs(ImmutableList.of(// Read.out
createInstructionInput(0, 0), // Read.out
createInstructionInput(0, 0), // Read.out
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);
// Assert that the three parallel edges are maintained
assertEquals(3, network.edgesConnecting(readOutputNode, flattenNode).size());
InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
}
use of com.google.api.services.dataflow.model.ParallelInstruction in project beam by apache.
the class NodesTest method testParallelInstructionNode.
@Test
public void testParallelInstructionNode() {
ParallelInstruction param = new ParallelInstruction();
Nodes.ExecutionLocation location = Nodes.ExecutionLocation.UNKNOWN;
assertSame(param, ParallelInstructionNode.create(param, location).getParallelInstruction());
assertSame(Nodes.ExecutionLocation.UNKNOWN, ParallelInstructionNode.create(param, location).getExecutionLocation());
assertNotEquals(ParallelInstructionNode.create(param, location), ParallelInstructionNode.create(param, location));
}
use of com.google.api.services.dataflow.model.ParallelInstruction in project beam by apache.
the class ReplacePgbkWithPrecombineFunctionTest method testNormalPgbkIsNotReplaced.
@Test
public void testNormalPgbkIsNotReplaced() throws Exception {
// Network:
// out1 --> pgbk --> out2
Node out1 = createInstructionOutputNode("out1");
Node pgbk = createPrecombinePgbkNode("pgbk", null);
Node out2 = createInstructionOutputNode("out2");
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(out1);
network.addNode(pgbk);
network.addNode(out2);
network.addEdge(out1, pgbk, DefaultEdge.create());
network.addEdge(pgbk, out2, DefaultEdge.create());
Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
network = new ReplacePgbkWithPrecombineFunction().apply(network);
// Assert that network is unchanged (identical paths).
List<List<Node>> oldPaths = Networks.allPathsFromRootsToLeaves(inputNetwork);
List<List<Node>> newPaths = Networks.allPathsFromRootsToLeaves(network);
assertThat(oldPaths, containsInAnyOrder(newPaths.toArray()));
// Assert that the pgbk node is still present and unchanged.
for (Node node : network.nodes()) {
if (node instanceof ParallelInstructionNode) {
ParallelInstructionNode newNode = (ParallelInstructionNode) node;
ParallelInstruction parallelInstruction = newNode.getParallelInstruction();
assertEquals(parallelInstruction, ((ParallelInstructionNode) pgbk).getParallelInstruction());
break;
}
}
}
Aggregations