@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) { PBegin input = PBegin.in(transform.getPipeline()); BoundedSource<T> source; try { source = ReadTranslation.boundedSourceFromTransform(transform); } catch (IOException e) { throw new RuntimeException(e); } return PTransformReplacement.of(input, bounded(source)); }
public static PTransformOverride boundedOverride() { return PTransformOverride.of(boundedMatcher(), new BoundedOverrideFactory<>()); }
@Override public PipelineResult run(Pipeline pipeline) { pipeline.replaceAll(ImmutableList.of(JavaReadViaImpulse.boundedOverride()));
@Test @Category(NeedsRunner.class) public void testReadToImpulseOverride() { BoundedSource<Long> source = CountingSource.upTo(10L); // Use an explicit read transform to ensure the override is exercised. PCollection<Long> input = p.apply(Read.from(source)); PAssert.that(input).containsInAnyOrder(0L, 1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L); p.replaceAll(Collections.singletonList(JavaReadViaImpulse.boundedOverride())); p.traverseTopologically( new Pipeline.PipelineVisitor() { @Override public void enterPipeline(Pipeline p) {} @Override public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { assertNotReadTransform(node.getTransform()); return CompositeBehavior.ENTER_TRANSFORM; } @Override public void leaveCompositeTransform(TransformHierarchy.Node node) {} @Override public void visitPrimitiveTransform(TransformHierarchy.Node node) { assertNotReadTransform(node.getTransform()); } @Override public void visitValue(PValue value, TransformHierarchy.Node producer) {} @Override public void leavePipeline(Pipeline pipeline) {} }); p.run(); }
pipeline.replaceAll(Collections.singletonList(JavaReadViaImpulse.boundedOverride()));
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Read.Bounded<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), JavaReadViaImpulse.bounded(transform.getTransform().getSource())); }
p.replaceAll(Collections.singletonList(JavaReadViaImpulse.boundedOverride()));
@Test public void testOutputCoder() { p.enableAbandonedNodeEnforcement(false); BoundedSource<Integer> fixedCoderSource = new BigEndianIntegerSource(); assertThat( p.apply(JavaReadViaImpulse.bounded(fixedCoderSource)).getCoder(), equalTo(BigEndianIntegerCoder.of())); }
p.replaceAll(Collections.singletonList(JavaReadViaImpulse.boundedOverride()));
@Test @Category(NeedsRunner.class) public void testBoundedRead() { PCollection<Long> read = p.apply(JavaReadViaImpulse.bounded(CountingSource.upTo(10L))); PAssert.that(read).containsInAnyOrder(0L, 9L, 8L, 1L, 2L, 7L, 6L, 3L, 4L, 5L); p.run(); }