Search in sources :

Example 1 with FileBasedReader

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());
    }
}
Also used : FileBasedReader(org.apache.beam.sdk.io.FileBasedSource.FileBasedReader) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) Matchers.containsString(org.hamcrest.Matchers.containsString) File(java.io.File) Test(org.junit.Test)

Aggregations

File (java.io.File)1 FileBasedReader (org.apache.beam.sdk.io.FileBasedSource.FileBasedReader)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 Matchers.containsString (org.hamcrest.Matchers.containsString)1 Test (org.junit.Test)1