use of org.apache.beam.sdk.io.FileBasedSource.FileBasedReader in project beam by apache.
the class CompressedSourceTest method testSplittableProgress.
@Test
public void testSplittableProgress() throws IOException {
File tmpFile = tmpFolder.newFile("nonempty.txt");
String filename = tmpFile.toPath().toString();
Files.write(new byte[2], tmpFile);
PipelineOptions options = PipelineOptionsFactory.create();
CompressedSource<Byte> source = CompressedSource.from(new ByteSource(filename, 1));
try (BoundedReader<Byte> readerOrig = source.createReader(options)) {
assertThat(readerOrig, not(instanceOf(CompressedReader.class)));
assertThat(readerOrig, instanceOf(FileBasedReader.class));
FileBasedReader<Byte> reader = (FileBasedReader<Byte>) readerOrig;
// Check preconditions before starting
assertEquals(0.0, reader.getFractionConsumed(), 1e-6);
assertEquals(0, reader.getSplitPointsConsumed());
assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining());
// First record: none consumed, unknown remaining.
assertTrue(reader.start());
assertEquals(0, reader.getSplitPointsConsumed());
assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining());
// Second record: 1 consumed, know that we're on the last record.
assertTrue(reader.advance());
assertEquals(1, reader.getSplitPointsConsumed());
assertEquals(1, reader.getSplitPointsRemaining());
// Confirm empty and check post-conditions
assertFalse(reader.advance());
assertEquals(1.0, reader.getFractionConsumed(), 1e-6);
assertEquals(2, reader.getSplitPointsConsumed());
assertEquals(0, reader.getSplitPointsRemaining());
}
}
Aggregations