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

Example 21 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node 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 22 with Node

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

the class MapTaskToNetworkFunctionTest method testEmptyMapTask.

@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));
}

Example 23 with Node

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

Example 24 with Node

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

Example 25 with Node

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

the class CreateRegisterFnOperationFunction method rewireAcrossSdkRunnerPortNode.

/**
 * Rewires the given set of predecessors and successors across a gRPC port surrounded by output
 * nodes. Edges to the remaining successors are copied over to the new output node that is placed
 * before the port node. For example:
 *
 * <pre><code>
 * predecessors --> outputNode --> successors
 *                            \--> existingSuccessors
 * </pre></code> becomes:
 *
 * <pre><code>
 *
 *                                outputNode -------------------------------\
 *                                          \                                \
 *                                           |-> existingSuccessors           \
 *                                          /                                  \
 * predecessors --> newPredecessorOutputNode --> portNode --> portOutputNode --> successors}.
 * </code></pre>
 */
private Node rewireAcrossSdkRunnerPortNode(MutableNetwork<Node, Edge> network, InstructionOutputNode outputNode, Set<Node> predecessors, Set<Node> successors) {
    InstructionOutputNode newPredecessorOutputNode = InstructionOutputNode.create(outputNode.getInstructionOutput(), outputNode.getPcollectionId());
    InstructionOutputNode portOutputNode = InstructionOutputNode.create(outputNode.getInstructionOutput(), outputNode.getPcollectionId());
    Node portNode = portSupplier.get();
    network.addNode(newPredecessorOutputNode);
    network.addNode(portNode);
    for (Node predecessor : predecessors) {
        for (Edge edge : ImmutableList.copyOf(network.edgesConnecting(predecessor, outputNode))) {
            network.removeEdge(edge);
            network.addEdge(predecessor, newPredecessorOutputNode, edge);
        }
    }
    // Maintain edges for existing successors.
    List<Node> existingSuccessors = ImmutableList.copyOf(Sets.difference(network.successors(outputNode), successors));
    for (Node existingSuccessor : existingSuccessors) {
        List<Edge> existingSuccessorEdges = ImmutableList.copyOf(network.edgesConnecting(outputNode, existingSuccessor));
        for (Edge existingSuccessorEdge : existingSuccessorEdges) {
            network.addEdge(newPredecessorOutputNode, existingSuccessor, existingSuccessorEdge.clone());
        }
    }
    // Rewire the requested successors over the port node.
    network.addEdge(newPredecessorOutputNode, portNode, MultiOutputInfoEdge.create(new MultiOutputInfo().setTag(idGenerator.getId())));
    network.addEdge(portNode, portOutputNode, MultiOutputInfoEdge.create(new MultiOutputInfo().setTag(idGenerator.getId())));
    for (Node successor : successors) {
        for (Edge edge : ImmutableList.copyOf(network.edgesConnecting(outputNode, successor))) {
            network.addEdge(portOutputNode, successor, edge.clone());
        }
    }
    return portNode;
}