Examples with ParallelInstruction com.google.api.services.dataflow.model.ParallelInstruction used on opensource projects

Example 66 with ParallelInstruction

use of com.google.api.services.dataflow.model.ParallelInstruction 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 67 with ParallelInstruction

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

the class MapTaskToNetworkFunctionTest method createParallelInstruction.

private static ParallelInstruction createParallelInstruction(String name, InstructionOutput... outputs) {
    ParallelInstruction rval = new ParallelInstruction();
    rval.setName(name);
    rval.setOutputs(Arrays.asList(outputs));
    return rval;
}

Example 68 with ParallelInstruction

use of com.google.api.services.dataflow.model.ParallelInstruction 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 69 with ParallelInstruction

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

the class ReplacePgbkWithPrecombineFunctionTest method testPrecombinePgbkIsReplaced.

@Test
public void testPrecombinePgbkIsReplaced() throws Exception {
    // Network:
    // out1 --> precombine_pgbk --> out2
    Map<String, Object> valueCombiningFn = new HashMap<>();
    Node out1 = createInstructionOutputNode("out1");
    String pgbkName = "precombine_pgbk";
    Node precombinePgbk = createPrecombinePgbkNode(pgbkName, valueCombiningFn);
    Node out2 = createInstructionOutputNode("out2");
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(out1);
    network.addNode(precombinePgbk);
    network.addNode(out2);
    network.addEdge(out1, precombinePgbk, DefaultEdge.create());
    network.addEdge(precombinePgbk, out2, DefaultEdge.create());
    Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network);
    network = new ReplacePgbkWithPrecombineFunction().apply(network);
    // Assert that network has same structure (same number of nodes and paths).
    assertEquals(inputNetwork.nodes().size(), network.nodes().size());
    assertEquals(inputNetwork.edges().size(), network.edges().size());
    List<List<Node>> oldPaths = Networks.allPathsFromRootsToLeaves(inputNetwork);
    List<List<Node>> newPaths = Networks.allPathsFromRootsToLeaves(network);
    assertEquals(oldPaths.size(), newPaths.size());
    // Assert that the pgbk node has been replaced.
    for (Node node : network.nodes()) {
        if (node instanceof ParallelInstructionNode) {
            ParallelInstructionNode createdCombineNode = (ParallelInstructionNode) node;
            ParallelInstruction parallelInstruction = createdCombineNode.getParallelInstruction();
            assertEquals(parallelInstruction.getName(), pgbkName);
            assertNull(parallelInstruction.getPartialGroupByKey());
            assertNotNull(parallelInstruction.getParDo());
            ParDoInstruction parDoInstruction = parallelInstruction.getParDo();
            assertEquals(parDoInstruction.getUserFn(), valueCombiningFn);
            break;
        }
    }
}

Example 70 with ParallelInstruction

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

the class LengthPrefixUnknownCodersTest method testLengthPrefixReadInstructionCoder.

@Test
public void testLengthPrefixReadInstructionCoder() throws Exception {
    ReadInstruction readInstruction = new ReadInstruction();
    readInstruction.setSource(new Source().setCodec(CloudObjects.asCloudObject(windowedValueCoder, /*sdkComponents=*/
    null)));
    instruction.setRead(readInstruction);
    ParallelInstruction prefixedInstruction = forParallelInstruction(instruction, false);
    assertEqualsAsJson(CloudObjects.asCloudObject(prefixedWindowedValueCoder, /*sdkComponents=*/
    null), prefixedInstruction.getRead().getSource().getCodec());
    // Should not mutate the instruction.
    assertEqualsAsJson(readInstruction.getSource().getCodec(), CloudObjects.asCloudObject(windowedValueCoder, /*sdkComponents=*/
    null));
}