use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.
the class MapTaskToNetworkFunctionTest method testParDo.
@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();
// Read.out
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());
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.
the class MapTaskToNetworkFunctionTest method testFlatten.
@Test
public void testFlatten() {
// ReadA --\
// |--> Flatten
// ReadB --/
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(// ReadA.out
createInstructionInput(0, 0), // ReadB.out
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());
// Make sure the successors for both ReadA and ReadB output PCollections are the same
assertEquals(network.successors(readOutputANode), network.successors(readOutputBNode));
ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputANode);
InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.
the class MapTaskToNetworkFunctionTest method testPartialGroupByKey.
@Test
public void testPartialGroupByKey() {
// Read --> PGBK --> Write
InstructionOutput readOutput = createInstructionOutput("Read.out");
ParallelInstruction read = createParallelInstruction("Read", readOutput);
read.setRead(new ReadInstruction());
PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction();
// Read.out
pgbkInstruction.setInput(createInstructionInput(0, 0));
InstructionOutput pgbkOutput = createInstructionOutput("PGBK.out");
ParallelInstruction pgbk = createParallelInstruction("PGBK", pgbkOutput);
pgbk.setPartialGroupByKey(pgbkInstruction);
WriteInstruction writeInstruction = new WriteInstruction();
// PGBK.out
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());
getOnlySuccessor(network, pgbkOutputNode);
assertNotNull(write);
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.
the class MapTaskToNetworkFunctionTest method testWrite.
@Test
public void testWrite() {
InstructionOutput readOutput = createInstructionOutput("Read.out");
ParallelInstruction read = createParallelInstruction("Read", readOutput);
read.setRead(new ReadInstruction());
WriteInstruction writeInstruction = new WriteInstruction();
// Read.out
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);
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.
the class MapTaskToNetworkFunctionTest method assertNetworkProperties.
private static void assertNetworkProperties(Network<Node, Edge> network) {
assertTrue(network.isDirected());
assertFalse(network.allowsSelfLoops());
for (Node node : network.nodes()) {
if (node instanceof ParallelInstructionNode) {
assertNull(((ParallelInstructionNode) node).getParallelInstruction().getOutputs());
// except for ParDoInstructions
for (Node successor : network.successors(node)) {
assertThat(successor, instanceOf(InstructionOutputNode.class));
// Assert that all outgoing edges for a ParDo are MultiOutputInfoEdges
if (((ParallelInstructionNode) node).getParallelInstruction().getParDo() != null) {
for (Edge edge : network.edgesConnecting(node, successor)) {
assertThat(edge, instanceOf(MultiOutputInfoEdge.class));
}
} else {
for (Edge edge : network.edgesConnecting(node, successor)) {
assertThat(edge, instanceOf(DefaultEdge.class));
}
}
}
} else if (node instanceof InstructionOutputNode) {
assertThat(network.inDegree(node), greaterThanOrEqualTo(1));
// Validate that all successors are instructions with DefaultEdge outgoing edges
for (Node successor : network.successors(node)) {
assertThat(successor, instanceOf(ParallelInstructionNode.class));
for (Edge edge : network.edgesConnecting(node, successor)) {
assertThat(edge, instanceOf(DefaultEdge.class));
}
}
}
}
}
Aggregations