Examples with AppliedPTransform org.apache.beam.sdk.runners.AppliedPTransform used on opensource projects

Example 21 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class WatermarkManagerTest method getWatermarkMultiIdenticalInput.

/**
   * Demonstrates that getWatermark for a transform that takes multiple inputs is held to the
   * minimum watermark across all of its inputs.
   */
@Test
public void getWatermarkMultiIdenticalInput() {
    PCollection<Integer> created = p.apply(Create.of(1, 2, 3));
    PCollection<Integer> multiConsumer = PCollectionList.of(created).and(created).apply(Flatten.<Integer>pCollections());
    DirectGraphVisitor graphVisitor = new DirectGraphVisitor();
    p.traverseTopologically(graphVisitor);
    DirectGraph graph = graphVisitor.getGraph();
    AppliedPTransform<?, ?, ?> theFlatten = graph.getProducer(multiConsumer);
    WatermarkManager tstMgr = WatermarkManager.create(clock, graph);
    CommittedBundle<Void> root = bundleFactory.<Void>createRootBundle().add(WindowedValue.<Void>valueInGlobalWindow(null)).commit(clock.now());
    CommittedBundle<Integer> createBundle = bundleFactory.createBundle(created).add(WindowedValue.timestampedValueInGlobalWindow(1, new Instant(33536))).commit(clock.now());
    Map<AppliedPTransform<?, ?, ?>, Collection<CommittedBundle<?>>> initialInputs = ImmutableMap.<AppliedPTransform<?, ?, ?>, Collection<CommittedBundle<?>>>builder().put(graph.getProducer(created), Collections.<CommittedBundle<?>>singleton(root)).build();
    tstMgr.initialize(initialInputs);
    tstMgr.updateWatermarks(root, TimerUpdate.empty(), CommittedResult.create(StepTransformResult.withoutHold(graph.getProducer(created)).build(), root.withElements(Collections.<WindowedValue<Void>>emptyList()), Collections.singleton(createBundle), EnumSet.allOf(OutputType.class)), BoundedWindow.TIMESTAMP_MAX_VALUE);
    tstMgr.refreshAll();
    TransformWatermarks flattenWms = tstMgr.getWatermarks(theFlatten);
    assertThat(flattenWms.getInputWatermark(), equalTo(new Instant(33536)));
    tstMgr.updateWatermarks(createBundle, TimerUpdate.empty(), CommittedResult.create(StepTransformResult.withoutHold(theFlatten).build(), createBundle.withElements(Collections.<WindowedValue<Integer>>emptyList()), Collections.<CommittedBundle<?>>emptyList(), EnumSet.allOf(OutputType.class)), BoundedWindow.TIMESTAMP_MAX_VALUE);
    tstMgr.refreshAll();
    assertThat(flattenWms.getInputWatermark(), equalTo(new Instant(33536)));
    tstMgr.updateWatermarks(createBundle, TimerUpdate.empty(), CommittedResult.create(StepTransformResult.withoutHold(theFlatten).build(), createBundle.withElements(Collections.<WindowedValue<Integer>>emptyList()), Collections.<CommittedBundle<?>>emptyList(), EnumSet.allOf(OutputType.class)), BoundedWindow.TIMESTAMP_MAX_VALUE);
    tstMgr.refreshAll();
    assertThat(flattenWms.getInputWatermark(), equalTo(BoundedWindow.TIMESTAMP_MAX_VALUE));
}