Examples with OffsetRange org.apache.beam.sdk.io.range.OffsetRange used on opensource projects

Example 21 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange 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(Integer.MAX_VALUE);
    Instant base = Instant.now();
    int baseIndex = 42;
    ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester = new ProcessFnTester<>(base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class), VoidCoder.of(), 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 22 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class SplittableParDoProcessFnTest method testResumeCarriesOverState.

@Test
public void testResumeCarriesOverState() throws Exception {
    DoFn<Integer, String> fn = new CounterFn(1);
    Instant base = Instant.now();
    dateTimeProvider.setDateTimeFixed(base.getMillis());
    ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester = new ProcessFnTester<>(base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class), VoidCoder.of(), MAX_OUTPUTS_PER_BUNDLE, MAX_BUNDLE_DURATION);
    tester.startElement(42, new OffsetRange(0, 3));
    assertThat(tester.takeOutputElements(), contains("42"));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), contains("43"));
    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    assertThat(tester.takeOutputElements(), contains("44"));
    // Should not resume the null residual.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
    // After outputting all 3 items, should not output anything more.
    assertEquals(0, tester.takeOutputElements().size());
    // Should also not ask to resume.
    assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
}

Example 23 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange 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, Long, Instant> tester = new ProcessFnTester<>(base, fn, InstantCoder.of(), SerializableCoder.of(OffsetRange.class), InstantCoder.of(), 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 24 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method testInvokeProcessElementOutputDisallowedAfterFailedTryClaim.

@Test
public void testInvokeProcessElementOutputDisallowedAfterFailedTryClaim() throws Exception {
    DoFn<Void, String> brokenFn = new DoFn<Void, String>() {

        @ProcessElement
        public void process(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
            assertFalse(tracker.tryClaim(6L));
            c.output("foo");
        }

        @GetInitialRestriction
        public OffsetRange getInitialRestriction(@Element Void element) {
            throw new UnsupportedOperationException("Should not be called in this test");
        }
    };
    e.expectMessage("Output is not allowed after a failed tryClaim()");
    runTest(brokenFn, new OffsetRange(0, 5));
}

Example 25 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method testInvokeProcessElementVoluntaryReturnResume.

@Test
public void testInvokeProcessElementVoluntaryReturnResume() throws Exception {
    SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result res = runTest(10, Duration.ZERO, 5, Duration.millis(100));
    assertTrue(res.getContinuation().shouldResume());
    assertEquals(new OffsetRange(5, 10), res.getResidualRestriction());
}