use of org.apache.beam.runners.dataflow.worker.graph.Edges.Edge in project beam by apache.
the class DeduceNodeLocationsFunctionTest method testSingleNodeWithRunnerParDo.
@Test
public void testSingleNodeWithRunnerParDo() throws Exception {
Node unknown = createParDoNode("Unknown", "RunnerDoFn");
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(unknown);
Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
network = new DeduceNodeLocationsFunction().apply(network);
assertThatNetworksAreIdentical(inputNetwork, network);
for (Node node : ImmutableList.copyOf(network.nodes())) {
assertNodesIdenticalExceptForExecutionLocation(unknown, node);
assertThatLocationIsProperlyDeduced(node, ExecutionLocation.RUNNER_HARNESS);
}
}
use of org.apache.beam.runners.dataflow.worker.graph.Edges.Edge in project beam by apache.
the class DeduceNodeLocationsFunctionTest method testGraphWithNonDeducibleNodes.
/**
* Tests that graphs with deducible and non-deducible nodes are maintained correctly.
*/
@Test
public void testGraphWithNonDeducibleNodes() throws Exception {
// A --> out1 --\
// --> Flatten --> D
// B --> out2 --/-->C
Node a = createReadNode("A", CUSTOM_SOURCE);
Node out1 = InstructionOutputNode.create(new InstructionOutput(), "fakeId");
Node b = createReadNode("B", RUNNER_SOURCE);
Node out2 = InstructionOutputNode.create(new InstructionOutput(), "fakeId");
Node c = createParDoNode("C", "RunnerDoFn");
Node flatten = ParallelInstructionNode.create(new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN);
Node d = createParDoNode("D", DO_FN);
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(a);
network.addNode(out1);
network.addNode(b);
network.addNode(out2);
network.addNode(c);
network.addNode(flatten);
network.addNode(d);
network.addEdge(a, out1, DefaultEdge.create());
network.addEdge(b, out2, DefaultEdge.create());
network.addEdge(out1, flatten, DefaultEdge.create());
network.addEdge(out2, flatten, DefaultEdge.create());
network.addEdge(out2, c, DefaultEdge.create());
network.addEdge(flatten, d, DefaultEdge.create());
Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
network = new DeduceNodeLocationsFunction().apply(network);
assertThatNetworksAreIdentical(inputNetwork, network);
assertAllNodesDeducedExceptFlattens(network);
}
use of org.apache.beam.runners.dataflow.worker.graph.Edges.Edge in project beam by apache.
the class DeduceNodeLocationsFunctionTest method testSingleNodeWithRunnerRead.
@Test
public void testSingleNodeWithRunnerRead() throws Exception {
Node unknown = createReadNode("Unknown", RUNNER_SOURCE);
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(unknown);
Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
network = new DeduceNodeLocationsFunction().apply(network);
assertThatNetworksAreIdentical(inputNetwork, network);
for (Node node : ImmutableList.copyOf(network.nodes())) {
assertNodesIdenticalExceptForExecutionLocation(unknown, node);
assertThatLocationIsProperlyDeduced(node, ExecutionLocation.RUNNER_HARNESS);
}
}
use of org.apache.beam.runners.dataflow.worker.graph.Edges.Edge in project beam by apache.
the class InsertFetchAndFilterStreamingSideInputNodesTest method testSdkParDoWithoutSideInput.
@Test
public void testSdkParDoWithoutSideInput() throws Exception {
Pipeline p = Pipeline.create();
PCollection<String> pc = p.apply(Create.of("a", "b", "c"));
pc.apply(ParDo.of(new TestDoFn()));
RunnerApi.Pipeline pipeline = PipelineTranslation.toProto(p);
Node predecessor = createParDoNode("predecessor");
Node mainInput = InstructionOutputNode.create(new InstructionOutput(), "fakeId");
Node sideInputParDo = createParDoNode("noSideInput");
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(predecessor);
network.addNode(mainInput);
network.addNode(sideInputParDo);
network.addEdge(predecessor, mainInput, DefaultEdge.create());
network.addEdge(mainInput, sideInputParDo, DefaultEdge.create());
Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
network = InsertFetchAndFilterStreamingSideInputNodes.with(pipeline).forNetwork(network);
assertThatNetworksAreIdentical(inputNetwork, network);
}
use of org.apache.beam.runners.dataflow.worker.graph.Edges.Edge 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());
}
Aggregations