@Override public void translate(Pipeline pipeline) { // Ensure all outputs of all reads are consumed. UnconsumedReads.ensureAllReadsConsumed(pipeline); super.translate(pipeline); }
public static void ensureAllReadsConsumed(Pipeline pipeline) { final Set<PCollection<?>> unconsumed = new HashSet<>(); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { unconsumed.removeAll(node.getInputs().values()); } @Override public void visitValue(PValue value, Node producer) { String urn = PTransformTranslation.urnForTransformOrNull(producer.getTransform()); if (PTransformTranslation.READ_TRANSFORM_URN.equals(urn)) { unconsumed.add((PCollection<?>) value); } } }); int i = 0; for (PCollection<?> unconsumedPCollection : unconsumed) { consume(unconsumedPCollection, i); i++; } }
@Override public void translate(Pipeline pipeline) { // Ensure all outputs of all reads are consumed. UnconsumedReads.ensureAllReadsConsumed(pipeline); super.translate(pipeline); }
@Override public void translate(Pipeline pipeline) { // Ensure all outputs of all reads are consumed. UnconsumedReads.ensureAllReadsConsumed(pipeline); super.translate(pipeline); }
@VisibleForTesting protected void replaceTransforms(Pipeline pipeline) { boolean streaming = options.isStreaming() || containsUnboundedPCollection(pipeline); // Ensure all outputs of all reads are consumed before potentially replacing any // Read PTransforms UnconsumedReads.ensureAllReadsConsumed(pipeline); pipeline.replaceAll(getOverrides(streaming)); }
@Test public void matcherProducesUnconsumedValueUnboundedRead() { Unbounded<Long> transform = Read.from(CountingSource.unbounded()); pipeline.apply(transform); UnconsumedReads.ensureAllReadsConsumed(pipeline); validateConsumed(); }
@Test public void matcherProducesUnconsumedValueBoundedRead() { Bounded<Long> transform = Read.from(CountingSource.upTo(20L)); pipeline.apply(transform); UnconsumedReads.ensureAllReadsConsumed(pipeline); validateConsumed(); }
@Test public void doesNotConsumeAlreadyConsumedRead() { Unbounded<Long> transform = Read.from(CountingSource.unbounded()); final PCollection<Long> output = pipeline.apply(transform); final Flatten.PCollections<Long> consumer = Flatten.pCollections(); PCollectionList.of(output).apply(consumer); UnconsumedReads.ensureAllReadsConsumed(pipeline); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { // The output should only be consumed by a single consumer if (node.getInputs().values().contains(output)) { assertThat(node.getTransform(), Matchers.is(consumer)); } } }); }