| /* |
| * 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.io.FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE; |
| import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument; |
| import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull; |
| import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState; |
| |
| import java.io.ByteArrayInputStream; |
| import java.io.EOFException; |
| import java.io.IOException; |
| import java.io.InputStream; |
| import java.io.InvalidObjectException; |
| import java.io.ObjectInputStream; |
| import java.io.ObjectStreamException; |
| import java.io.PushbackInputStream; |
| import java.io.Serializable; |
| import java.nio.ByteBuffer; |
| import java.nio.channels.Channels; |
| import java.nio.channels.ReadableByteChannel; |
| import java.nio.channels.SeekableByteChannel; |
| import java.nio.charset.StandardCharsets; |
| 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.Nullable; |
| 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.coders.Coder; |
| import org.apache.beam.sdk.io.fs.EmptyMatchTreatment; |
| 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.options.ValueProvider; |
| import org.apache.beam.sdk.transforms.SerializableFunction; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.base.MoreObjects; |
| 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.8.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.8</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 = 2L * DataFileConstants.DEFAULT_SYNC_INTERVAL; |
| |
| // Use cases of AvroSource are: |
| // 1) AvroSource<GenericRecord> Reading GenericRecord records with a specified schema. |
| // 2) AvroSource<Foo> Reading records of a generated Avro class Foo. |
| // 3) AvroSource<T> Reading GenericRecord records with an unspecified schema |
| // and converting them to type T. |
| // | Case 1 | Case 2 | Case 3 | |
| // type | GenericRecord | Foo | GenericRecord | |
| // readerSchemaString | non-null | non-null | null | |
| // parseFn | null | null | non-null | |
| // outputCoder | null | null | non-null | |
| private static class Mode<T> implements Serializable { |
| private final Class<?> type; |
| |
| // The JSON schema used to decode records. |
| @Nullable private String readerSchemaString; |
| |
| @Nullable private final SerializableFunction<GenericRecord, T> parseFn; |
| |
| @Nullable private final Coder<T> outputCoder; |
| |
| private Mode( |
| Class<?> type, |
| @Nullable String readerSchemaString, |
| @Nullable SerializableFunction<GenericRecord, T> parseFn, |
| @Nullable Coder<T> outputCoder) { |
| this.type = type; |
| this.readerSchemaString = internSchemaString(readerSchemaString); |
| this.parseFn = parseFn; |
| this.outputCoder = outputCoder; |
| } |
| |
| private void readObject(ObjectInputStream is) throws IOException, ClassNotFoundException { |
| is.defaultReadObject(); |
| readerSchemaString = internSchemaString(readerSchemaString); |
| } |
| |
| private Coder<T> getOutputCoder() { |
| if (parseFn == null) { |
| return AvroCoder.of((Class<T>) type, internOrParseSchemaString(readerSchemaString)); |
| } else { |
| return outputCoder; |
| } |
| } |
| |
| private void validate() { |
| if (parseFn == null) { |
| checkArgument( |
| readerSchemaString != null, |
| "schema must be specified using withSchema() when not using a parse fn"); |
| } |
| } |
| } |
| |
| private static Mode<GenericRecord> readGenericRecordsWithSchema(String schema) { |
| return new Mode<>(GenericRecord.class, schema, null, null); |
| } |
| |
| private static <T> Mode<T> readGeneratedClasses(Class<T> clazz) { |
| return new Mode<>(clazz, ReflectData.get().getSchema(clazz).toString(), null, null); |
| } |
| |
| private static <T> Mode<T> parseGenericRecords( |
| SerializableFunction<GenericRecord, T> parseFn, Coder<T> outputCoder) { |
| return new Mode<>(GenericRecord.class, null, parseFn, outputCoder); |
| } |
| |
| private final Mode<T> mode; |
| |
| /** |
| * Reads from the given file name or pattern ("glob"). The returned source needs to be further |
| * configured by calling {@link #withSchema} to return a type other than {@link GenericRecord}. |
| */ |
| public static AvroSource<GenericRecord> from(ValueProvider<String> fileNameOrPattern) { |
| return new AvroSource<>( |
| fileNameOrPattern, |
| EmptyMatchTreatment.DISALLOW, |
| DEFAULT_MIN_BUNDLE_SIZE, |
| readGenericRecordsWithSchema(null /* will need to be specified in withSchema */)); |
| } |
| |
| public static AvroSource<GenericRecord> from(Metadata metadata) { |
| return new AvroSource<>( |
| metadata, |
| DEFAULT_MIN_BUNDLE_SIZE, |
| 0, |
| metadata.sizeBytes(), |
| readGenericRecordsWithSchema(null /* will need to be specified in withSchema */)); |
| } |
| |
| /** Like {@link #from(ValueProvider)}. */ |
| public static AvroSource<GenericRecord> from(String fileNameOrPattern) { |
| return from(ValueProvider.StaticValueProvider.of(fileNameOrPattern)); |
| } |
| |
| public AvroSource<T> withEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment) { |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), emptyMatchTreatment, getMinBundleSize(), mode); |
| } |
| |
| /** Reads files containing records that conform to the given schema. */ |
| public AvroSource<GenericRecord> withSchema(String schema) { |
| checkArgument(schema != null, "schema can not be null"); |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), |
| getEmptyMatchTreatment(), |
| getMinBundleSize(), |
| readGenericRecordsWithSchema(schema)); |
| } |
| |
| /** Like {@link #withSchema(String)}. */ |
| public AvroSource<GenericRecord> withSchema(Schema schema) { |
| checkArgument(schema != null, "schema can not be null"); |
| return withSchema(schema.toString()); |
| } |
| |
| /** Reads files containing records of the given class. */ |
| public <X> AvroSource<X> withSchema(Class<X> clazz) { |
| checkArgument(clazz != null, "clazz can not be null"); |
| if (getMode() == SINGLE_FILE_OR_SUBRANGE) { |
| return new AvroSource<>( |
| getSingleFileMetadata(), |
| getMinBundleSize(), |
| getStartOffset(), |
| getEndOffset(), |
| readGeneratedClasses(clazz)); |
| } |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), |
| getEmptyMatchTreatment(), |
| getMinBundleSize(), |
| readGeneratedClasses(clazz)); |
| } |
| |
| /** |
| * Reads {@link GenericRecord} of unspecified schema and maps them to instances of a custom type |
| * using the given {@code parseFn} and encoded using the given coder. |
| */ |
| public <X> AvroSource<X> withParseFn( |
| SerializableFunction<GenericRecord, X> parseFn, Coder<X> coder) { |
| checkArgument(parseFn != null, "parseFn can not be null"); |
| checkArgument(coder != null, "coder can not be null"); |
| if (getMode() == SINGLE_FILE_OR_SUBRANGE) { |
| return new AvroSource<>( |
| getSingleFileMetadata(), |
| getMinBundleSize(), |
| getStartOffset(), |
| getEndOffset(), |
| parseGenericRecords(parseFn, coder)); |
| } |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), |
| getEmptyMatchTreatment(), |
| getMinBundleSize(), |
| parseGenericRecords(parseFn, coder)); |
| } |
| |
| /** |
| * Sets the minimum bundle size. Refer to {@link OffsetBasedSource} for a description of {@code |
| * minBundleSize} and its use. |
| */ |
| public AvroSource<T> withMinBundleSize(long minBundleSize) { |
| if (getMode() == SINGLE_FILE_OR_SUBRANGE) { |
| return new AvroSource<>( |
| getSingleFileMetadata(), minBundleSize, getStartOffset(), getEndOffset(), mode); |
| } |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), getEmptyMatchTreatment(), minBundleSize, mode); |
| } |
| |
| /** Constructor for FILEPATTERN mode. */ |
| private AvroSource( |
| ValueProvider<String> fileNameOrPattern, |
| EmptyMatchTreatment emptyMatchTreatment, |
| long minBundleSize, |
| Mode<T> mode) { |
| super(fileNameOrPattern, emptyMatchTreatment, minBundleSize); |
| this.mode = mode; |
| } |
| |
| /** Constructor for SINGLE_FILE_OR_SUBRANGE mode. */ |
| private AvroSource( |
| Metadata metadata, long minBundleSize, long startOffset, long endOffset, Mode<T> mode) { |
| super(metadata, minBundleSize, startOffset, endOffset); |
| this.mode = mode; |
| } |
| |
| @Override |
| public void validate() { |
| super.validate(); |
| mode.validate(); |
| } |
| |
| /** |
| * Used by the Dataflow worker. Do not introduce new usages. Do not delete without confirming that |
| * Dataflow ValidatesRunner tests pass. |
| * |
| * @deprecated Used by Dataflow worker |
| */ |
| @Deprecated |
| 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) { |
| return new AvroSource<>(fileMetadata, getMinBundleSize(), start, end, mode); |
| } |
| |
| @Override |
| protected BlockBasedReader<T> createSingleFileReader(PipelineOptions options) { |
| return new AvroReader<>(this); |
| } |
| |
| @Override |
| public Coder<T> getOutputCoder() { |
| return mode.getOutputCoder(); |
| } |
| |
| @VisibleForTesting |
| @Nullable |
| String getReaderSchemaString() { |
| return mode.readerSchemaString; |
| } |
| |
| /** Avro file metadata. */ |
| @VisibleForTesting |
| static class AvroMetadata { |
| private final byte[] syncMarker; |
| private final String codec; |
| private final String schemaString; |
| |
| AvroMetadata(byte[] syncMarker, String codec, String schemaString) { |
| this.syncMarker = checkNotNull(syncMarker, "syncMarker"); |
| this.codec = checkNotNull(codec, "codec"); |
| this.schemaString = internSchemaString(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, StandardCharsets.UTF_8); |
| } else if (key.equals(DataFileConstants.SCHEMA)) { |
| schemaString = new String(bytes, StandardCharsets.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); |
| } |
| |
| // 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(), mode); |
| case FILEPATTERN: |
| return new AvroSource<>( |
| getFileOrPatternSpecProvider(), getEmptyMatchTreatment(), getMinBundleSize(), mode); |
| 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> { |
| private final Mode<T> mode; |
| |
| // The number of records in the block. |
| private final long numRecords; |
| |
| // The current record in the block. Initialized in readNextRecord. |
| @Nullable 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<?> 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, Mode<T> mode, String writerSchemaString, String codec) |
| throws IOException { |
| this.mode = mode; |
| this.numRecords = numRecords; |
| checkNotNull(writerSchemaString, "writerSchemaString"); |
| Schema writerSchema = internOrParseSchemaString(writerSchemaString); |
| Schema readerSchema = |
| internOrParseSchemaString( |
| MoreObjects.firstNonNull(mode.readerSchemaString, writerSchemaString)); |
| this.reader = |
| (mode.type == GenericRecord.class) |
| ? new GenericDatumReader<T>(writerSchema, readerSchema) |
| : new ReflectDatumReader<T>(writerSchema, readerSchema); |
| this.decoder = DecoderFactory.get().binaryDecoder(decodeAsInputStream(data, codec), null); |
| } |
| |
| @Override |
| public T getCurrentRecord() { |
| return currentRecord; |
| } |
| |
| @Override |
| public boolean readNextRecord() throws IOException { |
| if (currentRecordIndex >= numRecords) { |
| return false; |
| } |
| Object record = reader.read(null, decoder); |
| currentRecord = |
| (mode.parseFn == null) ? ((T) record) : mode.parseFn.apply((GenericRecord) record); |
| 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> { |
| // Initialized in startReading. |
| @Nullable private AvroMetadata metadata; |
| |
| // The current block. |
| // Initialized in readNextRecord. |
| @Nullable 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. |
| // Initialized in startReading. |
| @Nullable 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. |
| // Initialized in startReading. |
| @Nullable private CountingInputStream countStream; |
| |
| // Caches the Avro DirectBinaryDecoder used to decode binary-encoded values from the buffer. |
| // Initialized in readNextBlock. |
| @Nullable 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().mode, |
| metadata.getSchemaString(), |
| metadata.getCodec()); |
| |
| // Read the end of this block, which MUST be a sync marker for correctness. |
| byte[] syncMarker = metadata.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), metadata.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 { |
| try { |
| metadata = readMetadataFromFile(getCurrentSource().getSingleFileMetadata().resourceId()); |
| } catch (IOException e) { |
| throw new RuntimeException( |
| "Error reading metadata from file " + getCurrentSource().getSingleFileMetadata(), e); |
| } |
| |
| long startOffset = getCurrentSource().getStartOffset(); |
| byte[] syncMarker = metadata.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; |
| } |
| } |
| } |
| } |
