| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.sdk.io; |
| |
| import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem; |
| import static org.hamcrest.Matchers.containsInAnyOrder; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertSame; |
| import static org.junit.Assert.assertThat; |
| import static org.junit.Assert.assertTrue; |
| |
| import java.io.ByteArrayInputStream; |
| import java.io.File; |
| import java.io.FileOutputStream; |
| import java.io.IOException; |
| import java.io.PushbackInputStream; |
| import java.util.ArrayList; |
| import java.util.Collections; |
| import java.util.List; |
| import java.util.NoSuchElementException; |
| import java.util.Objects; |
| import java.util.Random; |
| import org.apache.avro.Schema; |
| import org.apache.avro.file.CodecFactory; |
| import org.apache.avro.file.DataFileConstants; |
| import org.apache.avro.file.DataFileWriter; |
| import org.apache.avro.generic.GenericDatumWriter; |
| import org.apache.avro.generic.GenericRecord; |
| import org.apache.avro.io.DatumWriter; |
| import org.apache.avro.reflect.AvroDefault; |
| import org.apache.avro.reflect.Nullable; |
| import org.apache.avro.reflect.ReflectData; |
| import org.apache.avro.reflect.ReflectDatumWriter; |
| import org.apache.beam.sdk.coders.AvroCoder; |
| import org.apache.beam.sdk.coders.DefaultCoder; |
| import org.apache.beam.sdk.io.AvroSource.AvroMetadata; |
| import org.apache.beam.sdk.io.AvroSource.AvroReader; |
| import org.apache.beam.sdk.io.AvroSource.AvroReader.Seeker; |
| import org.apache.beam.sdk.io.BlockBasedSource.BlockBasedReader; |
| import org.apache.beam.sdk.io.BoundedSource.BoundedReader; |
| import org.apache.beam.sdk.io.fs.MatchResult.Metadata; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.testing.SourceTestUtils; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.util.SerializableUtils; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.base.MoreObjects; |
| import org.hamcrest.Matchers; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.ExpectedException; |
| import org.junit.rules.TemporaryFolder; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** Tests for AvroSource. */ |
| @RunWith(JUnit4.class) |
| public class AvroSourceTest { |
| @Rule public TemporaryFolder tmpFolder = new TemporaryFolder(); |
| |
| @Rule public ExpectedException expectedException = ExpectedException.none(); |
| |
| private enum SyncBehavior { |
| SYNC_REGULAR, // Sync at regular, user defined intervals |
| SYNC_RANDOM, // Sync at random intervals |
| SYNC_DEFAULT // Sync at default intervals (i.e., no manual syncing). |
| } |
| |
| private static final int DEFAULT_RECORD_COUNT = 1000; |
| |
| /** |
| * Generates an input Avro file containing the given records in the temporary directory and |
| * returns the full path of the file. |
| */ |
| private <T> String generateTestFile( |
| String filename, |
| List<T> elems, |
| SyncBehavior syncBehavior, |
| int syncInterval, |
| AvroCoder<T> coder, |
| String codec) |
| throws IOException { |
| Random random = new Random(0); |
| File tmpFile = tmpFolder.newFile(filename); |
| String path = tmpFile.toString(); |
| |
| FileOutputStream os = new FileOutputStream(tmpFile); |
| DatumWriter<T> datumWriter = |
| coder.getType().equals(GenericRecord.class) |
| ? new GenericDatumWriter<>(coder.getSchema()) |
| : new ReflectDatumWriter<>(coder.getSchema()); |
| try (DataFileWriter<T> writer = new DataFileWriter<>(datumWriter)) { |
| writer.setCodec(CodecFactory.fromString(codec)); |
| writer.create(coder.getSchema(), os); |
| |
| int recordIndex = 0; |
| int syncIndex = syncBehavior == SyncBehavior.SYNC_RANDOM ? random.nextInt(syncInterval) : 0; |
| |
| for (T elem : elems) { |
| writer.append(elem); |
| recordIndex++; |
| |
| switch (syncBehavior) { |
| case SYNC_REGULAR: |
| if (recordIndex == syncInterval) { |
| recordIndex = 0; |
| writer.sync(); |
| } |
| break; |
| case SYNC_RANDOM: |
| if (recordIndex == syncIndex) { |
| recordIndex = 0; |
| writer.sync(); |
| syncIndex = random.nextInt(syncInterval); |
| } |
| break; |
| case SYNC_DEFAULT: |
| default: |
| } |
| } |
| } |
| return path; |
| } |
| |
| @Test |
| public void testReadWithDifferentCodecs() throws Exception { |
| // Test reading files generated using all codecs. |
| String[] codecs = { |
| DataFileConstants.NULL_CODEC, |
| DataFileConstants.BZIP2_CODEC, |
| DataFileConstants.DEFLATE_CODEC, |
| DataFileConstants.SNAPPY_CODEC, |
| DataFileConstants.XZ_CODEC, |
| }; |
| // As Avro's default block size is 64KB, write 64K records to ensure at least one full block. |
| // We could make this smaller than 64KB assuming each record is at least B bytes, but then the |
| // test could silently stop testing the failure condition from BEAM-422. |
| List<Bird> expected = createRandomRecords(1 << 16); |
| |
| for (String codec : codecs) { |
| String filename = |
| generateTestFile( |
| codec, expected, SyncBehavior.SYNC_DEFAULT, 0, AvroCoder.of(Bird.class), codec); |
| AvroSource<Bird> source = AvroSource.from(filename).withSchema(Bird.class); |
| List<Bird> actual = SourceTestUtils.readFromSource(source, null); |
| assertThat(expected, containsInAnyOrder(actual.toArray())); |
| } |
| } |
| |
| @Test |
| public void testSplitAtFraction() throws Exception { |
| // A reduced dataset is enough here. |
| List<FixedRecord> expected = createFixedRecords(DEFAULT_RECORD_COUNT); |
| // Create an AvroSource where each block is 1/10th of the total set of records. |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| expected, |
| SyncBehavior.SYNC_REGULAR, |
| DEFAULT_RECORD_COUNT / 10 /* max records per block */, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| File file = new File(filename); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| List<? extends BoundedSource<FixedRecord>> splits = source.split(file.length() / 3, null); |
| for (BoundedSource<FixedRecord> subSource : splits) { |
| int items = SourceTestUtils.readFromSource(subSource, null).size(); |
| // Shouldn't split while unstarted. |
| SourceTestUtils.assertSplitAtFractionFails(subSource, 0, 0.0, null); |
| SourceTestUtils.assertSplitAtFractionFails(subSource, 0, 0.7, null); |
| SourceTestUtils.assertSplitAtFractionSucceedsAndConsistent(subSource, 1, 0.7, null); |
| SourceTestUtils.assertSplitAtFractionSucceedsAndConsistent( |
| subSource, DEFAULT_RECORD_COUNT / 100, 0.7, null); |
| SourceTestUtils.assertSplitAtFractionSucceedsAndConsistent( |
| subSource, DEFAULT_RECORD_COUNT / 10, 0.1, null); |
| SourceTestUtils.assertSplitAtFractionFails( |
| subSource, DEFAULT_RECORD_COUNT / 10 + 1, 0.1, null); |
| SourceTestUtils.assertSplitAtFractionFails(subSource, DEFAULT_RECORD_COUNT / 3, 0.3, null); |
| SourceTestUtils.assertSplitAtFractionFails(subSource, items, 0.9, null); |
| SourceTestUtils.assertSplitAtFractionFails(subSource, items, 1.0, null); |
| SourceTestUtils.assertSplitAtFractionSucceedsAndConsistent(subSource, items, 0.999, null); |
| } |
| } |
| |
| @Test |
| public void testGetProgressFromUnstartedReader() throws Exception { |
| List<FixedRecord> records = createFixedRecords(DEFAULT_RECORD_COUNT); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| records, |
| SyncBehavior.SYNC_DEFAULT, |
| 1000, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| File file = new File(filename); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| try (BoundedSource.BoundedReader<FixedRecord> reader = source.createReader(null)) { |
| assertEquals(Double.valueOf(0.0), reader.getFractionConsumed()); |
| } |
| |
| List<? extends BoundedSource<FixedRecord>> splits = source.split(file.length() / 3, null); |
| for (BoundedSource<FixedRecord> subSource : splits) { |
| try (BoundedSource.BoundedReader<FixedRecord> reader = subSource.createReader(null)) { |
| assertEquals(Double.valueOf(0.0), reader.getFractionConsumed()); |
| } |
| } |
| } |
| |
| @Test |
| public void testProgress() throws Exception { |
| // 5 records, 2 per block. |
| List<FixedRecord> records = createFixedRecords(5); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| records, |
| SyncBehavior.SYNC_REGULAR, |
| 2, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| try (BoundedSource.BoundedReader<FixedRecord> readerOrig = source.createReader(null)) { |
| assertThat(readerOrig, Matchers.instanceOf(BlockBasedReader.class)); |
| BlockBasedReader<FixedRecord> reader = (BlockBasedReader<FixedRecord>) readerOrig; |
| |
| // Before starting |
| assertEquals(0.0, reader.getFractionConsumed(), 1e-6); |
| assertEquals(0, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| |
| // First 2 records are in the same block. |
| assertTrue(reader.start()); |
| assertTrue(reader.isAtSplitPoint()); |
| assertEquals(0, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| // continued |
| assertTrue(reader.advance()); |
| assertFalse(reader.isAtSplitPoint()); |
| assertEquals(0, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| |
| // Second block -> parallelism consumed becomes 1. |
| assertTrue(reader.advance()); |
| assertTrue(reader.isAtSplitPoint()); |
| assertEquals(1, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| // continued |
| assertTrue(reader.advance()); |
| assertFalse(reader.isAtSplitPoint()); |
| assertEquals(1, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| |
| // Third and final block -> parallelism consumed becomes 2, remaining becomes 1. |
| assertTrue(reader.advance()); |
| assertTrue(reader.isAtSplitPoint()); |
| assertEquals(2, reader.getSplitPointsConsumed()); |
| assertEquals(1, reader.getSplitPointsRemaining()); |
| |
| // Done |
| assertFalse(reader.advance()); |
| assertEquals(3, reader.getSplitPointsConsumed()); |
| assertEquals(0, reader.getSplitPointsRemaining()); |
| assertEquals(1.0, reader.getFractionConsumed(), 1e-6); |
| } |
| } |
| |
| @Test |
| public void testProgressEmptySource() throws Exception { |
| // 0 records, 20 per block. |
| List<FixedRecord> records = Collections.emptyList(); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| records, |
| SyncBehavior.SYNC_REGULAR, |
| 2, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| try (BoundedSource.BoundedReader<FixedRecord> readerOrig = source.createReader(null)) { |
| assertThat(readerOrig, Matchers.instanceOf(BlockBasedReader.class)); |
| BlockBasedReader<FixedRecord> reader = (BlockBasedReader<FixedRecord>) readerOrig; |
| |
| // before starting |
| assertEquals(0.0, reader.getFractionConsumed(), 1e-6); |
| assertEquals(0, reader.getSplitPointsConsumed()); |
| assertEquals(BoundedReader.SPLIT_POINTS_UNKNOWN, reader.getSplitPointsRemaining()); |
| |
| // confirm empty |
| assertFalse(reader.start()); |
| |
| // after reading empty source |
| assertEquals(0, reader.getSplitPointsConsumed()); |
| assertEquals(0, reader.getSplitPointsRemaining()); |
| assertEquals(1.0, reader.getFractionConsumed(), 1e-6); |
| } |
| } |
| |
| @Test |
| public void testGetCurrentFromUnstartedReader() throws Exception { |
| List<FixedRecord> records = createFixedRecords(DEFAULT_RECORD_COUNT); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| records, |
| SyncBehavior.SYNC_DEFAULT, |
| 1000, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| try (BlockBasedSource.BlockBasedReader<FixedRecord> reader = |
| (BlockBasedSource.BlockBasedReader<FixedRecord>) source.createReader(null)) { |
| assertEquals(null, reader.getCurrentBlock()); |
| |
| expectedException.expect(NoSuchElementException.class); |
| expectedException.expectMessage("No block has been successfully read from"); |
| reader.getCurrent(); |
| } |
| } |
| |
| @Test |
| public void testSplitAtFractionExhaustive() throws Exception { |
| // A small-sized input is sufficient, because the test verifies that splitting is non-vacuous. |
| List<FixedRecord> expected = createFixedRecords(20); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| expected, |
| SyncBehavior.SYNC_REGULAR, |
| 5, |
| AvroCoder.of(FixedRecord.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<FixedRecord> source = AvroSource.from(filename).withSchema(FixedRecord.class); |
| SourceTestUtils.assertSplitAtFractionExhaustive(source, null); |
| } |
| |
| @Test |
| public void testSplitsWithSmallBlocks() throws Exception { |
| PipelineOptions options = PipelineOptionsFactory.create(); |
| // Test reading from an object file with many small random-sized blocks. |
| // The file itself doesn't have to be big; we can use a decreased record count. |
| List<Bird> expected = createRandomRecords(DEFAULT_RECORD_COUNT); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| expected, |
| SyncBehavior.SYNC_RANDOM, |
| DEFAULT_RECORD_COUNT / 20 /* max records/block */, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| File file = new File(filename); |
| |
| // Small minimum bundle size |
| AvroSource<Bird> source = |
| AvroSource.from(filename).withSchema(Bird.class).withMinBundleSize(100L); |
| |
| // Assert that the source produces the expected records |
| assertEquals(expected, SourceTestUtils.readFromSource(source, options)); |
| |
| List<? extends BoundedSource<Bird>> splits; |
| int nonEmptySplits; |
| |
| // Split with the minimum bundle size |
| splits = source.split(100L, options); |
| assertTrue(splits.size() > 2); |
| SourceTestUtils.assertSourcesEqualReferenceSource(source, splits, options); |
| nonEmptySplits = 0; |
| for (BoundedSource<Bird> subSource : splits) { |
| if (SourceTestUtils.readFromSource(subSource, options).size() > 0) { |
| nonEmptySplits += 1; |
| } |
| } |
| assertTrue(nonEmptySplits > 2); |
| |
| // Split with larger bundle size |
| splits = source.split(file.length() / 4, options); |
| assertTrue(splits.size() > 2); |
| SourceTestUtils.assertSourcesEqualReferenceSource(source, splits, options); |
| nonEmptySplits = 0; |
| for (BoundedSource<Bird> subSource : splits) { |
| if (SourceTestUtils.readFromSource(subSource, options).size() > 0) { |
| nonEmptySplits += 1; |
| } |
| } |
| assertTrue(nonEmptySplits > 2); |
| |
| // Split with the file length |
| splits = source.split(file.length(), options); |
| assertTrue(splits.size() == 1); |
| SourceTestUtils.assertSourcesEqualReferenceSource(source, splits, options); |
| } |
| |
| @Test |
| public void testMultipleFiles() throws Exception { |
| String baseName = "tmp-"; |
| List<Bird> expected = new ArrayList<>(); |
| for (int i = 0; i < 10; i++) { |
| List<Bird> contents = createRandomRecords(DEFAULT_RECORD_COUNT / 10); |
| expected.addAll(contents); |
| generateTestFile( |
| baseName + i, |
| contents, |
| SyncBehavior.SYNC_DEFAULT, |
| 0, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| } |
| |
| AvroSource<Bird> source = |
| AvroSource.from(new File(tmpFolder.getRoot().toString(), baseName + "*").toString()) |
| .withSchema(Bird.class); |
| List<Bird> actual = SourceTestUtils.readFromSource(source, null); |
| assertThat(actual, containsInAnyOrder(expected.toArray())); |
| } |
| |
| @Test |
| public void testCreationWithSchema() throws Exception { |
| List<Bird> expected = createRandomRecords(100); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| expected, |
| SyncBehavior.SYNC_DEFAULT, |
| 0, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| |
| // Create a source with a schema object |
| Schema schema = ReflectData.get().getSchema(Bird.class); |
| AvroSource<GenericRecord> source = AvroSource.from(filename).withSchema(schema); |
| List<GenericRecord> records = SourceTestUtils.readFromSource(source, null); |
| assertEqualsWithGeneric(expected, records); |
| |
| // Create a source with a JSON schema |
| String schemaString = ReflectData.get().getSchema(Bird.class).toString(); |
| source = AvroSource.from(filename).withSchema(schemaString); |
| records = SourceTestUtils.readFromSource(source, null); |
| assertEqualsWithGeneric(expected, records); |
| } |
| |
| @Test |
| public void testSchemaUpdate() throws Exception { |
| List<Bird> birds = createRandomRecords(100); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| birds, |
| SyncBehavior.SYNC_DEFAULT, |
| 0, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<FancyBird> source = AvroSource.from(filename).withSchema(FancyBird.class); |
| List<FancyBird> actual = SourceTestUtils.readFromSource(source, null); |
| |
| List<FancyBird> expected = new ArrayList<>(); |
| for (Bird bird : birds) { |
| expected.add( |
| new FancyBird( |
| bird.number, bird.species, bird.quality, bird.quantity, null, "MAXIMUM OVERDRIVE")); |
| } |
| |
| assertThat(actual, containsInAnyOrder(expected.toArray())); |
| } |
| |
| @Test |
| public void testSchemaStringIsInterned() throws Exception { |
| List<Bird> birds = createRandomRecords(100); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| birds, |
| SyncBehavior.SYNC_DEFAULT, |
| 0, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| Metadata fileMetadata = FileSystems.matchSingleFileSpec(filename); |
| String schema = AvroSource.readMetadataFromFile(fileMetadata.resourceId()).getSchemaString(); |
| // Add "" to the schema to make sure it is not interned. |
| AvroSource<GenericRecord> sourceA = AvroSource.from(filename).withSchema("" + schema); |
| AvroSource<GenericRecord> sourceB = AvroSource.from(filename).withSchema("" + schema); |
| assertSame(sourceA.getReaderSchemaString(), sourceB.getReaderSchemaString()); |
| |
| // Ensure that deserialization still goes through interning |
| AvroSource<GenericRecord> sourceC = SerializableUtils.clone(sourceB); |
| assertSame(sourceA.getReaderSchemaString(), sourceC.getReaderSchemaString()); |
| } |
| |
| @Test |
| public void testParseFn() throws Exception { |
| List<Bird> expected = createRandomRecords(100); |
| String filename = |
| generateTestFile( |
| "tmp.avro", |
| expected, |
| SyncBehavior.SYNC_DEFAULT, |
| 0, |
| AvroCoder.of(Bird.class), |
| DataFileConstants.NULL_CODEC); |
| |
| AvroSource<Bird> source = |
| AvroSource.from(filename) |
| .withParseFn( |
| input -> |
| new Bird( |
| (long) input.get("number"), |
| input.get("species").toString(), |
| input.get("quality").toString(), |
| (long) input.get("quantity")), |
| AvroCoder.of(Bird.class)); |
| List<Bird> actual = SourceTestUtils.readFromSource(source, null); |
| assertThat(actual, containsInAnyOrder(expected.toArray())); |
| } |
| |
| private void assertEqualsWithGeneric(List<Bird> expected, List<GenericRecord> actual) { |
| assertEquals(expected.size(), actual.size()); |
| for (int i = 0; i < expected.size(); i++) { |
| Bird fixed = expected.get(i); |
| GenericRecord generic = actual.get(i); |
| assertEquals(fixed.number, generic.get("number")); |
| assertEquals(fixed.quality, generic.get("quality").toString()); // From Avro util.Utf8 |
| assertEquals(fixed.quantity, generic.get("quantity")); |
| assertEquals(fixed.species, generic.get("species").toString()); |
| } |
| } |
| |
| /** |
| * Creates a haystack byte array of the give size with a needle that starts at the given position. |
| */ |
| private byte[] createHaystack(byte[] needle, int position, int size) { |
| byte[] haystack = new byte[size]; |
| for (int i = position, j = 0; i < size && j < needle.length; i++, j++) { |
| haystack[i] = needle[j]; |
| } |
| return haystack; |
| } |
| |
| /** |
| * Asserts that advancePastNextSyncMarker advances an input stream past a sync marker and |
| * correctly returns the number of bytes consumed from the stream. Creates a haystack of size |
| * bytes and places a 16-byte sync marker at the position specified. |
| */ |
| private void testAdvancePastNextSyncMarkerAt(int position, int size) throws IOException { |
| byte sentinel = (byte) 0xFF; |
| byte[] marker = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}; |
| byte[] haystack = createHaystack(marker, position, size); |
| PushbackInputStream stream = |
| new PushbackInputStream(new ByteArrayInputStream(haystack), marker.length); |
| if (position + marker.length < size) { |
| haystack[position + marker.length] = sentinel; |
| assertEquals(position + marker.length, AvroReader.advancePastNextSyncMarker(stream, marker)); |
| assertEquals(sentinel, (byte) stream.read()); |
| } else { |
| assertEquals(size, AvroReader.advancePastNextSyncMarker(stream, marker)); |
| assertEquals(-1, stream.read()); |
| } |
| } |
| |
| @Test |
| public void testAdvancePastNextSyncMarker() throws IOException { |
| // Test placing the sync marker at different locations at the start and in the middle of the |
| // buffer. |
| for (int i = 0; i <= 16; i++) { |
| testAdvancePastNextSyncMarkerAt(i, 1000); |
| testAdvancePastNextSyncMarkerAt(160 + i, 1000); |
| } |
| // Test placing the sync marker at the end of the buffer. |
| testAdvancePastNextSyncMarkerAt(983, 1000); |
| // Test placing the sync marker so that it begins at the end of the buffer. |
| testAdvancePastNextSyncMarkerAt(984, 1000); |
| testAdvancePastNextSyncMarkerAt(985, 1000); |
| testAdvancePastNextSyncMarkerAt(999, 1000); |
| // Test with no sync marker. |
| testAdvancePastNextSyncMarkerAt(1000, 1000); |
| } |
| |
| // Tests for Seeker. |
| @Test |
| public void testSeekerFind() { |
| byte[] marker = {0, 1, 2, 3}; |
| byte[] buffer; |
| Seeker s; |
| s = new Seeker(marker); |
| |
| buffer = new byte[] {0, 1, 2, 3, 4, 5, 6, 7}; |
| assertEquals(3, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {0, 0, 0, 0, 0, 1, 2, 3}; |
| assertEquals(7, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {0, 1, 2, 0, 0, 1, 2, 3}; |
| assertEquals(7, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {0, 1, 2, 3}; |
| assertEquals(3, s.find(buffer, buffer.length)); |
| } |
| |
| @Test |
| public void testSeekerFindResume() { |
| byte[] marker = {0, 1, 2, 3}; |
| byte[] buffer; |
| Seeker s; |
| s = new Seeker(marker); |
| |
| buffer = new byte[] {0, 0, 0, 0, 0, 0, 0, 0}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {1, 2, 3, 0, 0, 0, 0, 0}; |
| assertEquals(2, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {0, 0, 0, 0, 0, 0, 1, 2}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {3, 0, 1, 2, 3, 0, 1, 2}; |
| assertEquals(0, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {0}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {1}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {2}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {3}; |
| assertEquals(0, s.find(buffer, buffer.length)); |
| } |
| |
| @Test |
| public void testSeekerUsesBufferLength() { |
| byte[] marker = {0, 0, 1}; |
| byte[] buffer; |
| Seeker s; |
| s = new Seeker(marker); |
| |
| buffer = new byte[] {0, 0, 0, 1}; |
| assertEquals(-1, s.find(buffer, 3)); |
| |
| s = new Seeker(marker); |
| buffer = new byte[] {0, 0}; |
| assertEquals(-1, s.find(buffer, 1)); |
| buffer = new byte[] {1, 0}; |
| assertEquals(-1, s.find(buffer, 1)); |
| |
| s = new Seeker(marker); |
| buffer = new byte[] {0, 2}; |
| assertEquals(-1, s.find(buffer, 1)); |
| buffer = new byte[] {0, 2}; |
| assertEquals(-1, s.find(buffer, 1)); |
| buffer = new byte[] {1, 2}; |
| assertEquals(0, s.find(buffer, 1)); |
| } |
| |
| @Test |
| public void testSeekerFindPartial() { |
| byte[] marker = {0, 0, 1}; |
| byte[] buffer; |
| Seeker s; |
| s = new Seeker(marker); |
| |
| buffer = new byte[] {0, 0, 0, 1}; |
| assertEquals(3, s.find(buffer, buffer.length)); |
| |
| marker = new byte[] {1, 1, 1, 2}; |
| s = new Seeker(marker); |
| |
| buffer = new byte[] {1, 1, 1, 1, 1}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {1, 1, 2}; |
| assertEquals(2, s.find(buffer, buffer.length)); |
| |
| buffer = new byte[] {1, 1, 1, 1, 1}; |
| assertEquals(-1, s.find(buffer, buffer.length)); |
| buffer = new byte[] {2, 1, 1, 1, 2}; |
| assertEquals(0, s.find(buffer, buffer.length)); |
| } |
| |
| @Test |
| public void testSeekerFindAllLocations() { |
| byte[] marker = {1, 1, 2}; |
| byte[] allOnes = new byte[] {1, 1, 1, 1}; |
| byte[] findIn = new byte[] {1, 1, 1, 1}; |
| Seeker s = new Seeker(marker); |
| |
| for (int i = 0; i < findIn.length; i++) { |
| assertEquals(-1, s.find(allOnes, allOnes.length)); |
| findIn[i] = 2; |
| assertEquals(i, s.find(findIn, findIn.length)); |
| findIn[i] = 1; |
| } |
| } |
| |
| @Test |
| public void testDisplayData() { |
| AvroSource<Bird> source = |
| AvroSource.from("foobar.txt").withSchema(Bird.class).withMinBundleSize(1234); |
| |
| DisplayData displayData = DisplayData.from(source); |
| assertThat(displayData, hasDisplayItem("filePattern", "foobar.txt")); |
| assertThat(displayData, hasDisplayItem("minBundleSize", 1234)); |
| } |
| |
| @Test |
| public void testReadMetadataWithCodecs() throws Exception { |
| // Test reading files generated using all codecs. |
| String[] codecs = { |
| DataFileConstants.NULL_CODEC, |
| DataFileConstants.BZIP2_CODEC, |
| DataFileConstants.DEFLATE_CODEC, |
| DataFileConstants.SNAPPY_CODEC, |
| DataFileConstants.XZ_CODEC |
| }; |
| List<Bird> expected = createRandomRecords(DEFAULT_RECORD_COUNT); |
| |
| for (String codec : codecs) { |
| String filename = |
| generateTestFile( |
| codec, expected, SyncBehavior.SYNC_DEFAULT, 0, AvroCoder.of(Bird.class), codec); |
| |
| Metadata fileMeta = FileSystems.matchSingleFileSpec(filename); |
| AvroMetadata metadata = AvroSource.readMetadataFromFile(fileMeta.resourceId()); |
| assertEquals(codec, metadata.getCodec()); |
| } |
| } |
| |
| @Test |
| public void testReadSchemaString() throws Exception { |
| List<Bird> expected = createRandomRecords(DEFAULT_RECORD_COUNT); |
| String codec = DataFileConstants.NULL_CODEC; |
| String filename = |
| generateTestFile( |
| codec, expected, SyncBehavior.SYNC_DEFAULT, 0, AvroCoder.of(Bird.class), codec); |
| Metadata fileMeta = FileSystems.matchSingleFileSpec(filename); |
| AvroMetadata metadata = AvroSource.readMetadataFromFile(fileMeta.resourceId()); |
| // By default, parse validates the schema, which is what we want. |
| Schema schema = new Schema.Parser().parse(metadata.getSchemaString()); |
| assertEquals(4, schema.getFields().size()); |
| } |
| |
| @Test |
| public void testCreateFromMetadata() throws Exception { |
| List<Bird> expected = createRandomRecords(DEFAULT_RECORD_COUNT); |
| String codec = DataFileConstants.NULL_CODEC; |
| String filename = |
| generateTestFile( |
| codec, expected, SyncBehavior.SYNC_DEFAULT, 0, AvroCoder.of(Bird.class), codec); |
| Metadata fileMeta = FileSystems.matchSingleFileSpec(filename); |
| |
| AvroSource<GenericRecord> source = AvroSource.from(fileMeta); |
| AvroSource<Bird> sourceWithSchema = source.withSchema(Bird.class); |
| AvroSource<Bird> sourceWithSchemaWithMinBundleSize = sourceWithSchema.withMinBundleSize(1234); |
| |
| assertEquals(FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE, source.getMode()); |
| assertEquals(FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE, sourceWithSchema.getMode()); |
| assertEquals( |
| FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE, sourceWithSchemaWithMinBundleSize.getMode()); |
| } |
| |
| /** |
| * Class that will encode to a fixed size: 16 bytes. |
| * |
| * <p>Each object has a 15-byte array. Avro encodes an object of this type as a byte array, so |
| * each encoded object will consist of 1 byte that encodes the length of the array, followed by 15 |
| * bytes. |
| */ |
| @DefaultCoder(AvroCoder.class) |
| public static class FixedRecord { |
| private byte[] value = new byte[15]; |
| |
| public FixedRecord() { |
| this(0); |
| } |
| |
| public FixedRecord(int i) { |
| value[0] = (byte) i; |
| value[1] = (byte) (i >> 8); |
| value[2] = (byte) (i >> 16); |
| value[3] = (byte) (i >> 24); |
| } |
| |
| public int asInt() { |
| return value[0] | (value[1] << 8) | (value[2] << 16) | (value[3] << 24); |
| } |
| |
| @Override |
| public boolean equals(Object o) { |
| if (o instanceof FixedRecord) { |
| FixedRecord other = (FixedRecord) o; |
| return this.asInt() == other.asInt(); |
| } |
| return false; |
| } |
| |
| @Override |
| public int hashCode() { |
| return toString().hashCode(); |
| } |
| |
| @Override |
| public String toString() { |
| return Integer.toString(this.asInt()); |
| } |
| } |
| |
| /** Create a list of count 16-byte records. */ |
| private static List<FixedRecord> createFixedRecords(int count) { |
| List<FixedRecord> records = new ArrayList<>(); |
| for (int i = 0; i < count; i++) { |
| records.add(new FixedRecord(i)); |
| } |
| return records; |
| } |
| |
| /** Class used as the record type in tests. */ |
| @DefaultCoder(AvroCoder.class) |
| static class Bird { |
| long number; |
| String species; |
| String quality; |
| long quantity; |
| |
| public Bird() {} |
| |
| public Bird(long number, String species, String quality, long quantity) { |
| this.number = number; |
| this.species = species; |
| this.quality = quality; |
| this.quantity = quantity; |
| } |
| |
| @Override |
| public String toString() { |
| return MoreObjects.toStringHelper(Bird.class) |
| .addValue(number) |
| .addValue(species) |
| .addValue(quantity) |
| .addValue(quality) |
| .toString(); |
| } |
| |
| @Override |
| public boolean equals(Object obj) { |
| if (obj instanceof Bird) { |
| Bird other = (Bird) obj; |
| return Objects.equals(species, other.species) |
| && Objects.equals(quality, other.quality) |
| && quantity == other.quantity |
| && number == other.number; |
| } |
| return false; |
| } |
| |
| @Override |
| public int hashCode() { |
| return Objects.hash(number, species, quality, quantity); |
| } |
| } |
| |
| /** |
| * Class used as the record type in tests. |
| * |
| * <p>Contains nullable fields and fields with default values. Can be read using a file written |
| * with the Bird schema. |
| */ |
| @DefaultCoder(AvroCoder.class) |
| public static class FancyBird { |
| long number; |
| String species; |
| String quality; |
| long quantity; |
| |
| @Nullable String habitat; |
| |
| @AvroDefault("\"MAXIMUM OVERDRIVE\"") |
| String fancinessLevel; |
| |
| public FancyBird() {} |
| |
| public FancyBird( |
| long number, |
| String species, |
| String quality, |
| long quantity, |
| String habitat, |
| String fancinessLevel) { |
| this.number = number; |
| this.species = species; |
| this.quality = quality; |
| this.quantity = quantity; |
| this.habitat = habitat; |
| this.fancinessLevel = fancinessLevel; |
| } |
| |
| @Override |
| public String toString() { |
| return MoreObjects.toStringHelper(FancyBird.class) |
| .addValue(number) |
| .addValue(species) |
| .addValue(quality) |
| .addValue(quantity) |
| .addValue(habitat) |
| .addValue(fancinessLevel) |
| .toString(); |
| } |
| |
| @Override |
| public boolean equals(Object obj) { |
| if (obj instanceof FancyBird) { |
| FancyBird other = (FancyBird) obj; |
| return Objects.equals(species, other.species) |
| && Objects.equals(quality, other.quality) |
| && quantity == other.quantity |
| && number == other.number |
| && Objects.equals(fancinessLevel, other.fancinessLevel) |
| && Objects.equals(habitat, other.habitat); |
| } |
| return false; |
| } |
| |
| @Override |
| public int hashCode() { |
| return Objects.hash(number, species, quality, quantity, habitat, fancinessLevel); |
| } |
| } |
| |
| /** Create a list of n random records. */ |
| private static List<Bird> createRandomRecords(long n) { |
| String[] qualities = { |
| "miserable", "forelorn", "fidgity", "squirrelly", "fanciful", "chipper", "lazy" |
| }; |
| String[] species = {"pigeons", "owls", "gulls", "hawks", "robins", "jays"}; |
| Random random = new Random(0); |
| |
| List<Bird> records = new ArrayList<>(); |
| for (long i = 0; i < n; i++) { |
| Bird bird = new Bird(); |
| bird.quality = qualities[random.nextInt(qualities.length)]; |
| bird.species = species[random.nextInt(species.length)]; |
| bird.number = i; |
| bird.quantity = random.nextLong(); |
| records.add(bird); |
| } |
| return records; |
| } |
| } |
