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

Example 1 with OffsetRangeTracker

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

the class SplittableParDoProcessFnTest method testUpdatesWatermark.

@Test
public void testUpdatesWatermark() throws Exception {
    DoFn<Instant, String> fn = new WatermarkUpdateFn();
    Instant base = Instant.now();
    ProcessFnTester<Instant, String, OffsetRange, OffsetRangeTracker> tester = new ProcessFnTester<>(base, fn, InstantCoder.of(), SerializableCoder.of(OffsetRange.class), 3, MAX_BUNDLE_DURATION);
    tester.startElement(base, new OffsetRange(0, 8));
    assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
    assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), hasItems("3", "4", "5"));
    assertEquals(base.plus(Duration.standardSeconds(5)), tester.getWatermarkHold());
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), hasItems("6", "7"));
    assertEquals(null, tester.getWatermarkHold());
}

Example 2 with OffsetRangeTracker

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

the class SplittableParDoProcessFnTest method testCheckpointsAfterDuration.

@Test
public void testCheckpointsAfterDuration() throws Exception {
    // Don't bound number of outputs.
    int max = Integer.MAX_VALUE;
    // But bound bundle duration - the bundle should terminate.
    Duration maxBundleDuration = Duration.standardSeconds(1);
    // Create an fn that attempts to 2x output more than checkpointing allows.
    DoFn<Integer, String> fn = new CounterFn();
    Instant base = Instant.now();
    int baseIndex = 42;
    ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester = new ProcessFnTester<>(base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class), max, maxBundleDuration);
    List<String> elements;
    tester.startElement(baseIndex, new OffsetRange(0, Long.MAX_VALUE));
    // Bundle should terminate, and should do at least some processing.
    elements = tester.takeOutputElements();
    assertFalse(elements.isEmpty());
    // Bundle should have run for at least the requested duration.
    assertThat(Instant.now().getMillis() - base.getMillis(), greaterThanOrEqualTo(maxBundleDuration.getMillis()));
}

Example 3 with OffsetRangeTracker

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

the class SplittableParDoProcessFnTest method testCheckpointsAfterNumOutputs.

@Test
public void testCheckpointsAfterNumOutputs() throws Exception {
    int max = 100;
    DoFn<Integer, String> fn = new CounterFn();
    Instant base = Instant.now();
    int baseIndex = 42;
    ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester = new ProcessFnTester<>(base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class), max, MAX_BUNDLE_DURATION);
    List<String> elements;
    // Create an fn that attempts to 2x output more than checkpointing allows.
    tester.startElement(baseIndex, new OffsetRange(0, 2 * max + max / 2));
    elements = tester.takeOutputElements();
    assertEquals(max, elements.size());
    // Should output the range [0, max)
    assertThat(elements, hasItem(String.valueOf(baseIndex)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + max - 1)));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    elements = tester.takeOutputElements();
    assertEquals(max, elements.size());
    // Should output the range [max, 2*max)
    assertThat(elements, hasItem(String.valueOf(baseIndex + max)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max - 1)));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    elements = tester.takeOutputElements();
    assertEquals(max / 2, elements.size());
    // Should output the range [2*max, 2*max + max/2)
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max)));
    assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max + max / 2 - 1)));
    assertThat(elements, not(hasItem((String.valueOf(baseIndex + 2 * max + max / 2)))));
}

Example 4 with OffsetRangeTracker

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method runTest.

private SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result<Integer, String, OffsetRange, OffsetRangeTracker> runTest(int count, Duration sleepPerElement) {
    SomeFn fn = new SomeFn(sleepPerElement);
    SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker> invoker = new OutputAndTimeBoundedSplittableProcessElementInvoker<>(fn, PipelineOptionsFactory.create(), new OutputWindowedValue<String>() {

        @Override
        public void outputWindowedValue(String output, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) {
        }

        @Override
        public <AdditionalOutputT> void outputWindowedValue(TupleTag<AdditionalOutputT> tag, AdditionalOutputT output, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) {
        }
    }, NullSideInputReader.empty(), Executors.newSingleThreadScheduledExecutor(), 1000, Duration.standardSeconds(3));
    return invoker.invokeProcessElement(DoFnInvokers.invokerFor(fn), WindowedValue.of(count, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING), new OffsetRangeTracker(new OffsetRange(0, count)));
}

Example 5 with OffsetRangeTracker

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method testInvokeProcessElementTimeBounded.

@Test
public void testInvokeProcessElementTimeBounded() throws Exception {
    SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result<Integer, String, OffsetRange, OffsetRangeTracker> res = runTest(10000, Duration.millis(100));
    OffsetRange residualRange = res.getResidualRestriction();
    // Should process ideally around 30 elements - but due to timing flakiness, we can't enforce
    // that precisely. Just test that it's not egregiously off.
    assertThat(residualRange.getFrom(), greaterThan(10L));
    assertThat(residualRange.getFrom(), lessThan(100L));
    assertEquals(10000, residualRange.getTo());
}