Java tutorial
/* * 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 com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; import com.google.common.annotations.VisibleForTesting; import java.io.ByteArrayInputStream; import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.io.InvalidObjectException; import java.io.ObjectStreamException; import java.io.PushbackInputStream; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.channels.ReadableByteChannel; import java.nio.channels.SeekableByteChannel; import java.util.Arrays; import java.util.Map; import java.util.WeakHashMap; import java.util.zip.Inflater; import java.util.zip.InflaterInputStream; import javax.annotation.concurrent.GuardedBy; import org.apache.avro.Schema; import org.apache.avro.file.CodecFactory; import org.apache.avro.file.DataFileConstants; import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericRecord; import org.apache.avro.io.BinaryDecoder; import org.apache.avro.io.DatumReader; import org.apache.avro.io.DecoderFactory; import org.apache.avro.reflect.ReflectData; import org.apache.avro.reflect.ReflectDatumReader; import org.apache.beam.sdk.PipelineRunner; import org.apache.beam.sdk.annotations.Experimental; import org.apache.beam.sdk.coders.AvroCoder; import org.apache.beam.sdk.io.fs.MatchResult.Metadata; import org.apache.beam.sdk.io.fs.ResourceId; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.values.PCollection; import org.apache.commons.compress.compressors.bzip2.BZip2CompressorInputStream; import org.apache.commons.compress.compressors.snappy.SnappyCompressorInputStream; import org.apache.commons.compress.compressors.xz.XZCompressorInputStream; import org.apache.commons.compress.utils.CountingInputStream; import org.apache.commons.compress.utils.IOUtils; // CHECKSTYLE.OFF: JavadocStyle /** * Do not use in pipelines directly: most users should use {@link AvroIO.Read}. * * <p>A {@link FileBasedSource} for reading Avro files. * * <p>To read a {@link PCollection} of objects from one or more Avro files, use * {@link AvroSource#from} to specify the path(s) of the files to read. The {@link AvroSource} that * is returned will read objects of type {@link GenericRecord} with the schema(s) that were written * at file creation. To further configure the {@link AvroSource} to read with a user-defined schema, * or to return records of a type other than {@link GenericRecord}, use * {@link AvroSource#withSchema(Schema)} (using an Avro {@link Schema}), * {@link AvroSource#withSchema(String)} (using a JSON schema), or * {@link AvroSource#withSchema(Class)} (to return objects of the Avro-generated class specified). * * <p>An {@link AvroSource} can be read from using the {@link Read} transform. For example: * * <pre> * {@code * AvroSource<MyType> source = AvroSource.from(file.toPath()).withSchema(MyType.class); * PCollection<MyType> records = Read.from(mySource); * } * </pre> * * <p>This class's implementation is based on the <a * href="https://avro.apache.org/docs/1.7.7/spec.html">Avro 1.7.7</a> specification and implements * parsing of some parts of Avro Object Container Files. The rationale for doing so is that the Avro * API does not provide efficient ways of computing the precise offsets of blocks within a file, * which is necessary to support dynamic work rebalancing. However, whenever it is possible to use * the Avro API in a way that supports maintaining precise offsets, this class uses the Avro API. * * <p>Avro Object Container files store records in blocks. Each block contains a collection of * records. Blocks may be encoded (e.g., with bzip2, deflate, snappy, etc.). Blocks are delineated * from one another by a 16-byte sync marker. * * <p>An {@link AvroSource} for a subrange of a single file contains records in the blocks such that * the start offset of the block is greater than or equal to the start offset of the source and less * than the end offset of the source. * * <p>To use XZ-encoded Avro files, please include an explicit dependency on {@code xz-1.5.jar}, * which has been marked as optional in the Maven {@code sdk/pom.xml}. * * <pre>{@code * <dependency> * <groupId>org.tukaani</groupId> * <artifactId>xz</artifactId> * <version>1.5</version> * </dependency> * }</pre> * * <h3>Permissions</h3> * * <p>Permission requirements depend on the {@link PipelineRunner} that is used to execute the * pipeline. Please refer to the documentation of corresponding {@link PipelineRunner}s for * more details. * * @param <T> The type of records to be read from the source. */ // CHECKSTYLE.ON: JavadocStyle @Experimental(Experimental.Kind.SOURCE_SINK) public class AvroSource<T> extends BlockBasedSource<T> { // Default minimum bundle size (chosen as two default-size Avro blocks to attempt to // ensure that every source has at least one block of records). // The default sync interval is 64k. private static final long DEFAULT_MIN_BUNDLE_SIZE = 2 * DataFileConstants.DEFAULT_SYNC_INTERVAL; // The JSON schema used to encode records. private final String readSchemaString; // The JSON schema that was used to write the source Avro file (may differ from the schema we will // use to read from it). private final String fileSchemaString; // The type of the records contained in the file. private final Class<T> type; // The following metadata fields are not user-configurable. They are extracted from the object // container file header upon subsource creation. // The codec used to encode the blocks in the Avro file. String value drawn from those in // https://avro.apache.org/docs/1.7.7/api/java/org/apache/avro/file/CodecFactory.html private final String codec; // The object container file's 16-byte sync marker. private final byte[] syncMarker; // Default output coder, lazily initialized. private transient AvroCoder<T> coder = null; // Schema of the file, lazily initialized. private transient Schema fileSchema; // Schema used to encode records, lazily initialized. private transient Schema readSchema; /** * Creates an {@link AvroSource} that reads from the given file name or pattern ("glob"). The * returned source can be further configured by calling {@link #withSchema} to return a type other * than {@link GenericRecord}. */ public static AvroSource<GenericRecord> from(String fileNameOrPattern) { return new AvroSource<>(fileNameOrPattern, DEFAULT_MIN_BUNDLE_SIZE, null, GenericRecord.class, null, null); } /** * Returns an {@link AvroSource} that's like this one but reads files containing records that * conform to the given schema. * * <p>Does not modify this object. */ public AvroSource<GenericRecord> withSchema(String schema) { return new AvroSource<>(getFileOrPatternSpec(), getMinBundleSize(), schema, GenericRecord.class, codec, syncMarker); } /** * Returns an {@link AvroSource} that's like this one but reads files containing records that * conform to the given schema. * * <p>Does not modify this object. */ public AvroSource<GenericRecord> withSchema(Schema schema) { return new AvroSource<>(getFileOrPatternSpec(), getMinBundleSize(), schema.toString(), GenericRecord.class, codec, syncMarker); } /** * Returns an {@link AvroSource} that's like this one but reads files containing records of the * type of the given class. * * <p>Does not modify this object. */ public <X> AvroSource<X> withSchema(Class<X> clazz) { return new AvroSource<X>(getFileOrPatternSpec(), getMinBundleSize(), ReflectData.get().getSchema(clazz).toString(), clazz, codec, syncMarker); } /** * Returns an {@link AvroSource} that's like this one but uses the supplied minimum bundle size. * Refer to {@link OffsetBasedSource} for a description of {@code minBundleSize} and its use. * * <p>Does not modify this object. */ public AvroSource<T> withMinBundleSize(long minBundleSize) { return new AvroSource<>(getFileOrPatternSpec(), minBundleSize, readSchemaString, type, codec, syncMarker); } private AvroSource(String fileNameOrPattern, long minBundleSize, String schema, Class<T> type, String codec, byte[] syncMarker) { super(fileNameOrPattern, minBundleSize); this.readSchemaString = internSchemaString(schema); this.codec = codec; this.syncMarker = syncMarker; this.type = type; this.fileSchemaString = null; } private AvroSource(Metadata metadata, long minBundleSize, long startOffset, long endOffset, String schema, Class<T> type, String codec, byte[] syncMarker, String fileSchema) { super(metadata, minBundleSize, startOffset, endOffset); this.readSchemaString = internSchemaString(schema); this.codec = codec; this.syncMarker = syncMarker; this.type = type; this.fileSchemaString = internSchemaString(fileSchema); } @Override public void validate() { // AvroSource objects do not need to be configured with more than a file pattern. Overridden to // make this explicit. super.validate(); } @Deprecated // Added to let DataflowRunner migrate off of this; to be deleted. public BlockBasedSource<T> createForSubrangeOfFile(String fileName, long start, long end) throws IOException { return createForSubrangeOfFile(FileSystems.matchSingleFileSpec(fileName), start, end); } @Override public BlockBasedSource<T> createForSubrangeOfFile(Metadata fileMetadata, long start, long end) { byte[] syncMarker = this.syncMarker; String codec = this.codec; String readSchemaString = this.readSchemaString; String fileSchemaString = this.fileSchemaString; // codec and syncMarker are initially null when the source is created, as they differ // across input files and must be read from the file. Here, when we are creating a source // for a subrange of a file, we can initialize these values. When the resulting AvroSource // is further split, they do not need to be read again. if (codec == null || syncMarker == null || fileSchemaString == null) { AvroMetadata metadata; try { metadata = readMetadataFromFile(fileMetadata.resourceId()); } catch (IOException e) { throw new RuntimeException("Error reading metadata from file " + fileMetadata, e); } codec = metadata.getCodec(); syncMarker = metadata.getSyncMarker(); fileSchemaString = metadata.getSchemaString(); // If the source was created with a null schema, use the schema that we read from the file's // metadata. if (readSchemaString == null) { readSchemaString = metadata.getSchemaString(); } } // Note that if the fileSchemaString is equivalent to the readSchemaString, "intern"ing // the string will occur within the constructor and return the same reference as the // readSchemaString. This allows for Java to have an efficient serialization since it // will only encode the schema once while just storing pointers to the encoded version // within this source. return new AvroSource<>(fileMetadata, getMinBundleSize(), start, end, readSchemaString, type, codec, syncMarker, fileSchemaString); } @Override protected BlockBasedReader<T> createSingleFileReader(PipelineOptions options) { return new AvroReader<>(this); } @Override public AvroCoder<T> getDefaultOutputCoder() { if (coder == null) { coder = AvroCoder.of(type, internOrParseSchemaString(readSchemaString)); } return coder; } public String getSchema() { return readSchemaString; } @VisibleForTesting Schema getReadSchema() { if (readSchemaString == null) { return null; } // If the schema has not been parsed, parse it. if (readSchema == null) { readSchema = internOrParseSchemaString(readSchemaString); } return readSchema; } @VisibleForTesting Schema getFileSchema() { if (fileSchemaString == null) { return null; } // If the schema has not been parsed, parse it. if (fileSchema == null) { fileSchema = internOrParseSchemaString(fileSchemaString); } return fileSchema; } private byte[] getSyncMarker() { return syncMarker; } private String getCodec() { return codec; } /** * Avro file metadata. */ @VisibleForTesting static class AvroMetadata { private byte[] syncMarker; private String codec; private String schemaString; AvroMetadata(byte[] syncMarker, String codec, String schemaString) { this.syncMarker = checkNotNull(syncMarker, "syncMarker"); this.codec = checkNotNull(codec, "codec"); this.schemaString = checkNotNull(schemaString, "schemaString"); } /** * The JSON-encoded <a href="https://avro.apache.org/docs/1.7.7/spec.html#schemas">schema</a> * string for the file. */ public String getSchemaString() { return schemaString; } /** * The <a href="https://avro.apache.org/docs/1.7.7/spec.html#Required+Codecs">codec</a> of the * file. */ public String getCodec() { return codec; } /** * The 16-byte sync marker for the file. See the documentation for * <a href="https://avro.apache.org/docs/1.7.7/spec.html#Object+Container+Files">Object * Container File</a> for more information. */ public byte[] getSyncMarker() { return syncMarker; } } /** * Reads the {@link AvroMetadata} from the header of an Avro file. * * <p>This method parses the header of an Avro * <a href="https://avro.apache.org/docs/1.7.7/spec.html#Object+Container+Files"> * Object Container File</a>. * * @throws IOException if the file is an invalid format. */ @VisibleForTesting static AvroMetadata readMetadataFromFile(ResourceId fileResource) throws IOException { String codec = null; String schemaString = null; byte[] syncMarker; try (InputStream stream = Channels.newInputStream(FileSystems.open(fileResource))) { BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(stream, null); // The header of an object container file begins with a four-byte magic number, followed // by the file metadata (including the schema and codec), encoded as a map. Finally, the // header ends with the file's 16-byte sync marker. // See https://avro.apache.org/docs/1.7.7/spec.html#Object+Container+Files for details on // the encoding of container files. // Read the magic number. byte[] magic = new byte[DataFileConstants.MAGIC.length]; decoder.readFixed(magic); if (!Arrays.equals(magic, DataFileConstants.MAGIC)) { throw new IOException("Missing Avro file signature: " + fileResource); } // Read the metadata to find the codec and schema. ByteBuffer valueBuffer = ByteBuffer.allocate(512); long numRecords = decoder.readMapStart(); while (numRecords > 0) { for (long recordIndex = 0; recordIndex < numRecords; recordIndex++) { String key = decoder.readString(); // readBytes() clears the buffer and returns a buffer where: // - position is the start of the bytes read // - limit is the end of the bytes read valueBuffer = decoder.readBytes(valueBuffer); byte[] bytes = new byte[valueBuffer.remaining()]; valueBuffer.get(bytes); if (key.equals(DataFileConstants.CODEC)) { codec = new String(bytes, "UTF-8"); } else if (key.equals(DataFileConstants.SCHEMA)) { schemaString = new String(bytes, "UTF-8"); } } numRecords = decoder.mapNext(); } if (codec == null) { codec = DataFileConstants.NULL_CODEC; } // Finally, read the sync marker. syncMarker = new byte[DataFileConstants.SYNC_SIZE]; decoder.readFixed(syncMarker); } checkState(schemaString != null, "No schema present in Avro file metadata %s", fileResource); return new AvroMetadata(syncMarker, codec, schemaString); } private DatumReader<T> createDatumReader() { Schema readSchema = getReadSchema(); Schema fileSchema = getFileSchema(); checkNotNull(readSchema, "No read schema has been initialized for source %s", this); checkNotNull(fileSchema, "No file schema has been initialized for source %s", this); if (type == GenericRecord.class) { return new GenericDatumReader<>(fileSchema, readSchema); } else { return new ReflectDatumReader<>(fileSchema, readSchema); } } // A logical reference cache used to store schemas and schema strings to allow us to // "intern" values and reduce the number of copies of equivalent objects. private static final Map<String, Schema> schemaLogicalReferenceCache = new WeakHashMap<>(); private static final Map<String, String> schemaStringLogicalReferenceCache = new WeakHashMap<>(); // We avoid String.intern() because depending on the JVM, these may be added to the PermGenSpace // which we want to avoid otherwise we could run out of PermGenSpace. private static synchronized String internSchemaString(String schema) { String internSchema = schemaStringLogicalReferenceCache.get(schema); if (internSchema != null) { return internSchema; } schemaStringLogicalReferenceCache.put(schema, schema); return schema; } private static synchronized Schema internOrParseSchemaString(String schemaString) { Schema schema = schemaLogicalReferenceCache.get(schemaString); if (schema != null) { return schema; } Schema.Parser parser = new Schema.Parser(); schema = parser.parse(schemaString); schemaLogicalReferenceCache.put(schemaString, schema); return schema; } // Reading the object from Java serialization typically does not go through the constructor, // we use readResolve to replace the constructed instance with one which uses the constructor // allowing us to intern any schemas. @SuppressWarnings("unused") private Object readResolve() throws ObjectStreamException { switch (getMode()) { case SINGLE_FILE_OR_SUBRANGE: return new AvroSource<>(getSingleFileMetadata(), getMinBundleSize(), getStartOffset(), getEndOffset(), readSchemaString, type, codec, syncMarker, fileSchemaString); case FILEPATTERN: return new AvroSource<>(getFileOrPatternSpec(), getMinBundleSize(), readSchemaString, type, codec, syncMarker); default: throw new InvalidObjectException(String.format("Unknown mode %s for AvroSource %s", getMode(), this)); } } /** * A {@link BlockBasedSource.Block} of Avro records. * * @param <T> The type of records stored in the block. */ @Experimental(Experimental.Kind.SOURCE_SINK) static class AvroBlock<T> extends Block<T> { // The number of records in the block. private final long numRecords; // The current record in the block. private T currentRecord; // The index of the current record in the block. private long currentRecordIndex = 0; // A DatumReader to read records from the block. private final DatumReader<T> reader; // A BinaryDecoder used by the reader to decode records. private final BinaryDecoder decoder; /** * Decodes a byte array as an InputStream. The byte array may be compressed using some * codec. Reads from the returned stream will result in decompressed bytes. * * <p>This supports the same codecs as Avro's {@link CodecFactory}, namely those defined in * {@link DataFileConstants}. * * <ul> * <li>"snappy" : Google's Snappy compression * <li>"deflate" : deflate compression * <li>"bzip2" : Bzip2 compression * <li>"xz" : xz compression * <li>"null" (the string, not the value): Uncompressed data * </ul> */ private static InputStream decodeAsInputStream(byte[] data, String codec) throws IOException { ByteArrayInputStream byteStream = new ByteArrayInputStream(data); switch (codec) { case DataFileConstants.SNAPPY_CODEC: return new SnappyCompressorInputStream(byteStream, 1 << 16 /* Avro uses 64KB blocks */); case DataFileConstants.DEFLATE_CODEC: // nowrap == true: Do not expect ZLIB header or checksum, as Avro does not write them. Inflater inflater = new Inflater(true); return new InflaterInputStream(byteStream, inflater); case DataFileConstants.XZ_CODEC: return new XZCompressorInputStream(byteStream); case DataFileConstants.BZIP2_CODEC: return new BZip2CompressorInputStream(byteStream); case DataFileConstants.NULL_CODEC: return byteStream; default: throw new IllegalArgumentException("Unsupported codec: " + codec); } } AvroBlock(byte[] data, long numRecords, AvroSource<T> source) throws IOException { this.numRecords = numRecords; this.reader = source.createDatumReader(); this.decoder = DecoderFactory.get().binaryDecoder(decodeAsInputStream(data, source.getCodec()), null); } @Override public T getCurrentRecord() { return currentRecord; } @Override public boolean readNextRecord() throws IOException { if (currentRecordIndex >= numRecords) { return false; } currentRecord = reader.read(null, decoder); currentRecordIndex++; return true; } @Override public double getFractionOfBlockConsumed() { return ((double) currentRecordIndex) / numRecords; } } /** * A {@link BlockBasedSource.BlockBasedReader} for reading blocks from Avro files. * * <p>An Avro Object Container File consists of a header followed by a 16-bit sync marker * and then a sequence of blocks, where each block begins with two encoded longs representing * the total number of records in the block and the block's size in bytes, followed by the * block's (optionally-encoded) records. Each block is terminated by a 16-bit sync marker. * * @param <T> The type of records contained in the block. */ @Experimental(Experimental.Kind.SOURCE_SINK) public static class AvroReader<T> extends BlockBasedReader<T> { // The current block. private AvroBlock<T> currentBlock; // A lock used to synchronize block offsets for getRemainingParallelism private final Object progressLock = new Object(); // Offset of the current block. @GuardedBy("progressLock") private long currentBlockOffset = 0; // Size of the current block. @GuardedBy("progressLock") private long currentBlockSizeBytes = 0; // Stream used to read from the underlying file. // A pushback stream is used to restore bytes buffered during seeking. private PushbackInputStream stream; // Counts the number of bytes read. Used only to tell how many bytes are taken up in // a block's variable-length header. private CountingInputStream countStream; // Caches the Avro DirectBinaryDecoder used to decode binary-encoded values from the buffer. private BinaryDecoder decoder; /** * Reads Avro records of type {@code T} from the specified source. */ public AvroReader(AvroSource<T> source) { super(source); } @Override public synchronized AvroSource<T> getCurrentSource() { return (AvroSource<T>) super.getCurrentSource(); } // Precondition: the stream is positioned after the sync marker in the current (about to be // previous) block. currentBlockSize equals the size of the current block, or zero if this // reader was just started. // // Postcondition: same as above, but for the new current (formerly next) block. @Override public boolean readNextBlock() throws IOException { long startOfNextBlock; synchronized (progressLock) { startOfNextBlock = currentBlockOffset + currentBlockSizeBytes; } // Before reading the variable-sized block header, record the current number of bytes read. long preHeaderCount = countStream.getBytesRead(); decoder = DecoderFactory.get().directBinaryDecoder(countStream, decoder); long numRecords; try { numRecords = decoder.readLong(); } catch (EOFException e) { // Expected for the last block, at which the start position is the EOF. The way to detect // stream ending is to try reading from it. return false; } long blockSize = decoder.readLong(); // Mark header size as the change in the number of bytes read. long headerSize = countStream.getBytesRead() - preHeaderCount; // Create the current block by reading blockSize bytes. Block sizes permitted by the Avro // specification are [32, 2^30], so the cast is safe. byte[] data = new byte[(int) blockSize]; int bytesRead = IOUtils.readFully(stream, data); checkState(blockSize == bytesRead, "Only able to read %s/%s bytes in the block before EOF reached.", bytesRead, blockSize); currentBlock = new AvroBlock<>(data, numRecords, getCurrentSource()); // Read the end of this block, which MUST be a sync marker for correctness. byte[] syncMarker = getCurrentSource().getSyncMarker(); byte[] readSyncMarker = new byte[syncMarker.length]; long syncMarkerOffset = startOfNextBlock + headerSize + blockSize; bytesRead = IOUtils.readFully(stream, readSyncMarker); checkState(bytesRead == syncMarker.length, "Only able to read %s/%s bytes of Avro sync marker at position %s before EOF reached.", bytesRead, syncMarker.length, syncMarkerOffset); if (!Arrays.equals(syncMarker, readSyncMarker)) { throw new IllegalStateException( String.format("Expected the bytes [%d,%d) in file %s to be a sync marker, but found %s", syncMarkerOffset, syncMarkerOffset + syncMarker.length, getCurrentSource().getFileOrPatternSpec(), Arrays.toString(readSyncMarker))); } // Atomically update both the position and offset of the new block. synchronized (progressLock) { currentBlockOffset = startOfNextBlock; // Total block size includes the header, block content, and trailing sync marker. currentBlockSizeBytes = headerSize + blockSize + syncMarker.length; } return true; } @Override public AvroBlock<T> getCurrentBlock() { return currentBlock; } @Override public long getCurrentBlockOffset() { synchronized (progressLock) { return currentBlockOffset; } } @Override public long getCurrentBlockSize() { synchronized (progressLock) { return currentBlockSizeBytes; } } @Override public long getSplitPointsRemaining() { if (isDone()) { return 0; } synchronized (progressLock) { if (currentBlockOffset + currentBlockSizeBytes >= getCurrentSource().getEndOffset()) { // This block is known to be the last block in the range. return 1; } } return super.getSplitPointsRemaining(); } /** * Creates a {@link PushbackInputStream} that has a large enough pushback buffer to be able * to push back the syncBuffer. */ private PushbackInputStream createStream(ReadableByteChannel channel) { return new PushbackInputStream(Channels.newInputStream(channel), getCurrentSource().getSyncMarker().length); } // Postcondition: the stream is positioned at the beginning of the first block after the start // of the current source, and currentBlockOffset is that position. Additionally, // currentBlockSizeBytes will be set to 0 indicating that the previous block was empty. @Override protected void startReading(ReadableByteChannel channel) throws IOException { long startOffset = getCurrentSource().getStartOffset(); byte[] syncMarker = getCurrentSource().getSyncMarker(); long syncMarkerLength = syncMarker.length; if (startOffset != 0) { // Rewind order to find the sync marker ending the previous block. long position = Math.max(0, startOffset - syncMarkerLength); ((SeekableByteChannel) channel).position(position); startOffset = position; } // Satisfy the post condition. stream = createStream(channel); countStream = new CountingInputStream(stream); synchronized (progressLock) { currentBlockOffset = startOffset + advancePastNextSyncMarker(stream, syncMarker); currentBlockSizeBytes = 0; } } /** * Advances to the first byte after the next occurrence of the sync marker in the * stream when reading from the current offset. Returns the number of bytes consumed * from the stream. Note that this method requires a PushbackInputStream with a buffer * at least as big as the marker it is seeking for. */ static long advancePastNextSyncMarker(PushbackInputStream stream, byte[] syncMarker) throws IOException { Seeker seeker = new Seeker(syncMarker); byte[] syncBuffer = new byte[syncMarker.length]; long totalBytesConsumed = 0; // Seek until either a sync marker is found or we reach the end of the file. int mark = -1; // Position of the last byte in the sync marker. int read; // Number of bytes read. do { read = stream.read(syncBuffer); if (read >= 0) { mark = seeker.find(syncBuffer, read); // Update the currentOffset with the number of bytes read. totalBytesConsumed += read; } } while (mark < 0 && read > 0); // If the sync marker was found, unread block data and update the current offsets. if (mark >= 0) { // The current offset after this call should be just past the sync marker, so we should // unread the remaining buffer contents and update the currentOffset accordingly. stream.unread(syncBuffer, mark + 1, read - (mark + 1)); totalBytesConsumed = totalBytesConsumed - (read - (mark + 1)); } return totalBytesConsumed; } /** * A {@link Seeker} looks for a given marker within a byte buffer. Uses naive string matching * with a sliding window, as sync markers are small and random. */ static class Seeker { // The marker to search for. private byte[] marker; // Buffer used for the sliding window. private byte[] searchBuffer; // Number of bytes available to be matched in the buffer. private int available = 0; /** * Create a {@link Seeker} that looks for the given marker. */ public Seeker(byte[] marker) { this.marker = marker; this.searchBuffer = new byte[marker.length]; } /** * Find the marker in the byte buffer. Returns the index of the end of the marker in the * buffer. If the marker is not found, returns -1. * * <p>State is maintained between calls. If the marker was partially matched, a subsequent * call to find will resume matching the marker. * * @param buffer * @return the index of the end of the marker within the buffer, or -1 if the buffer was not * found. */ public int find(byte[] buffer, int length) { for (int i = 0; i < length; i++) { System.arraycopy(searchBuffer, 1, searchBuffer, 0, searchBuffer.length - 1); searchBuffer[searchBuffer.length - 1] = buffer[i]; available = Math.min(available + 1, searchBuffer.length); if (ByteBuffer.wrap(searchBuffer, searchBuffer.length - available, available) .equals(ByteBuffer.wrap(marker))) { available = 0; return i; } } return -1; } } } }