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

Example 16 with InstructionOutputNode

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

Example 17 with InstructionOutputNode

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

Example 18 with InstructionOutputNode

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

Example 19 with InstructionOutputNode

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

Example 20 with InstructionOutputNode

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

the class LengthPrefixUnknownCodersTest method testLengthPrefixAndReplaceForRunnerNetwork.

@Test
public void testLengthPrefixAndReplaceForRunnerNetwork() throws Exception {
    Node readNode = createReadNode("Read", "Source", windowedValueCoder);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out", windowedValueCoder);
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(readNode);
    network.addNode(readNodeOut);
    network.addEdge(readNode, readNodeOut, readNodeEdge);
    ParallelInstructionNode prefixedReadNode = createReadNode("Read", "Source", prefixedAndReplacedWindowedValueCoder);
    InstructionOutputNode prefixedReadNodeOut = createInstructionOutputNode("Read.out", prefixedAndReplacedWindowedValueCoder);
    MutableNetwork<Node, Edge> prefixedNetwork = andReplaceForRunnerNetwork(network);
    ImmutableSet.Builder<GenericJson> prefixedInstructions = ImmutableSet.builder();
    for (Node node : prefixedNetwork.nodes()) {
        if (node instanceof ParallelInstructionNode) {
            prefixedInstructions.add(((ParallelInstructionNode) node).getParallelInstruction());
        } else if (node instanceof InstructionOutputNode) {
            prefixedInstructions.add(((InstructionOutputNode) node).getInstructionOutput());
        }
    }
    assertThat(prefixedInstructions.build(), containsInAnyOrder(jsonOf(prefixedReadNodeOut.getInstructionOutput()), jsonOf(prefixedReadNode.getParallelInstruction())));
}