| /* |
| * 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 |
| * |
| * https://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.avro.mapreduce; |
| |
| import java.io.IOException; |
| |
| import org.apache.avro.Schema; |
| import org.apache.avro.file.DataFileReader; |
| import org.apache.avro.file.SeekableInput; |
| import org.apache.avro.generic.GenericData; |
| import org.apache.avro.hadoop.io.AvroSerialization; |
| import org.apache.avro.io.DatumReader; |
| import org.apache.avro.mapred.FsInput; |
| import org.apache.hadoop.conf.Configuration; |
| import org.apache.hadoop.fs.Path; |
| import org.apache.hadoop.mapreduce.InputSplit; |
| import org.apache.hadoop.mapreduce.RecordReader; |
| import org.apache.hadoop.mapreduce.TaskAttemptContext; |
| import org.apache.hadoop.mapreduce.lib.input.FileSplit; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Abstract base class for <code>RecordReader</code>s that read Avro container |
| * files. |
| * |
| * @param <K> The type of key the record reader should generate. |
| * @param <V> The type of value the record reader should generate. |
| * @param <T> The type of the entries within the Avro container file being read. |
| */ |
| public abstract class AvroRecordReaderBase<K, V, T> extends RecordReader<K, V> { |
| private static final Logger LOG = LoggerFactory.getLogger(AvroRecordReaderBase.class); |
| |
| /** The reader schema for the records within the input Avro container file. */ |
| private final Schema mReaderSchema; |
| |
| /** The current record from the Avro container file being read. */ |
| private T mCurrentRecord; |
| |
| /** A reader for the Avro container file containing the current input split. */ |
| private DataFileReader<T> mAvroFileReader; |
| |
| /** |
| * The byte offset into the Avro container file where the first block that fits |
| * completely within the current input split begins. |
| */ |
| private long mStartPosition; |
| |
| /** |
| * The byte offset into the Avro container file where the current input split |
| * ends. |
| */ |
| private long mEndPosition; |
| |
| /** |
| * Constructor. |
| * |
| * @param readerSchema The reader schema for the records of the Avro container |
| * file. |
| */ |
| protected AvroRecordReaderBase(Schema readerSchema) { |
| mReaderSchema = readerSchema; |
| mCurrentRecord = null; |
| } |
| |
| /** {@inheritDoc} */ |
| @Override |
| public void initialize(InputSplit inputSplit, TaskAttemptContext context) throws IOException, InterruptedException { |
| if (!(inputSplit instanceof FileSplit)) { |
| throw new IllegalArgumentException("Only compatible with FileSplits."); |
| } |
| FileSplit fileSplit = (FileSplit) inputSplit; |
| |
| // Open a seekable input stream to the Avro container file. |
| SeekableInput seekableFileInput = createSeekableInput(context.getConfiguration(), fileSplit.getPath()); |
| |
| // Wrap the seekable input stream in an Avro DataFileReader. |
| Configuration conf = context.getConfiguration(); |
| GenericData dataModel = AvroSerialization.createDataModel(conf); |
| DatumReader<T> datumReader = dataModel.createDatumReader(mReaderSchema); |
| mAvroFileReader = createAvroFileReader(seekableFileInput, datumReader); |
| |
| // Initialize the start and end offsets into the file based on the boundaries of |
| // the |
| // input split we're responsible for. We will read the first block that begins |
| // after the input split start boundary. We will read up to but not including |
| // the |
| // first block that starts after input split end boundary. |
| |
| // Sync to the closest block/record boundary just after beginning of our input |
| // split. |
| mAvroFileReader.sync(fileSplit.getStart()); |
| |
| // Initialize the start position to the beginning of the first block of the |
| // input split. |
| mStartPosition = mAvroFileReader.previousSync(); |
| |
| // Initialize the end position to the end of the input split (this isn't |
| // necessarily |
| // on a block boundary so using this for reporting progress will be approximate. |
| mEndPosition = fileSplit.getStart() + fileSplit.getLength(); |
| } |
| |
| /** {@inheritDoc} */ |
| @Override |
| public boolean nextKeyValue() throws IOException, InterruptedException { |
| assert null != mAvroFileReader; |
| |
| if (mAvroFileReader.hasNext() && !mAvroFileReader.pastSync(mEndPosition)) { |
| mCurrentRecord = mAvroFileReader.next(mCurrentRecord); |
| return true; |
| } |
| return false; |
| } |
| |
| /** {@inheritDoc} */ |
| @Override |
| public float getProgress() throws IOException, InterruptedException { |
| assert null != mAvroFileReader; |
| |
| if (mEndPosition == mStartPosition) { |
| // Trivial empty input split. |
| return 0.0f; |
| } |
| long bytesRead = mAvroFileReader.previousSync() - mStartPosition; |
| long bytesTotal = mEndPosition - mStartPosition; |
| LOG.debug("Progress: bytesRead=" + bytesRead + ", bytesTotal=" + bytesTotal); |
| return Math.min(1.0f, (float) bytesRead / (float) bytesTotal); |
| } |
| |
| /** {@inheritDoc} */ |
| @Override |
| public void close() throws IOException { |
| if (null != mAvroFileReader) { |
| try { |
| mAvroFileReader.close(); |
| } finally { |
| mAvroFileReader = null; |
| } |
| } |
| } |
| |
| /** |
| * Gets the current record read from the Avro container file. |
| * |
| * <p> |
| * Calling <code>nextKeyValue()</code> moves this to the next record. |
| * </p> |
| * |
| * @return The current Avro record (may be null if no record has been read). |
| */ |
| protected T getCurrentRecord() { |
| return mCurrentRecord; |
| } |
| |
| /** |
| * Creates a seekable input stream to an Avro container file. |
| * |
| * @param conf The hadoop configuration. |
| * @param path The path to the avro container file. |
| * @throws IOException If there is an error reading from the path. |
| */ |
| protected SeekableInput createSeekableInput(Configuration conf, Path path) throws IOException { |
| return new FsInput(path, conf); |
| } |
| |
| /** |
| * Creates an Avro container file reader from a seekable input stream. |
| * |
| * @param input The input containing the Avro container file. |
| * @param datumReader The reader to use for the individual records in the Avro |
| * container file. |
| * @throws IOException If there is an error reading from the input stream. |
| */ |
| protected DataFileReader<T> createAvroFileReader(SeekableInput input, DatumReader<T> datumReader) throws IOException { |
| return new DataFileReader<>(input, datumReader); |
| } |
| } |