Examples with InMemoryReader org.apache.beam.runners.dataflow.worker.InMemoryReader used on opensource projects

Example 1 with InMemoryReader

use of org.apache.beam.runners.dataflow.worker.InMemoryReader in project beam by apache.

the class ReadOperationTest method testCheckpoint.

@Test
public void testCheckpoint() throws Exception {
    InMemoryReader<String> reader = new InMemoryReader<>(Arrays.asList("0", "1", "2", "3", "4", "5", "6", "7", "8", "9"), 0, 10, StringUtf8Coder.of());
    final ReadOperation[] readOperationHolder = new ReadOperation[1];
    OutputReceiver receiver = new OutputReceiver() {

        @Override
        public void process(Object elem) throws Exception {
            ReadOperation readOperation = readOperationHolder[0];
            if ("1".equals(elem)) {
                NativeReader.DynamicSplitResultWithPosition split = (NativeReader.DynamicSplitResultWithPosition) readOperation.requestCheckpoint();
                assertNotNull(split);
                assertEquals(positionAtIndex(2L), toCloudPosition(split.getAcceptedPosition()));
                // Check that the progress has been recomputed.
                ApproximateReportedProgress progress = readerProgressToCloudProgress(readOperation.getProgress());
                assertEquals(1, progress.getPosition().getRecordIndex().longValue());
            }
        }
    };
    ReadOperation readOperation = ReadOperation.forTest(reader, receiver, context);
    readOperation.setProgressUpdatePeriodMs(ReadOperation.UPDATE_ON_EACH_ITERATION);
    readOperationHolder[0] = readOperation;
    // An unstarted ReadOperation refuses checkpoint requests.
    assertNull(readOperation.requestCheckpoint());
    readOperation.start();
    readOperation.finish();
    // Operation is now finished. Check that it refuses a checkpoint request.
    assertNull(readOperation.requestCheckpoint());
}

Example 2 with InMemoryReader

use of org.apache.beam.runners.dataflow.worker.InMemoryReader in project beam by apache.

the class ReadOperationTest method testDynamicSplit.

@Test
public void testDynamicSplit() throws Exception {
    InMemoryReader<String> reader = new InMemoryReader<>(Arrays.asList("0", "1", "2", "3", "4", "5", "6", "7", "8", "9"), 0, 10, StringUtf8Coder.of());
    final ReadOperation[] operationHolder = new ReadOperation[1];
    OutputReceiver receiver = new OutputReceiver() {

        @Override
        public void process(Object elem) throws Exception {
            ReadOperation readOperation = operationHolder[0];
            if ("1".equals(elem)) {
                NativeReader.DynamicSplitResultWithPosition split = (NativeReader.DynamicSplitResultWithPosition) readOperation.requestDynamicSplit(splitRequestAtIndex(8L));
                assertNotNull(split);
                assertEquals(positionAtIndex(8L), toCloudPosition(split.getAcceptedPosition()));
                // Check that the progress has been recomputed.
                ApproximateReportedProgress progress = readerProgressToCloudProgress(readOperation.getProgress());
                assertEquals(1, progress.getPosition().getRecordIndex().longValue());
            } else if ("3".equals(elem)) {
                // Should accept a split at an earlier position than previously requested.
                // Should reject a split at a later position than previously requested.
                // Note that here we're testing our own MockReaderIterator class, so it's
                // kind of pointless, but we're also testing that ReadOperation correctly
                // relays the request to the iterator.
                NativeReader.DynamicSplitResultWithPosition split = (NativeReader.DynamicSplitResultWithPosition) readOperation.requestDynamicSplit(splitRequestAtIndex(6L));
                assertNotNull(split);
                assertEquals(positionAtIndex(6L), toCloudPosition(split.getAcceptedPosition()));
                split = (NativeReader.DynamicSplitResultWithPosition) readOperation.requestDynamicSplit(splitRequestAtIndex(6L));
                assertNull(split);
            }
        }
    };
    ReadOperation readOperation = ReadOperation.forTest(reader, receiver, context);
    readOperation.setProgressUpdatePeriodMs(ReadOperation.UPDATE_ON_EACH_ITERATION);
    operationHolder[0] = readOperation;
    // An unstarted ReadOperation refuses split requests.
    assertNull(readOperation.requestDynamicSplit(splitRequestAtIndex(8L)));
    readOperation.start();
    readOperation.finish();
    // Operation is now finished. Check that it refuses a split request.
    assertNull(readOperation.requestDynamicSplit(splitRequestAtIndex(5L)));
}