Search in sources :

Example 36 with InstructionOutput

use of com.google.api.services.dataflow.model.InstructionOutput in project beam by apache.

the class RemoveFlattenInstructionsFunctionTest method testRemoveFlatten.

@Test
public void testRemoveFlatten() {
    Node a = ParallelInstructionNode.create(new ParallelInstruction().setName("A"), Nodes.ExecutionLocation.UNKNOWN);
    Node aPCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out"), PCOLLECTION_ID);
    Edge aOutput = DefaultEdge.create();
    Node b = ParallelInstructionNode.create(new ParallelInstruction().setName("B"), Nodes.ExecutionLocation.UNKNOWN);
    Edge bOutput = DefaultEdge.create();
    Node bPCollection = InstructionOutputNode.create(new InstructionOutput().setName("B.out"), PCOLLECTION_ID);
    Node flatten = ParallelInstructionNode.create(new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN);
    Node flattenPCollection = InstructionOutputNode.create(new InstructionOutput().setName("Flatten.out"), PCOLLECTION_ID);
    Node c = ParallelInstructionNode.create(new ParallelInstruction().setName("C"), Nodes.ExecutionLocation.UNKNOWN);
    Edge cOutput = DefaultEdge.create();
    Node cPCollection = InstructionOutputNode.create(new InstructionOutput().setName("C.out"), PCOLLECTION_ID);
    // A --\
    // Flatten --> C
    // B --/
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(a);
    network.addNode(aPCollection);
    network.addNode(b);
    network.addNode(bPCollection);
    network.addNode(flatten);
    network.addNode(flattenPCollection);
    network.addNode(c);
    network.addNode(cPCollection);
    network.addEdge(a, aPCollection, aOutput);
    network.addEdge(aPCollection, flatten, DefaultEdge.create());
    network.addEdge(b, bPCollection, bOutput);
    network.addEdge(bPCollection, flatten, DefaultEdge.create());
    network.addEdge(flatten, flattenPCollection, DefaultEdge.create());
    network.addEdge(flattenPCollection, c, DefaultEdge.create());
    network.addEdge(c, cPCollection, cOutput);
    // A --\
    // C
    // B --/
    assertThatFlattenIsProperlyRemoved(network);
}
Also used : ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) Node(org.apache.beam.runners.dataflow.worker.graph.Nodes.Node) InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) ParallelInstructionNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode) InstructionOutput(com.google.api.services.dataflow.model.InstructionOutput) Edge(org.apache.beam.runners.dataflow.worker.graph.Edges.Edge) MultiOutputInfoEdge(org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge) DefaultEdge(org.apache.beam.runners.dataflow.worker.graph.Edges.DefaultEdge) FlattenInstruction(com.google.api.services.dataflow.model.FlattenInstruction) Test(org.junit.Test)

Aggregations

InstructionOutput (com.google.api.services.dataflow.model.InstructionOutput)36 ParallelInstruction (com.google.api.services.dataflow.model.ParallelInstruction)27 Test (org.junit.Test)20 InstructionOutputNode (org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode)19 DefaultEdge (org.apache.beam.runners.dataflow.worker.graph.Edges.DefaultEdge)17 Edge (org.apache.beam.runners.dataflow.worker.graph.Edges.Edge)17 Node (org.apache.beam.runners.dataflow.worker.graph.Nodes.Node)17 ParallelInstructionNode (org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode)17 MultiOutputInfoEdge (org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge)14 ReadInstruction (com.google.api.services.dataflow.model.ReadInstruction)13 CloudObject (org.apache.beam.runners.dataflow.util.CloudObject)12 ParDoInstruction (com.google.api.services.dataflow.model.ParDoInstruction)10 InstructionInput (com.google.api.services.dataflow.model.InstructionInput)9 FlattenInstruction (com.google.api.services.dataflow.model.FlattenInstruction)8 MultiOutputInfo (com.google.api.services.dataflow.model.MultiOutputInfo)8 MapTask (com.google.api.services.dataflow.model.MapTask)7 ArrayList (java.util.ArrayList)6 SdkComponents (org.apache.beam.runners.core.construction.SdkComponents)6 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)6 HashMap (java.util.HashMap)5