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

Example 1 with AvroByteFileIterator

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

the class AvroByteReaderTest method testGetProgress.

@Test
public void testGetProgress() throws Exception {
    // Ensure that AvroByteReader reports progress from the underlying AvroSource.
    // 4 blocks with 4 split points.
    List<List<String>> elements = generateInputBlocks(4, 10, 10);
    Coder<String> coder = StringUtf8Coder.of();
    AvroFileInfo<String> fileInfo = initInputFile(elements, coder);
    AvroByteReader<String> reader = new AvroByteReader<String>(fileInfo.filename, 0L, Long.MAX_VALUE, coder, PipelineOptionsFactory.create());
    AvroByteFileIterator iterator = reader.iterator();
    Assert.assertTrue(iterator.start());
    ApproximateReportedProgress progress = readerProgressToCloudProgress(reader.iterator().getProgress());
    Assert.assertEquals(0.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(0.0, progress.getFractionConsumed(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());
    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(1.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());
    // Advance to the end of last block and check consumed parallelism along the way.
    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(2.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());
    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(3.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(1.0, progress.getRemainingParallelism().getValue(), 1e-6);
    Assert.assertFalse(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(4.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(0.0, progress.getRemainingParallelism().getValue(), 1e-6);
    Assert.assertEquals(1.0, progress.getFractionConsumed(), 1e-6);
}