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

Example 6 with OffsetRange

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

the class ReadFromKafkaDoFnTest method testProcessElementWhenTopicPartitionIsRemoved.

@Test
public void testProcessElementWhenTopicPartitionIsRemoved() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    consumer.setRemoved();
    consumer.setNumOfRecordsPerPoll(10);
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(0L, Long.MAX_VALUE));
    ProcessContinuation result = dofnInstance.processElement(KafkaSourceDescriptor.of(topicPartition, null, null, null, null, null), tracker, null, (OutputReceiver) receiver);
    assertEquals(ProcessContinuation.stop(), result);
}

Example 7 with OffsetRange

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

the class OffsetRangeTracker method trySplit.

@Override
public SplitResult<OffsetRange> trySplit(double fractionOfRemainder) {
    // Convert to BigDecimal in computation to prevent overflow, which may result in loss of
    // precision.
    BigDecimal cur = (lastAttemptedOffset == null) ? BigDecimal.valueOf(range.getFrom()).subtract(BigDecimal.ONE, MathContext.DECIMAL128) : BigDecimal.valueOf(lastAttemptedOffset);
    // split = cur + max(1, (range.getTo() - cur) * fractionOfRemainder)
    BigDecimal splitPos = cur.add(BigDecimal.valueOf(range.getTo()).subtract(cur, MathContext.DECIMAL128).multiply(BigDecimal.valueOf(fractionOfRemainder), MathContext.DECIMAL128).max(BigDecimal.ONE), MathContext.DECIMAL128);
    long split = splitPos.longValue();
    if (split >= range.getTo()) {
        return null;
    }
    OffsetRange res = new OffsetRange(split, range.getTo());
    this.range = new OffsetRange(range.getFrom(), split);
    return SplitResult.of(range, res);
}

Example 8 with OffsetRange

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

the class OffsetRangeTrackerTest method testLargeRange.

@Test
public void testLargeRange() throws Exception {
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(Long.MIN_VALUE, Long.MAX_VALUE));
    Progress progress = tracker.getProgress();
    assertEquals(0, progress.getWorkCompleted(), 0.001);
    assertEquals(BigDecimal.valueOf(Long.MAX_VALUE).subtract(BigDecimal.valueOf(Long.MIN_VALUE), MathContext.DECIMAL128).doubleValue(), progress.getWorkRemaining(), 0.001);
    SplitResult res = tracker.trySplit(0);
    assertEquals(new OffsetRange(Long.MIN_VALUE, Long.MIN_VALUE), res.getPrimary());
    assertEquals(new OffsetRange(Long.MIN_VALUE, Long.MAX_VALUE), res.getResidual());
}

Example 9 with OffsetRange

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

the class OffsetRangeTrackerTest method testBacklogFinished.

@Test
public void testBacklogFinished() {
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(0, 200));
    tracker.tryClaim(300L);
    Progress progress = tracker.getProgress();
    assertEquals(200, progress.getWorkCompleted(), 0.001);
    assertEquals(0, progress.getWorkRemaining(), 0.001);
    tracker = new OffsetRangeTracker(new OffsetRange(100, 200));
    tracker.tryClaim(300L);
    progress = tracker.getProgress();
    assertEquals(100, progress.getWorkCompleted(), 0.001);
    assertEquals(0, progress.getWorkRemaining(), 0.001);
}

Example 10 with OffsetRange

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

the class OffsetRangeTrackerTest method testCheckDoneAfterTryClaimPastEndOfRange.

@Test
public void testCheckDoneAfterTryClaimPastEndOfRange() {
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(100, 200));
    assertTrue(tracker.tryClaim(150L));
    assertTrue(tracker.tryClaim(175L));
    assertFalse(tracker.tryClaim(220L));
    tracker.checkDone();
}