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