1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.jobs.mapreduce
10 
11 import org.apache.hadoop.fs.{Path, Seekable}
12 import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
13 import org.geotools.api.feature.simple.SimpleFeature
14 import org.locationtech.geomesa.features.ScalaSimpleFeature
15 import org.locationtech.geomesa.features.avro.io.AvroDataFileReader
16 import org.locationtech.geomesa.utils.geotools.SimpleFeatureTypes
17 
18 import java.io.{Closeable, InputStream}
19 
20 /**
21  * Class for reading avro files written using
22  * <code>org.locationtech.geomesa.features.avro.AvroDataFileWriter</code>.
23  */
24 class AvroFileInputFormat extends FileStreamInputFormat {
25   override def createRecordReader(): FileStreamRecordReader = new AvroFileRecordReader
26 }
27 
28 object AvroFileInputFormat {
29   object Counters {
30     val Group = "org.locationtech.geomesa.jobs.input.avro"
31     val Read  = "read"
32   }
33 
34   def setTypeName(job: Job, typeName: String): Unit =
35     job.getConfiguration.set(FileStreamInputFormat.TypeNameKey, typeName)
36 }
37 
38 class AvroFileRecordReader extends FileStreamRecordReader {
39 
40   import AvroFileInputFormat.Counters
41 
42   override def createIterator(stream: InputStream with Seekable,
43                               filePath: Path,
44                               context: TaskAttemptContext): Iterator[SimpleFeature] with Closeable = {
45     val typeName = context.getConfiguration.get(FileStreamInputFormat.TypeNameKey)
46     val reader = new AvroDataFileReader(stream)
47     val dataSft = reader.getSft
48     val counter = context.getCounter(Counters.Group, Counters.Read)
49 
50     if (typeName == null || dataSft.getTypeName == typeName) {
51       new Iterator[SimpleFeature] with Closeable {
52         override def hasNext: Boolean = reader.hasNext
53         override def next(): SimpleFeature = {
54           counter.increment(1)
55           reader.next()
56         }
57         override def close(): Unit = reader.close()
58       }
59     } else {
60       val sft = SimpleFeatureTypes.renameSft(dataSft, typeName)
61       new Iterator[SimpleFeature] with Closeable {
62         override def hasNext: Boolean = reader.hasNext
63         override def next(): SimpleFeature = {
64           val sf = reader.next()
65           counter.increment(1)
66           new ScalaSimpleFeature(sft, sf.getID, sf.getAttributes.toArray)
67         }
68         override def close(): Unit = reader.close()
69       }
70     }
71   }
72 }