content: /*
   * 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.extensions.avro.io;
  
  import static org.apache.beam.sdk.io.FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.InvalidObjectException;
  import java.io.ObjectInputStream;
  import java.io.ObjectStreamException;
  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.Iterator;
  import java.util.Map;
  import java.util.Optional;
  import java.util.WeakHashMap;
  import javax.annotation.concurrent.GuardedBy;
  import org.apache.avro.Schema;
  import org.apache.avro.file.DataFileConstants;
  import org.apache.avro.file.DataFileReader;
  import org.apache.avro.file.SeekableInput;
  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.beam.sdk.PipelineRunner;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
  import org.apache.beam.sdk.io.BlockBasedSource;
  import org.apache.beam.sdk.io.FileBasedSource;
  import org.apache.beam.sdk.io.FileSystems;
  import org.apache.beam.sdk.io.OffsetBasedSource;
  import org.apache.beam.sdk.io.Read;
  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.util.VarInt;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  // 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
  
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  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;
  
    @FunctionalInterface
    public interface DatumReaderFactory<T> extends Serializable {
      DatumReader<T> apply(Schema writer, Schema reader);
    }
  
    // 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         |     either    |    either   |   non-null    |
    // readerFactory       |     either    |    either   |    either     |
    private static class Mode<T> implements Serializable {
      private final Class<?> type;
  
      // The JSON schema used to decode records.
      private @Nullable String readerSchemaString;
  
      private final @Nullable SerializableFunction<GenericRecord, T> parseFn;
  
      private final @Nullable Coder<T> outputCoder;
  
      private final @Nullable DatumReaderFactory<?> readerFactory;
  
      private Mode() {
        /* Empty Mode. Will need to be configured using withSchema/withClass/withParseFunction */
        this(GenericRecord.class, null, null, null, null);
      }
  
      private Mode(
          Class<?> type,
          @Nullable String readerSchemaString,
          @Nullable SerializableFunction<GenericRecord, T> parseFn,
          @Nullable Coder<T> outputCoder,
          @Nullable DatumReaderFactory<?> readerFactory) {
        this.type = type;
        this.readerSchemaString = internSchemaString(readerSchemaString);
        this.parseFn = parseFn;
        this.outputCoder = outputCoder;
        this.readerFactory = readerFactory;
      }
  
      private void readObject(ObjectInputStream is) throws IOException, ClassNotFoundException {
        is.defaultReadObject();
        readerSchemaString = internSchemaString(readerSchemaString);
      }
  
      private Coder<T> getOutputCoder() {
        if (parseFn == null && outputCoder == null) {
          if (readerFactory != null && readerFactory instanceof AvroDatumFactory) {
            // create custom AvroCoder from the AvroDatumFactory
            return AvroCoder.of(
                (AvroDatumFactory<T>) readerFactory, internOrParseSchemaString(readerSchemaString));
          } else {
            // fallback with default avro coder for the type & schema
            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 Mode<GenericRecord> withSchema(String schema) {
        return new Mode<>(
            GenericRecord.class, schema, null, (Coder<GenericRecord>) outputCoder, readerFactory);
      }
  
      private <X> Mode<X> withClass(Class<X> clazz) {
        return new Mode<>(
            clazz,
            ReflectData.get().getSchema(clazz).toString(),
            null,
            (Coder<X>) outputCoder,
            readerFactory);
      }
  
      private <X> Mode<X> withParseFunction(
          SerializableFunction<GenericRecord, X> parseFn, Coder<X> outputCoder) {
        return new Mode<>(GenericRecord.class, null, parseFn, outputCoder, readerFactory);
      }
  
      private Mode<T> withCoder(Coder<T> coder) {
        return new Mode<>(type, readerSchemaString, parseFn, coder, readerFactory);
      }
  
      private Mode<T> withReaderFactory(DatumReaderFactory<?> factory) {
        return new Mode<>(type, readerSchemaString, parseFn, outputCoder, factory);
      }
    }
  
    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, new Mode<>());
    }
  
    public static AvroSource<GenericRecord> from(Metadata metadata) {
      return new AvroSource<>(
          metadata, DEFAULT_MIN_BUNDLE_SIZE, 0, metadata.sizeBytes(), new Mode<>());
    }
  
    /** 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");
      if (getMode() == SINGLE_FILE_OR_SUBRANGE) {
        return new AvroSource<>(
            getSingleFileMetadata(),
            getMinBundleSize(),
            getStartOffset(),
            getEndOffset(),
            mode.withSchema(schema));
      }
      return new AvroSource<>(
          getFileOrPatternSpecProvider(),
          getEmptyMatchTreatment(),
          getMinBundleSize(),
          mode.withSchema(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(),
            mode.withClass(clazz));
      }
      return new AvroSource<>(
          getFileOrPatternSpecProvider(),
          getEmptyMatchTreatment(),
          getMinBundleSize(),
          mode.withClass(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(),
            mode.withParseFunction(parseFn, coder));
      }
      return new AvroSource<>(
          getFileOrPatternSpecProvider(),
          getEmptyMatchTreatment(),
          getMinBundleSize(),
          mode.withParseFunction(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);
    }
  
    /**
     * Sets a custom {@link AvroSource.DatumReaderFactory} for reading. Pass a {@link
     * AvroDatumFactory} to also use the factory for the {@link AvroCoder}
     */
    public AvroSource<T> withDatumReaderFactory(DatumReaderFactory<?> factory) {
      checkArgument(factory != null, "factory can not be null");
      Mode<T> newMode = mode.withReaderFactory(factory);
      if (getMode() == SINGLE_FILE_OR_SUBRANGE) {
        return new AvroSource<>(
            getSingleFileMetadata(), getMinBundleSize(), getStartOffset(), getEndOffset(), newMode);
      }
      return new AvroSource<>(
          getFileOrPatternSpecProvider(), getEmptyMatchTreatment(), getMinBundleSize(), newMode);
    }
  
    /** Specifies the coder for the result of the {@code AvroSource}. */
    public AvroSource<T> withCoder(Coder<T> coder) {
      checkArgument(coder != null, "coder can not be null");
      Mode<T> newMode = mode.withCoder(coder);
      if (getMode() == SINGLE_FILE_OR_SUBRANGE) {
        return new AvroSource<>(
            getSingleFileMetadata(), getMinBundleSize(), getStartOffset(), getEndOffset(), newMode);
      }
      return new AvroSource<>(
          getFileOrPatternSpecProvider(), getEmptyMatchTreatment(), getMinBundleSize(), newMode);
    }
  
    /** 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;
    }
  
    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 Block} of Avro records.
     *
     * @param <T> The type of records stored in the block.
     */
    static class AvroBlock<T> extends Block<T> {
  
      // The current record in the block. Initialized in readNextRecord.
      private @Nullable T currentRecord;
  
      // The index of the current record in the block.
      private long currentRecordIndex = 0;
  
      private final Iterator<?> iterator;
  
      private final SerializableFunction<GenericRecord, T> parseFn;
  
      private final long numRecordsInBlock;
  
      AvroBlock(
          Iterator<?> iter, SerializableFunction<GenericRecord, T> parseFn, long numRecordsInBlock) {
        this.iterator = iter;
        this.parseFn = parseFn;
        this.numRecordsInBlock = numRecordsInBlock;
      }
  
      @Override
      public T getCurrentRecord() {
        return currentRecord;
      }
  
      @Override
      public boolean readNextRecord() {
        if (currentRecordIndex >= numRecordsInBlock) {
          return false;
        }
  
        Object record = iterator.next();
        currentRecord = (parseFn == null) ? ((T) record) : parseFn.apply((GenericRecord) record);
        currentRecordIndex++;
        return true;
      }
  
      @Override
      public double getFractionOfBlockConsumed() {
        return ((double) currentRecordIndex) / numRecordsInBlock;
      }
    }
  
    /**
     * A {@link 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.
     */
    public static class AvroReader<T> extends BlockBasedReader<T> {
  
      private static class SeekableChannelInput implements SeekableInput {
  
        private final SeekableByteChannel channel;
        private final InputStream input;
  
        SeekableChannelInput(SeekableByteChannel channel) {
          this.channel = channel;
          this.input = Channels.newInputStream(channel);
        }
  
        @Override
        public void seek(long p) throws IOException {
          channel.position(p);
        }
  
        @Override
        public long tell() throws IOException {
          return channel.position();
        }
  
        @Override
        public long length() throws IOException {
          return channel.size();
        }
  
        @Override
        public int read(byte[] b, int off, int len) throws IOException {
          return input.read(b, off, len);
        }
  
        @Override
        public void close() throws IOException {
          channel.close();
        }
      }
  
      // The current block.
      // Initialized in readNextRecord.
      private @Nullable AvroBlock<T> currentBlock;
  
      private @Nullable DataFileReader<?> dataFileReader;
  
      // 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;
  
      /** 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() {
        if (!dataFileReader.hasNext()) {
          return false;
        }
  
        long headerLength =
            (long) VarInt.getLength(dataFileReader.getBlockCount())
                + VarInt.getLength(dataFileReader.getBlockSize())
                + DataFileConstants.SYNC_SIZE;
  
        currentBlock =
            new AvroBlock<>(
                dataFileReader, getCurrentSource().mode.parseFn, dataFileReader.getBlockCount());
  
        // Atomically update both the position and offset of the new block.
        synchronized (progressLock) {
          currentBlockOffset = dataFileReader.previousSync();
          // Total block size includes the header, block content, and trailing sync marker.
          currentBlockSizeBytes = dataFileReader.getBlockSize() + headerLength;
        }
  
        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();
      }
  
      // 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 {
        SeekableChannelInput seekableChannelInput =
            new SeekableChannelInput((SeekableByteChannel) channel);
        // the channel needs to be at the beginning of the file in order for the DataFileReader to
        // read the header, etc, we'll seek it back to where it should be after creating the DFR.
        seekableChannelInput.seek(0);
  
        Schema readerSchema = null;
        String readerSchemaString = this.getCurrentSource().getReaderSchemaString();
        if (readerSchemaString != null) {
          readerSchema = AvroSource.internOrParseSchemaString(readerSchemaString);
        }
  
        DatumReader<?> reader =
            Optional.<DatumReaderFactory<?>>ofNullable(this.getCurrentSource().mode.readerFactory)
                .orElse(AvroDatumFactory.of(this.getCurrentSource().mode.type))
                .apply(readerSchema, readerSchema);
  
        dataFileReader = new DataFileReader<>(seekableChannelInput, reader);
  
        long startOffset = getCurrentSource().getStartOffset();
        if (startOffset != 0) {
          // the start offset may be in the middle of a sync marker, by rewinding SYNC_SIZE bytes we
          // ensure that we won't miss the block if so.
          dataFileReader.sync(Math.max(0, startOffset - DataFileConstants.SYNC_SIZE));
        }
  
        synchronized (progressLock) {
          currentBlockOffset = dataFileReader.previousSync();
          currentBlockSizeBytes = 0;
        }
      }
    }
  }
  