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

Example 1 with OffsetRange

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

the class OffsetBasedSource method split.

@Override
public List<? extends OffsetBasedSource<T>> split(long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
    // Split the range into bundles based on the desiredBundleSizeBytes. If the desired bundle
    // size is smaller than the minBundleSize of the source then minBundleSize will be used instead.
    long desiredBundleSizeOffsetUnits = Math.max(Math.max(1, desiredBundleSizeBytes / getBytesPerOffset()), minBundleSize);
    List<OffsetBasedSource<T>> subSources = new ArrayList<>();
    for (OffsetRange range : new OffsetRange(startOffset, Math.min(endOffset, getMaxEndOffset(options))).split(desiredBundleSizeOffsetUnits, minBundleSize)) {
        subSources.add(createSourceForSubrange(range.getFrom(), range.getTo()));
    }
    return subSources;
}

Example 2 with OffsetRange

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

the class ReadFromKafkaDoFnTest method testInitialRestrictionWhenHasStopOffset.

@Test
public void testInitialRestrictionWhenHasStopOffset() throws Exception {
    long expectedStartOffset = 10L;
    long expectedStopOffset = 20L;
    consumer.setStartOffsetForTime(15L, Instant.now());
    consumer.setStopOffsetForTime(18L, Instant.now());
    consumer.setCurrentPos(5L);
    OffsetRange result = dofnInstance.initialRestriction(KafkaSourceDescriptor.of(topicPartition, expectedStartOffset, null, expectedStopOffset, null, ImmutableList.of()));
    assertEquals(new OffsetRange(expectedStartOffset, expectedStopOffset), result);
}

Example 3 with OffsetRange

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

the class ReadFromKafkaDoFnTest method testProcessElementWithEmptyPoll.

@Test
public void testProcessElementWithEmptyPoll() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    consumer.setNumOfRecordsPerPoll(-1);
    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.resume(), result);
    assertTrue(receiver.getOutputs().isEmpty());
}

Example 4 with OffsetRange

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

the class ReadFromKafkaDoFnTest method testInitialRestrictionWhenHasStartOffset.

@Test
public void testInitialRestrictionWhenHasStartOffset() throws Exception {
    long expectedStartOffset = 10L;
    consumer.setStartOffsetForTime(15L, Instant.now());
    consumer.setCurrentPos(5L);
    OffsetRange result = dofnInstance.initialRestriction(KafkaSourceDescriptor.of(topicPartition, expectedStartOffset, null, null, null, ImmutableList.of()));
    assertEquals(new OffsetRange(expectedStartOffset, Long.MAX_VALUE), result);
}

Example 5 with OffsetRange

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

the class ReadFromKafkaDoFnTest method testProcessElementWhenTopicPartitionIsStopped.

@Test
public void testProcessElementWhenTopicPartitionIsStopped() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    ReadFromKafkaDoFn<String, String> instance = new ReadFromKafkaDoFn(makeReadSourceDescriptor(consumer).toBuilder().setCheckStopReadingFn(new SerializableFunction<TopicPartition, Boolean>() {

        @Override
        public Boolean apply(TopicPartition input) {
            assertTrue(input.equals(topicPartition));
            return true;
        }
    }).build());
    consumer.setNumOfRecordsPerPoll(10);
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(0L, Long.MAX_VALUE));
    ProcessContinuation result = instance.processElement(KafkaSourceDescriptor.of(topicPartition, null, null, null, null, null), tracker, null, (OutputReceiver) receiver);
    assertEquals(ProcessContinuation.stop(), result);
}