1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, 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  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.utils.bin
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
13 import org.locationtech.geomesa.utils.bin.BinaryEncodeCallback.{ByteArrayCallback, ByteStreamCallback}
14 import org.locationtech.geomesa.utils.bin.BinaryOutputEncoder.ToValues
15 import org.locationtech.geomesa.utils.collection.CloseableIterator
16 import org.locationtech.geomesa.utils.geotools.SimpleFeatureTypes
17 import org.locationtech.geomesa.utils.geotools.sft.SimpleFeatureSpec.ListAttributeSpec
18 import org.locationtech.geomesa.utils.geotools.sft.SimpleFeatureSpecParser
19 import org.locationtech.jts.geom.{Geometry, LineString, Point}
20 
21 import java.io.{ByteArrayOutputStream, OutputStream}
22 import java.nio.charset.StandardCharsets
23 import java.nio.{ByteBuffer, ByteOrder}
24 import java.util.Date
25 import scala.collection.JavaConverters._
26 
27 class BinaryOutputEncoder private (toValues: ToValues) {
28 
29   def encode(f: SimpleFeature): Array[Byte] = {
30     toValues(f, ByteArrayCallback)
31     ByteArrayCallback.result
32   }
33 
34   def encode(f: SimpleFeature, callback: BinaryOutputCallback): Unit = toValues(f, callback)
35 
36   def encode(f: Iterator[SimpleFeature], os: OutputStream): Long =
37     encode(f, os, sort = false)
38 
39   def encode(f: Iterator[SimpleFeature], os: OutputStream, sort: Boolean): Long = {
40     if (sort) {
41       val byteStream = new ByteArrayOutputStream
42       val callback = new ByteStreamCallback(byteStream)
43       f.foreach(toValues(_, callback))
44       val count = callback.result
45       val bytes = byteStream.toByteArray
46       val size = (bytes.length / count).toInt
47       bytes.grouped(size).toSeq.sorted(BinaryOutputEncoder.DateOrdering).foreach(os.write)
48       count
49     } else {
50       val callback = new ByteStreamCallback(os)
51       f.foreach(toValues(_, callback))
52       callback.result
53     }
54   }
55 }
56 
57 object BinaryOutputEncoder extends LazyLogging {
58 
59   import AxisOrder._
60   import org.locationtech.geomesa.utils.geotools.Conversions._
61   import org.locationtech.geomesa.utils.geotools.RichSimpleFeatureType.RichSimpleFeatureType
62 
63   val BinEncodedSft: SimpleFeatureType = SimpleFeatureTypes.createType("bin", "bin:Bytes,*geom:Point:srid=4326")
64   val BIN_ATTRIBUTE_INDEX = 0 // index of 'bin' attribute in BinEncodedSft
65 
66   // compares the 4 bytes representing the date in a bin array
67   private val DateOrdering = new Ordering[Array[Byte]] {
68     override def compare(x: Array[Byte], y: Array[Byte]): Int = {
69       val compare1 = Ordering.Byte.compare(x(4), y(4))
70       if (compare1 != 0) { return compare1 }
71       val compare2 = Ordering.Byte.compare(x(5), y(5))
72       if (compare2 != 0) { return compare2 }
73       val compare3 = Ordering.Byte.compare(x(6), y(6))
74       if (compare3 != 0) { return compare3 }
75       Ordering.Byte.compare(x(7), y(7))
76     }
77   }
78 
79   case class EncodingOptions(geomField: Option[Int],
80                              dtgField: Option[Int],
81                              trackIdField: Option[Int],
82                              labelField: Option[Int] = None,
83                              axisOrder: Option[AxisOrder] = None)
84 
85   case class EncodedValues(trackId: Int, lat: Float, lon: Float, dtg: Long, label: Long)
86 
87   /**
88     * BIN queries pack multiple records into each feature. To count the records, we have to count
89     * the total bytes coming back, instead of the number of features
90     *
91     * @param iter aggregated bin iter
92     * @param maxFeatures max features
93     * @param hasLabel bin results have labels (extended format) or not
94     */
95   class FeatureLimitingIterator(iter: CloseableIterator[SimpleFeature], maxFeatures: Int, hasLabel: Boolean)
96       extends CloseableIterator[SimpleFeature] {
97 
98     private val bytesPerHit = if (hasLabel) { 24 } else { 16 }
99     private var seen = 0L
100 
101     override def hasNext: Boolean = seen < maxFeatures && iter.hasNext
102 
103     override def next(): SimpleFeature = {
104       if (hasNext) {
105         val sf = iter.next()
106         val bytes = sf.getAttribute(0).asInstanceOf[Array[Byte]]
107         val count = bytes.length / bytesPerHit
108         seen += count
109         if (seen > maxFeatures) {
110           // remove the extra aggregated features so that we hit our exact feature limit
111           val trimmed = Array.ofDim[Byte]((count - (seen - maxFeatures).toInt) * bytesPerHit)
112           System.arraycopy(bytes, 0, trimmed, 0, trimmed.length)
113           sf.setAttribute(0, trimmed)
114         }
115         sf
116       } else {
117         Iterator.empty.next()
118       }
119     }
120 
121     override def close(): Unit = iter.close()
122   }
123 
124   def apply(sft: SimpleFeatureType, options: EncodingOptions): BinaryOutputEncoder =
125     new BinaryOutputEncoder(toValues(sft, options))
126 
127   def convertToTrack(f: SimpleFeature, i: Int): Int = convertToTrack(f.getAttribute(i))
128   def convertToTrack(track: AnyRef): Int = if (track == null) { 0 } else { track.hashCode }
129 
130   // TODO could use `.getDateAsLong` if we know we have a KryoBufferSimpleFeature
131   def convertToDate(f: SimpleFeature, i: Int): Long = convertToDate(f.getAttribute(i).asInstanceOf[Date])
132   def convertToDate(date: Date): Long = if (date == null) { 0L } else { date.getTime }
133 
134   def convertToLabel(f: SimpleFeature, i: Int): Long = convertToLabel(f.getAttribute(i))
135   def convertToLabel(label: AnyRef): Long = label match {
136     case null => 0L
137     case n: Number => n.longValue()
138     case _ =>
139       var sum = 0L
140       var i = 0
141       label.toString.getBytes(StandardCharsets.UTF_8).iterator.take(8).foreach { b =>
142         sum += (b & 0xffL) << (8 * i)
143         i += 1
144       }
145       sum
146   }
147 
148   /**
149     * Decodes a byte array
150     *
151     * @param encoded encoded byte array
152     * @param callback callback for results
153     */
154   def decode(encoded: Array[Byte], callback: BinaryOutputCallback): Unit = {
155     val buf = ByteBuffer.wrap(encoded).order(ByteOrder.LITTLE_ENDIAN)
156     val trackId = buf.getInt
157     val time = buf.getInt * 1000L
158     val lat = buf.getFloat
159     val lon = buf.getFloat
160     if (encoded.length > 16) {
161       val label = buf.getLong
162       callback(trackId, lat, lon, time, label)
163     } else {
164       callback(trackId, lat, lon, time)
165     }
166   }
167 
168   def decode(encoded: Array[Byte]): EncodedValues = {
169     var values: EncodedValues = null
170     decode(encoded, new BinaryOutputCallback() {
171       override def apply(trackId: Int, lat: Float, lon: Float, dtg: Long): Unit =
172         values = EncodedValues(trackId, lat, lon, dtg, -1L)
173       override def apply(trackId: Int, lat: Float, lon: Float, dtg: Long, label: Long): Unit =
174         values = EncodedValues(trackId, lat, lon, dtg, label)
175     })
176     values
177   }
178 
179   /**
180     * Creates the function to map a simple feature to a bin-encoded buffer
181     *
182     * @param sft simple feature type
183     * @param options encoding options
184     * @return
185     */
186   private def toValues(sft: SimpleFeatureType, options: EncodingOptions): ToValues = {
187 
188     val geomIndex = options.geomField.getOrElse(sft.getGeomIndex)
189     if (geomIndex == -1) {
190       throw new IllegalArgumentException(s"Invalid geometry field requested for feature type ${sft.getTypeName}")
191     }
192     val dtgIndex = options.dtgField.orElse(sft.getDtgIndex).getOrElse(-1)
193     if (dtgIndex == -1) {
194       throw new RuntimeException(s"Invalid date field requested for feature type ${sft.getTypeName}")
195     }
196     val isSingleDate = classOf[Date].isAssignableFrom(sft.getDescriptor(dtgIndex).getType.getBinding)
197     val axisOrder = options.axisOrder.getOrElse(AxisOrder.LonLat)
198 
199     val (isPoint, isLineString) = {
200       val binding = sft.getDescriptor(geomIndex).getType.getBinding
201       (binding == classOf[Point], binding == classOf[LineString])
202     }
203 
204     // noinspection ExistsEquals
205     if (options.trackIdField.exists(_ == -1)) {
206       throw new IllegalArgumentException(s"Invalid track field requested for feature type ${sft.getTypeName}")
207     } else if (options.labelField.exists(_ == -1)) {
208       throw new IllegalArgumentException(s"Invalid label field requested for feature type ${sft.getTypeName}")
209     } else if (!isSingleDate) {
210       if (isLineString) {
211         val dtgField = sft.getDescriptor(dtgIndex).getLocalName
212         val sftAttributes = SimpleFeatureSpecParser.parse(SimpleFeatureTypes.encodeType(sft)).attributes
213         sftAttributes.find(_.name == dtgField).foreach { spec =>
214           if (!spec.isInstanceOf[ListAttributeSpec] ||
215                 !classOf[Date].isAssignableFrom(spec.asInstanceOf[ListAttributeSpec].subClass)) {
216             throw new RuntimeException(s"Invalid date field requested for feature type ${sft.getTypeName}")
217           }
218         }
219       } else {
220         throw new RuntimeException(s"Invalid date field requested for feature type ${sft.getTypeName}")
221       }
222     }
223 
224     // gets the track id from a feature
225     val getTrackId: (SimpleFeature) => Int = options.trackIdField match {
226       case None => (f) => f.getID.hashCode
227       case Some(trackId) => convertToTrack(_, trackId)
228     }
229 
230     // gets the label from a feature
231     val getLabelOption: Option[(SimpleFeature) => Long] = options.labelField.map { labelIndex =>
232       convertToLabel(_, labelIndex)
233     }
234 
235     if (isLineString) {
236       // for linestrings, we return each point - use an array so we get constant-time lookup
237       // depending on srs requested and wfs versions, axis order can be flipped
238       val getLineLatLon: (SimpleFeature) => Array[(Float, Float)] = axisOrder match {
239         case LatLon => lineToXY(_, geomIndex)
240         case LonLat => lineToYX(_, geomIndex)
241       }
242 
243       if (isSingleDate) {
244         getLabelOption match {
245           case None => new ToValuesLines(getTrackId, getLineLatLon, dtgIndex)
246           case Some(getLabel) => new ToValuesLinesLabels(getTrackId, getLineLatLon, getLabel, dtgIndex)
247         }
248       } else {
249         // for line strings, we need an array of dates corresponding to the points in the line
250         val getLineDtg: (SimpleFeature) => Array[Long] = dateArray(_, dtgIndex)
251         getLabelOption match {
252           case None => new ToValuesLinesDates(getTrackId, getLineLatLon, getLineDtg)
253           case Some(getLabel) => new ToValuesLinesDatesLabels(getTrackId, getLineLatLon, getLineDtg, getLabel)
254         }
255       }
256     } else {
257       // get lat/lon as floats
258       // depending on srs requested and wfs versions, axis order can be flipped
259       val getLatLon: (SimpleFeature) => (Float, Float) = (isPoint, axisOrder) match {
260         case (true,  LatLon) => pointToXY(_, geomIndex)
261         case (true,  LonLat) => pointToYX(_, geomIndex)
262         case (false, LatLon) => geomToXY(_, geomIndex)
263         case (false, LonLat) => geomToYX(_, geomIndex)
264       }
265 
266       getLabelOption match {
267         case None => new ToValuesPoints(getTrackId, getLatLon, dtgIndex)
268         case Some(getLabel) => new ToValuesPointsLabels(getTrackId, getLatLon, getLabel, dtgIndex)
269       }
270     }
271   }
272 
273   private def pointToXY(p: Point): (Float, Float) = (p.getX.toFloat, p.getY.toFloat)
274   private def pointToYX(p: Point): (Float, Float) = (p.getY.toFloat, p.getX.toFloat)
275 
276   private def pointToXY(f: SimpleFeature, i: Int): (Float, Float) =
277     pointToXY(f.getAttribute(i).asInstanceOf[Point])
278   private def pointToYX(f: SimpleFeature, i: Int): (Float, Float) =
279     pointToYX(f.getAttribute(i).asInstanceOf[Point])
280 
281   private def geomToXY(f: SimpleFeature, i: Int): (Float, Float) =
282     pointToXY(f.getAttribute(i).asInstanceOf[Geometry].safeCentroid())
283   private def geomToYX(f: SimpleFeature, i: Int): (Float, Float) =
284     pointToYX(f.getAttribute(i).asInstanceOf[Geometry].safeCentroid())
285 
286   private def lineToXY(f: SimpleFeature, i: Int): Array[(Float, Float)] = {
287     val line = f.getAttribute(i).asInstanceOf[LineString]
288     Array.tabulate(line.getNumPoints)(i => pointToXY(line.getPointN(i)))
289   }
290   private def lineToYX(f: SimpleFeature, i: Int): Array[(Float, Float)] = {
291     val line = f.getAttribute(i).asInstanceOf[LineString]
292     Array.tabulate(line.getNumPoints)(i => pointToYX(line.getPointN(i)))
293   }
294 
295   private def dateArray(f: SimpleFeature, i: Int): Array[Long] = {
296     val dates = f.getAttribute(i).asInstanceOf[java.util.List[Date]]
297     if (dates == null) { Array.empty } else { dates.asScala.map(_.getTime).toArray }
298   }
299 
300   private trait ToValues {
301     def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit
302   }
303 
304   private class ToValuesPoints(getTrackId: (SimpleFeature) => Int,
305                                getLatLon: (SimpleFeature) => (Float, Float),
306                                dtgIndex: Int) extends ToValues {
307     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
308       val (lat, lon) = getLatLon(f)
309       callback(getTrackId(f), lat, lon, convertToDate(f, dtgIndex))
310     }
311   }
312 
313   private class ToValuesPointsLabels(getTrackId: (SimpleFeature) => Int,
314                                      getLatLon: (SimpleFeature) => (Float, Float),
315                                      getLabel: (SimpleFeature) => Long,
316                                      dtgIndex: Int) extends ToValues {
317     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
318       val (lat, lon) = getLatLon(f)
319       callback(getTrackId(f), lat, lon, convertToDate(f, dtgIndex), getLabel(f))
320     }
321   }
322 
323   private class ToValuesLines(getTrackId: (SimpleFeature) => Int,
324                               getLatLon: (SimpleFeature) => Array[(Float, Float)],
325                               dtgIndex: Int) extends ToValues {
326     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
327       val trackId = getTrackId(f)
328       val points = getLatLon(f)
329       val date = convertToDate(f, dtgIndex)
330       var i = 0
331       while (i < points.length) {
332         val (lat, lon) = points(i)
333         callback(trackId, lat, lon, date)
334         i += 1
335       }
336     }
337   }
338 
339   private class ToValuesLinesLabels(getTrackId: (SimpleFeature) => Int,
340                                     getLatLon: (SimpleFeature) => Array[(Float, Float)],
341                                     getLabel: (SimpleFeature) => Long,
342                                     dtgIndex: Int) extends ToValues {
343     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
344       val trackId = getTrackId(f)
345       val points = getLatLon(f)
346       val date = convertToDate(f, dtgIndex)
347       val label = getLabel(f)
348       var i = 0
349       while (i < points.length) {
350         val (lat, lon) = points(i)
351         callback(trackId, lat, lon, date, label)
352         i += 1
353       }
354     }
355   }
356 
357   private class ToValuesLinesDates(getTrackId: (SimpleFeature) => Int,
358                                    getLatLon: (SimpleFeature) => Array[(Float, Float)],
359                                    getLineDtg: (SimpleFeature) => Array[Long]) extends ToValues {
360     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
361       val trackId = getTrackId(f)
362       val points = getLatLon(f)
363       val dates = getLineDtg(f)
364       val size = if (points.length == dates.length) { points.length } else {
365         logger.warn(s"Mismatched geometries and dates for simple feature $f: ${points.toList} ${dates.toList}")
366         math.min(points.length, dates.length)
367       }
368       var i = 0
369       while (i < size) {
370         val (lat, lon) = points(i)
371         callback(trackId, lat, lon, dates(i))
372         i += 1
373       }
374     }
375   }
376 
377   private class ToValuesLinesDatesLabels(getTrackId: (SimpleFeature) => Int,
378                                          getLatLon: (SimpleFeature) => Array[(Float, Float)],
379                                          getLineDtg: (SimpleFeature) => Array[Long],
380                                          getLabel: (SimpleFeature) => Long) extends ToValues {
381     override def apply(f: SimpleFeature, callback: BinaryOutputCallback): Unit = {
382       val trackId = getTrackId(f)
383       val points = getLatLon(f)
384       val dates = getLineDtg(f)
385       val size = if (points.length == dates.length) { points.length } else {
386         logger.warn(s"Mismatched geometries and dates for simple feature $f: ${points.toList} ${dates.toList}")
387         math.min(points.length, dates.length)
388       }
389       val label = getLabel(f)
390       var i = 0
391       while (i < size) {
392         val (lat, lon) = points(i)
393         callback(trackId, lat, lon, dates(i), label)
394         i += 1
395       }
396     }
397   }
398 }