use of com.google.api.services.dataflow.model.ReadInstruction 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);
}
use of com.google.api.services.dataflow.model.ReadInstruction 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);
}
use of com.google.api.services.dataflow.model.ReadInstruction 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));
}
Aggregations