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.convert.avro
10 
11 import com.typesafe.config.Config
12 import org.apache.avro.Schema
13 import org.apache.avro.Schema.Parser
14 import org.apache.avro.file.DataFileStream
15 import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter, GenericRecord}
16 import org.apache.avro.io.{DecoderFactory, EncoderFactory}
17 import org.geotools.api.feature.simple.SimpleFeatureType
18 import org.locationtech.geomesa.convert.EvaluationContext
19 import org.locationtech.geomesa.convert.avro.AvroConverter._
20 import org.locationtech.geomesa.convert2.AbstractConverter.{BasicField, BasicOptions}
21 import org.locationtech.geomesa.convert2.transforms.Expression
22 import org.locationtech.geomesa.convert2.transforms.Expression.Column
23 import org.locationtech.geomesa.convert2.{AbstractConverter, ConverterConfig}
24 import org.locationtech.geomesa.utils.collection.CloseableIterator
25 import org.locationtech.geomesa.utils.io.CopyingInputStream
26 
27 import java.io.{ByteArrayOutputStream, InputStream}
28 
29 class AvroConverter(sft: SimpleFeatureType, config: AvroConfig, fields: Seq[BasicField], options: BasicOptions)
30     extends AbstractConverter[GenericRecord, AvroConfig, BasicField, BasicOptions](sft, config, fields, options) {
31 
32   private val schema = config.schema match {
33     case SchemaEmbedded => None
34     case SchemaString(s) => Some(new Parser().parse(s))
35     case SchemaFile(s) =>
36       val loader = Option(Thread.currentThread.getContextClassLoader).getOrElse(getClass.getClassLoader)
37       val res = Option(loader.getResourceAsStream(s)).orElse(Option(getClass.getResourceAsStream(s))).getOrElse {
38         throw new IllegalArgumentException(s"Could not load schema resource at $s")
39       }
40       Some(new Parser().parse(res))
41   }
42 
43   // if required, set the raw bytes in the result array
44   private val requiresBytes = {
45     val expressions = config.idField.toSeq ++ fields.flatMap(_.transforms) ++ config.userData.values
46     Expression.flatten(expressions).contains(Column(0))
47   }
48 
49   override protected def parse(is: InputStream, ec: EvaluationContext): CloseableIterator[GenericRecord] = {
50     schema match {
51       case Some(s) if requiresBytes => new GenericRecordBytesIterator(new CopyingInputStream(is), s, ec)
52       case Some(s)                  => new GenericRecordIterator(is, s, ec)
53       case None    if requiresBytes => new FileStreamBytesIterator(is, ec)
54       case None                     => new FileStreamIterator(is, ec)
55     }
56   }
57 
58   override protected def values(parsed: CloseableIterator[GenericRecord],
59                                 ec: EvaluationContext): CloseableIterator[Array[Any]] = {
60     val array = Array.ofDim[Any](2)
61     if (requiresBytes) {
62       parsed.map { record => array(0) = record.get(BytesField); array(1) = record; array }
63     } else {
64       parsed.map { record => array(1) = record; array }
65     }
66   }
67 }
68 
69 object AvroConverter {
70 
71   import scala.collection.JavaConverters._
72 
73   val BytesField = "__bytes__"
74 
75   /**
76     * Add a `__bytes__` field to the schema, for storing the raw bytes
77     *
78     * @param schema schema
79     * @return
80     */
81   def addBytes(schema: Schema): Schema = {
82     schema.getType match {
83       case Schema.Type.RECORD =>
84         val fields = new java.util.ArrayList[Schema.Field](schema.getFields.size() + 1)
85         schema.getFields.asScala.foreach { field =>
86           fields.add(new Schema.Field(field.name, field.schema, field.doc, field.defaultVal()))
87         }
88         fields.add(new Schema.Field(BytesField, Schema.create(Schema.Type.BYTES), "raw bytes", ""))
89 
90         val updated = Schema.createRecord(schema.getName, schema.getDoc, schema.getNamespace, schema.isError)
91         updated.setFields(fields)
92         updated
93 
94       case Schema.Type.UNION =>
95         Schema.createUnion(schema.getTypes.asScala.map(s => addBytes(s)).toSeq: _*)
96 
97       case _ =>
98         throw new NotImplementedError(
99           s"Raw Avro bytes (i.e. $$0) is not implemented for schema type ${schema.getType}")
100     }
101   }
102 
103   case class AvroConfig(
104       `type`: String,
105       schema: SchemaConfig,
106       idField: Option[Expression],
107       caches: Map[String, Config],
108       userData: Map[String, Expression]
109     ) extends ConverterConfig
110 
111   sealed trait SchemaConfig
112 
113   case class SchemaString(schema: String) extends SchemaConfig
114   case class SchemaFile(file: String) extends SchemaConfig
115   case object SchemaEmbedded extends SchemaConfig {
116     val name: String = "embedded"
117   }
118 
119   /**
120     * Reads avro records using a pre-defined schema
121     *
122     * @param is input stream
123     * @param schema schema
124     * @param ec evaluation context
125     */
126   class GenericRecordIterator private [AvroConverter] (is: InputStream, schema: Schema, ec: EvaluationContext)
127       extends CloseableIterator[GenericRecord] {
128 
129     private val reader = new GenericDatumReader[GenericRecord](schema)
130     private val decoder = DecoderFactory.get.binaryDecoder(is, null)
131     private var record: GenericRecord = _
132 
133     override def hasNext: Boolean = !decoder.isEnd
134 
135     override def next(): GenericRecord = {
136       ec.line += 1
137       record = reader.read(record, decoder)
138       record
139     }
140 
141     override def close(): Unit = is.close()
142   }
143 
144   /**
145     * Reads avro records using a pre-defined schema, setting the bytes for each record in a
146     * special `__bytes__` field
147     *
148     * @param is input stream
149     * @param schema schema
150     * @param ec evaluation context
151     */
152   class GenericRecordBytesIterator private [AvroConverter] (is: CopyingInputStream, schema: Schema, ec: EvaluationContext)
153       extends CloseableIterator[GenericRecord] {
154 
155     private val reader = new GenericDatumReader[GenericRecord](schema, addBytes(schema))
156     private val decoder = DecoderFactory.get.binaryDecoder(is, null)
157     private var record: GenericRecord = _
158 
159     override def hasNext: Boolean = !decoder.isEnd
160 
161     override def next(): GenericRecord = {
162       ec.line += 1
163       record = reader.read(record, decoder)
164       // parse out the bytes read and set them in the record
165       // check to see if the decoder buffered some bytes that weren't actually used
166       val buffered = decoder.inputStream().available()
167       record.put(BytesField, is.replay(is.copied - buffered))
168       record
169     }
170 
171     override def close(): Unit = is.close()
172   }
173 
174   /**
175     * Reads avro records from an avro file, with the schema embedded
176     *
177     * @param is input
178     * @param ec evaluation context
179     */
180   class FileStreamIterator private [AvroConverter] (is: InputStream, ec: EvaluationContext)
181       extends CloseableIterator[GenericRecord] {
182 
183     private val stream = new DataFileStream(is, new GenericDatumReader[GenericRecord]())
184     private var record: GenericRecord = _
185 
186     override def hasNext: Boolean = stream.hasNext
187 
188     override def next(): GenericRecord = {
189       ec.line += 1
190       record = stream.next(record)
191       record
192     }
193 
194     override def close(): Unit = stream.close()
195   }
196 
197   /**
198     * Reads avro records from an avro file, with the schema embedded, setting the bytes for
199     * each record in a special `__bytes__` field
200     *
201     * @param is input
202     * @param ec evaluation context
203     */
204   class FileStreamBytesIterator private [AvroConverter] (is: InputStream, ec: EvaluationContext)
205       extends CloseableIterator[GenericRecord] {
206 
207     private val reader = new GenericDatumReader[GenericRecord]()
208     private val stream = new DataFileStream(is, reader)
209     private var record: GenericRecord = _
210 
211     reader.setExpected(addBytes(reader.getSchema))
212 
213     // we can't tell which bytes correspond to which feature (due to buffering). if we could access the
214     // underlying avro encoder we could figure it out, but it is not exposed through DataFileStream. instead,
215     // re-serialize each record to get the raw bytes
216     private val out = new ByteArrayOutputStream()
217     private val writer = new GenericDatumWriter[GenericRecord](stream.getSchema)
218     private val encoder = EncoderFactory.get.binaryEncoder(out, null)
219 
220     override def hasNext: Boolean = stream.hasNext
221 
222     override def next(): GenericRecord = {
223       ec.line += 1
224       record = stream.next(record)
225       // regenerate the bytes read and set them in the record
226       out.reset()
227       writer.write(record, encoder)
228       encoder.flush()
229       record.put(BytesField, out.toByteArray)
230       record
231     }
232 
233     override def close(): Unit = stream.close()
234   }
235 }