Examples with WatermarkEstimator org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator used on opensource projects

Example 1 with WatermarkEstimator

use of org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator in project beam by apache.

the class DoFnInvokersTest method testTruncateFnWithHasDefaultMethodsWhenBounded.

@Test
public void testTruncateFnWithHasDefaultMethodsWhenBounded() throws Exception {
    class BoundedMockFn extends DoFn<String, String> {

        @ProcessElement
        public void processElement(ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {
        }

        @GetInitialRestriction
        public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) {
            return null;
        }

        @GetInitialWatermarkEstimatorState
        public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() {
            return null;
        }
    }
    BoundedMockFn fn = mock(BoundedMockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class));
    coderRegistry.registerCoderForClass(WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
    assertThat(invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
    assertThat(invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
    RestrictionTracker tracker = invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {

        @Override
        public Object restriction() {
            return new RestrictionWithBoundedDefaultTracker();
        }
    });
    assertThat(tracker, instanceOf(BoundedDefaultTracker.class));
    TruncateResult<?> result = invoker.invokeTruncateRestriction(new FakeArgumentProvider<String, String>() {

        @Override
        public RestrictionTracker restrictionTracker() {
            return tracker;
        }

        @Override
        public String element(DoFn<String, String> doFn) {
            return "blah";
        }

        @Override
        public Object restriction() {
            return "foo";
        }
    });
    assertEquals("foo", result.getTruncatedRestriction());
    assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
    assertThat(invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {

        @Override
        public Object watermarkEstimatorState() {
            return new WatermarkEstimatorStateWithDefaultWatermarkEstimator();
        }
    }), instanceOf(DefaultWatermarkEstimator.class));
}

Example 2 with WatermarkEstimator

use of org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator in project beam by apache.

the class DoFnInvokersTest method testTruncateFnWithHasDefaultMethodsWhenUnbounded.

@Test
public void testTruncateFnWithHasDefaultMethodsWhenUnbounded() throws Exception {
    class UnboundedMockFn extends DoFn<String, String> {

        @ProcessElement
        public void processElement(ProcessContext c, RestrictionTracker<RestrictionWithUnboundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {
        }

        @GetInitialRestriction
        public RestrictionWithUnboundedDefaultTracker getInitialRestriction(@Element String element) {
            return null;
        }

        @GetInitialWatermarkEstimatorState
        public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() {
            return null;
        }
    }
    UnboundedMockFn fn = mock(UnboundedMockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithUnboundedDefaultTracker.class, CoderForDefaultTracker.class));
    coderRegistry.registerCoderForClass(WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
    assertThat(invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
    assertThat(invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
    RestrictionTracker tracker = invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {

        @Override
        public Object restriction() {
            return new RestrictionWithUnboundedDefaultTracker();
        }
    });
    assertThat(tracker, instanceOf(UnboundedDefaultTracker.class));
    TruncateResult<?> result = invoker.invokeTruncateRestriction(new FakeArgumentProvider<String, String>() {

        @Override
        public RestrictionTracker restrictionTracker() {
            return tracker;
        }

        @Override
        public String element(DoFn<String, String> doFn) {
            return "blah";
        }

        @Override
        public Object restriction() {
            return "foo";
        }
    });
    assertNull(result);
    assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
    assertThat(invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {

        @Override
        public Object watermarkEstimatorState() {
            return new WatermarkEstimatorStateWithDefaultWatermarkEstimator();
        }
    }), instanceOf(DefaultWatermarkEstimator.class));
}

Example 3 with WatermarkEstimator

use of org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator in project beam by apache.

the class DoFnInvokersTest method testSplittableDoFnWithHasDefaultMethods.

@Test
public void testSplittableDoFnWithHasDefaultMethods() throws Exception {
    class MockFn extends DoFn<String, String> {

        @ProcessElement
        public void processElement(ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {
        }

        @GetInitialRestriction
        public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) {
            return null;
        }

        @GetInitialWatermarkEstimatorState
        public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() {
            return null;
        }
    }
    MockFn fn = mock(MockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    CoderRegistry coderRegistry = CoderRegistry.createDefault();
    coderRegistry.registerCoderProvider(CoderProviders.fromStaticMethods(RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class));
    coderRegistry.registerCoderForClass(WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
    assertThat(invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class));
    assertThat(invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
    invoker.invokeSplitRestriction(new FakeArgumentProvider<String, String>() {

        @Override
        public String element(DoFn<String, String> doFn) {
            return "blah";
        }

        @Override
        public Object restriction() {
            return "foo";
        }

        @Override
        public OutputReceiver<String> outputReceiver(DoFn<String, String> doFn) {
            return new DoFn.OutputReceiver<String>() {

                private boolean invoked;

                @Override
                public void output(String output) {
                    assertFalse(invoked);
                    invoked = true;
                    assertEquals("foo", output);
                }

                @Override
                public void outputWithTimestamp(String output, Instant instant) {
                    assertFalse(invoked);
                    invoked = true;
                    assertEquals("foo", output);
                }
            };
        }
    });
    assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
    assertThat(invoker.invokeNewTracker(new FakeArgumentProvider<String, String>() {

        @Override
        public Object restriction() {
            return new RestrictionWithBoundedDefaultTracker();
        }
    }), instanceOf(BoundedDefaultTracker.class));
    assertThat(invoker.invokeNewWatermarkEstimator(new FakeArgumentProvider<String, String>() {

        @Override
        public Object watermarkEstimatorState() {
            return new WatermarkEstimatorStateWithDefaultWatermarkEstimator();
        }
    }), instanceOf(DefaultWatermarkEstimator.class));
}