private DataflowPipeline buildPipeline(DataflowPipelineOptions options) {
DataflowPipeline p = DataflowPipeline.create(options);
p.apply(TextIO.Read.named("ReadMyFile").from("gs://bucket/object"))
.apply(TextIO.Write.named("WriteMyFile").to("gs://bucket/object"));
return p;
}
/** Returns a Step for a DoFn by creating and translating a pipeline. */
private static Step createPredefinedStep() throws Exception {
DataflowPipelineOptions options = buildPipelineOptions();
DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options);
DataflowPipeline pipeline = DataflowPipeline.create(options);
String stepName = "DoFn1";
pipeline
.apply(TextIO.Read.named("ReadMyFile").from("gs://bucket/in"))
.apply(ParDo.of(new NoOpFn()).named(stepName))
.apply(TextIO.Write.named("WriteMyFile").to("gs://bucket/out"));
Job job = translator.translate(pipeline, Collections.<DataflowPackage>emptyList()).getJob();
assertEquals(3, job.getSteps().size());
Step step = job.getSteps().get(1);
assertEquals(stepName, getString(step.getProperties(), PropertyNames.USER_NAME));
return step;
}
@Test
public void testPredefinedAddStep() throws Exception {
DataflowPipelineOptions options = buildPipelineOptions();
DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options);
DataflowPipelineTranslator.registerTransformTranslator(
EmbeddedTransform.class, new EmbeddedTranslator());
// Create a predefined step using another pipeline
Step predefinedStep = createPredefinedStep();
// Create a pipeline that the predefined step will be embedded into
DataflowPipeline pipeline = DataflowPipeline.create(options);
pipeline
.apply(TextIO.Read.named("ReadMyFile").from("gs://bucket/in"))
.apply(ParDo.of(new NoOpFn()))
.apply(new EmbeddedTransform(predefinedStep.clone()))
.apply(TextIO.Write.named("WriteMyFile").to("gs://bucket/out"));
Job job = translator.translate(pipeline, Collections.<DataflowPackage>emptyList()).getJob();
List<Step> steps = job.getSteps();
assertEquals(4, steps.size());
// The input to the embedded step should match the output of the step before
Map<String, Object> step1Out = getOutputPortReference(steps.get(1));
Map<String, Object> step2In =
getDictionary(steps.get(2).getProperties(), PropertyNames.PARALLEL_INPUT);
assertEquals(step1Out, step2In);
// The output from the embedded step should match the input of the step after
Map<String, Object> step2Out = getOutputPortReference(steps.get(2));
Map<String, Object> step3In =
getDictionary(steps.get(3).getProperties(), PropertyNames.PARALLEL_INPUT);
assertEquals(step2Out, step3In);
// The step should not have been modified other than remapping the input
Step predefinedStepClone = predefinedStep.clone();
Step embeddedStepClone = steps.get(2).clone();
predefinedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT);
embeddedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT);
assertEquals(predefinedStepClone, embeddedStepClone);
}
