Examples with ParallelInstructionNode org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode used on opensource projects

Example 6 with ParallelInstructionNode

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.

the class CloneAmbiguousFlattensFunctionTest method testConnectedFlattens.

/**
 * Tests that multiple connected ambiguous flattens in a network all get cloned without leaving
 * any ambiguous flattens left in the network and without changing the number of paths.
 */
@Test
public void testConnectedFlattens() throws Exception {
    // sdk1+out ----\
    // ambig_flatten1+out --> sdk3+out
    // runner1+out -/                    \                     /-> sdk4+out
    // sdk2+out ----\                    /-> ambig_flatten3+out
    // ambig_flatten2+out --> runner3+out        \-> runner4+out
    // runner2+out -/
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    Node sdk1 = createSdkNode("sdk1");
    Node sdk2 = createSdkNode("sdk2");
    Node sdk3 = createSdkNode("sdk3");
    Node sdk4 = createSdkNode("sdk4");
    Node sdk1Out = createPCollection("sdk1.out");
    Node sdk2Out = createPCollection("sdk2.out");
    Node sdk3Out = createPCollection("sdk3.out");
    Node sdk4Out = createPCollection("sdk4.out");
    Node runner1 = createRunnerNode("runner1");
    Node runner2 = createRunnerNode("runner2");
    Node runner3 = createRunnerNode("runner3");
    Node runner4 = createRunnerNode("runner4");
    Node runner1Out = createPCollection("runner1.out");
    Node runner2Out = createPCollection("runner2.out");
    Node runner3Out = createPCollection("runner3.out");
    Node runner4Out = createPCollection("runner4.out");
    Node ambiguousFlatten1 = createFlatten("ambiguous_flatten1", ExecutionLocation.AMBIGUOUS);
    Node ambiguousFlatten2 = createFlatten("ambiguous_flatten2", ExecutionLocation.AMBIGUOUS);
    Node ambiguousFlatten3 = createFlatten("ambiguous_flatten3", ExecutionLocation.AMBIGUOUS);
    Node ambiguousFlatten1Out = createPCollection("ambiguous_flatten1.out");
    Node ambiguousFlatten2Out = createPCollection("ambiguous_flatten2.out");
    Node ambiguousFlatten3Out = createPCollection("ambiguous_flatten3.out");
    network.addNode(sdk1);
    network.addNode(sdk2);
    network.addNode(sdk3);
    network.addNode(sdk4);
    network.addNode(sdk1Out);
    network.addNode(sdk2Out);
    network.addNode(sdk3Out);
    network.addNode(sdk4Out);
    network.addNode(runner1);
    network.addNode(runner2);
    network.addNode(runner3);
    network.addNode(runner4);
    network.addNode(runner1Out);
    network.addNode(runner2Out);
    network.addNode(runner3Out);
    network.addNode(runner4Out);
    network.addNode(ambiguousFlatten1);
    network.addNode(ambiguousFlatten2);
    network.addNode(ambiguousFlatten3);
    network.addNode(ambiguousFlatten1Out);
    network.addNode(ambiguousFlatten2Out);
    network.addNode(ambiguousFlatten3Out);
    network.addEdge(sdk1, sdk1Out, DefaultEdge.create());
    network.addEdge(sdk2, sdk2Out, DefaultEdge.create());
    network.addEdge(sdk3, sdk3Out, DefaultEdge.create());
    network.addEdge(sdk4, sdk4Out, DefaultEdge.create());
    network.addEdge(runner1, runner1Out, DefaultEdge.create());
    network.addEdge(runner2, runner2Out, DefaultEdge.create());
    network.addEdge(runner3, runner3Out, DefaultEdge.create());
    network.addEdge(runner4, runner4Out, DefaultEdge.create());
    network.addEdge(ambiguousFlatten1, ambiguousFlatten1Out, DefaultEdge.create());
    network.addEdge(ambiguousFlatten2, ambiguousFlatten2Out, DefaultEdge.create());
    network.addEdge(ambiguousFlatten3, ambiguousFlatten3Out, DefaultEdge.create());
    network.addEdge(sdk1Out, ambiguousFlatten1, DefaultEdge.create());
    network.addEdge(runner1Out, ambiguousFlatten1, DefaultEdge.create());
    network.addEdge(sdk2Out, ambiguousFlatten2, DefaultEdge.create());
    network.addEdge(runner2Out, ambiguousFlatten2, DefaultEdge.create());
    network.addEdge(ambiguousFlatten1Out, sdk3, DefaultEdge.create());
    network.addEdge(ambiguousFlatten1Out, ambiguousFlatten3, DefaultEdge.create());
    network.addEdge(ambiguousFlatten2Out, ambiguousFlatten3, DefaultEdge.create());
    network.addEdge(ambiguousFlatten2Out, runner3, DefaultEdge.create());
    network.addEdge(ambiguousFlatten3Out, sdk4, DefaultEdge.create());
    network.addEdge(ambiguousFlatten3Out, runner4, DefaultEdge.create());
    // Apply function and perform assertions
    List<List<Node>> originalPaths = Networks.allPathsFromRootsToLeaves(network);
    network = new CloneAmbiguousFlattensFunction().apply(network);
    for (Node node : network.nodes()) {
        if (node instanceof ParallelInstructionNode && ((ParallelInstructionNode) node).getParallelInstruction().getFlatten() != null) {
            ParallelInstructionNode castNode = ((ParallelInstructionNode) node);
            assertTrue("Ambiguous flatten not removed from network.", castNode.getExecutionLocation() != ExecutionLocation.AMBIGUOUS);
        }
    }
    assertEquals(originalPaths.size(), Networks.allPathsFromRootsToLeaves(network).size());
}

Example 7 with ParallelInstructionNode

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode in project beam by apache.

the class DeduceNodeLocationsFunctionTest method assertNodesIdenticalExceptForExecutionLocation.

/**
 * Asserts two nodes are identical except for ExecutionLocation, which can differ.
 */
private void assertNodesIdenticalExceptForExecutionLocation(Node expected, Node actual) {
    assertThat(expected, instanceOf(ParallelInstructionNode.class));
    assertThat(actual, instanceOf(ParallelInstructionNode.class));
    ParallelInstruction expectedContents = ((ParallelInstructionNode) expected).getParallelInstruction();
    ParallelInstruction actualContents = ((ParallelInstructionNode) actual).getParallelInstruction();
    assertEquals(expectedContents, actualContents);
}

Example 8 with ParallelInstructionNode

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode 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());
}

Example 9 with ParallelInstructionNode

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode 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());
}

Example 10 with ParallelInstructionNode

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode 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());
}