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

Example 36 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange 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(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, 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 37 with OffsetRange

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method testInvokeProcessElementTimeBoundedWithStartupDelay.

@Test
public void testInvokeProcessElementTimeBoundedWithStartupDelay() throws Exception {
    SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result res = runTest(10000, Duration.standardSeconds(3), Integer.MAX_VALUE, Duration.millis(100));
    assertFalse(res.getContinuation().shouldResume());
    OffsetRange residualRange = res.getResidualRestriction();
    // Same as above, but this time it counts from the time of the first tryClaim() call
    assertThat(residualRange.getFrom(), greaterThan(10L));
    assertThat(residualRange.getFrom(), lessThan(100L));
    assertEquals(10000, residualRange.getTo());
}

Example 38 with OffsetRange

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method testInvokeProcessElementOutputBounded.

@Test
public void testInvokeProcessElementOutputBounded() throws Exception {
    SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result res = runTest(10000, Duration.ZERO, Integer.MAX_VALUE, Duration.ZERO);
    assertFalse(res.getContinuation().shouldResume());
    OffsetRange residualRange = res.getResidualRestriction();
    // Should process the first 100 elements.
    assertEquals(1000, residualRange.getFrom());
    assertEquals(10000, residualRange.getTo());
}

Example 39 with OffsetRange

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method runTest.

private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result runTest(DoFn<Void, String> fn, OffsetRange initialRestriction) throws Exception {
    SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void> 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), () -> {
        throw new UnsupportedOperationException("BundleFinalizer not configured for test.");
    });
    SplittableProcessElementInvoker.Result rval = invoker.invokeProcessElement(DoFnInvokers.invokerFor(fn), WindowedValue.of(null, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING), new OffsetRangeTracker(initialRestriction), new WatermarkEstimator<Void>() {

        @Override
        public Instant currentWatermark() {
            return GlobalWindow.TIMESTAMP_MIN_VALUE;
        }

        @Override
        public Void getState() {
            return null;
        }
    }, Collections.emptyMap());
    return rval;
}

Example 40 with OffsetRange

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

the class OutputAndTimeBoundedSplittableProcessElementInvokerTest method runTest.

private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result runTest(int totalNumOutputs, Duration sleepBeforeFirstClaim, int numOutputsPerProcessCall, Duration sleepBeforeEachOutput) throws Exception {
    SomeFn fn = new SomeFn(sleepBeforeFirstClaim, numOutputsPerProcessCall, sleepBeforeEachOutput);
    OffsetRange initialRestriction = new OffsetRange(0, totalNumOutputs);
    return runTest(fn, initialRestriction);
}