DirectGraphVisitor extends PipelineVisitor.Defaults { public DirectGraph getGraph() { checkState(finalized, "Can't get a graph before the Pipeline has been completely traversed"); return DirectGraph.create( producers, viewWriters, perElementConsumers, rootTransforms, stepNames); }  @Override CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node); @Override void leaveCompositeTransform(TransformHierarchy.Node node); @Override void visitPrimitiveTransform(TransformHierarchy.Node node); @Override void visitValue(PValue value, TransformHierarchy.Node producer); DirectGraph getGraph();  }
@Test public void getViewsReturnsViews() { PCollectionView<List<String>> listView = p.apply("listCreate", Create.of("foo", "bar")) .apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(DoFn<String, String>.ProcessContext c) throws Exception { c.output(Integer.toString(c.element().length())); } })) .apply(View.asList()); PCollectionView<Object> singletonView = p.apply("singletonCreate", Create.<Object>of(1, 2, 3)).apply(View.asSingleton()); p.replaceAll( DirectRunner.fromOptions(TestPipeline.testingPipelineOptions()) .defaultTransformOverrides()); p.traverseTopologically(visitor); assertThat(visitor.getGraph().getViews(), Matchers.containsInAnyOrder(listView, singletonView)); }
@Test public void getRootTransformsContainsRootTransforms() { PCollection<byte[]> impulse = p.apply(Impulse.create()); impulse.apply(WithKeys.of("abc")); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); assertThat(graph.getRootTransforms(), hasSize(1)); assertThat( graph.getRootTransforms(), Matchers.containsInAnyOrder(new Object[] {graph.getProducer(impulse)})); for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) { assertThat(root.getInputs().entrySet(), emptyIterable()); assertThat( Iterables.getOnlyElement(root.getOutputs().values()), Matchers.<POutput>isOneOf(impulse)); } }
@Test public void getRootTransformsContainsEmptyFlatten() { PCollections<String> flatten = Flatten.pCollections(); PCollectionList<String> emptyList = PCollectionList.empty(p); PCollection<String> empty = emptyList.apply(flatten); empty.setCoder(StringUtf8Coder.of()); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); assertThat( graph.getRootTransforms(), Matchers.containsInAnyOrder(new Object[] {graph.getProducer(empty)})); AppliedPTransform<?, ?, ?> onlyRoot = Iterables.getOnlyElement(graph.getRootTransforms()); assertThat((Object) onlyRoot.getTransform(), equalTo(flatten)); assertThat(onlyRoot.getInputs().entrySet(), emptyIterable()); assertThat(onlyRoot.getOutputs(), equalTo(empty.expand())); }
@Test public void getValueToConsumersSucceeds() { PCollection<String> created = p.apply(Create.of("1", "2", "3")); PCollection<String> transformed = created.apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(DoFn<String, String>.ProcessContext c) throws Exception { c.output(Integer.toString(c.element().length())); } })); PCollection<String> flattened = PCollectionList.of(created).and(transformed).apply(Flatten.pCollections()); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); AppliedPTransform<?, ?, ?> transformedProducer = graph.getProducer(transformed); AppliedPTransform<?, ?, ?> flattenedProducer = graph.getProducer(flattened); assertThat( graph.getPerElementConsumers(created), Matchers.containsInAnyOrder(new Object[] {transformedProducer, flattenedProducer})); assertThat( graph.getPerElementConsumers(transformed), Matchers.containsInAnyOrder(new Object[] {flattenedProducer})); assertThat(graph.getPerElementConsumers(flattened), emptyIterable()); }
@Test public void getValueToConsumersWithDuplicateInputSucceeds() { PCollection<String> created = p.apply(Create.of("1", "2", "3")); PCollection<String> flattened = PCollectionList.of(created).and(created).apply(Flatten.pCollections()); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); AppliedPTransform<?, ?, ?> flattenedProducer = graph.getProducer(flattened); assertThat( graph.getPerElementConsumers(created), Matchers.containsInAnyOrder(new Object[] {flattenedProducer, flattenedProducer})); assertThat(graph.getPerElementConsumers(flattened), emptyIterable()); }
@Test public void getStepNamesContainsAllTransforms() { PCollection<String> created = p.apply(Create.of("1", "2", "3")); PCollection<String> transformed = created.apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(DoFn<String, String>.ProcessContext c) throws Exception { c.output(Integer.toString(c.element().length())); } })); PDone finished = transformed.apply( new PTransform<PInput, PDone>() { @Override public PDone expand(PInput input) { return PDone.in(input.getPipeline()); } }); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); int createdStepIndex = Integer.parseInt(graph.getStepName(graph.getProducer(created)).substring(1)); assertThat( graph.getStepName(graph.getProducer(transformed)), equalTo("s" + (createdStepIndex + 1))); }
@Test public void getGraphWithoutVisitingThrows() { thrown.expect(IllegalStateException.class); thrown.expectMessage("completely traversed"); thrown.expectMessage("get a graph"); visitor.getGraph(); }
@Test public void getRootTransformsContainsRootTransforms() { PCollection<String> created = p.apply(Create.of("foo", "bar")); PCollection<Long> counted = p.apply(Read.from(CountingSource.upTo(1234L))); PCollection<Long> unCounted = p.apply(GenerateSequence.from(0)); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); assertThat(graph.getRootTransforms(), hasSize(3)); assertThat( graph.getRootTransforms(), Matchers.containsInAnyOrder( new Object[] { graph.getProducer(created), graph.getProducer(counted), graph.getProducer(unCounted) })); for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) { assertThat(root.getInputs().entrySet(), emptyIterable()); assertThat( Iterables.getOnlyElement(root.getOutputs().values()), Matchers.<POutput>isOneOf(created, counted, unCounted)); } }
@Test public void getStepNamesContainsAllTransforms() { PCollection<String> created = p.apply(Create.of("1", "2", "3")); PCollection<String> transformed = created.apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(DoFn<String, String>.ProcessContext c) throws Exception { c.output(Integer.toString(c.element().length())); } })); PDone finished = transformed.apply( new PTransform<PInput, PDone>() { @Override public PDone expand(PInput input) { return PDone.in(input.getPipeline()); } }); p.traverseTopologically(visitor); DirectGraph graph = visitor.getGraph(); assertThat(graph.getStepName(graph.getProducer(created)), equalTo("s0")); assertThat(graph.getStepName(graph.getProducer(transformed)), equalTo("s1")); }