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.features.kryo
10 
11 import com.esotericsoftware.kryo.io.{Input, Output}
12 import org.geotools.api.feature.`type`.{AttributeDescriptor, Name}
13 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
14 import org.geotools.api.feature.{GeometryAttribute, Property}
15 import org.geotools.api.filter.identity.FeatureId
16 import org.geotools.api.geometry.BoundingBox
17 import org.geotools.geometry.jts.ReferencedEnvelope
18 import org.locationtech.geomesa.features.ScalaSimpleFeature
19 import org.locationtech.geomesa.features.SerializationOption.SerializationOption
20 import org.locationtech.geomesa.features.geotools.ImmutableFeatureId
21 import org.locationtech.geomesa.features.kryo.KryoBufferSimpleFeature._
22 import org.locationtech.geomesa.features.kryo.impl.KryoFeatureDeserialization.KryoLongReader
23 import org.locationtech.geomesa.features.kryo.impl.{IntBitSet, KryoFeatureDeserialization, KryoFeatureDeserializationV2, NonMutatingInput}
24 import org.locationtech.geomesa.features.kryo.serialization.KryoUserDataSerialization
25 import org.locationtech.geomesa.utils.geotools.Transform
26 import org.locationtech.geomesa.utils.geotools.Transform.{PropertyTransform, RenameTransform, Transforms}
27 import org.locationtech.jts.geom.Geometry
28 
29 class KryoBufferSimpleFeature(serializer: KryoFeatureDeserialization) extends SimpleFeature {
30 
31   import org.locationtech.geomesa.utils.geotools.RichSimpleFeatureType.RichSimpleFeatureType
32 
33   private val input = new NonMutatingInput()
34 
35   private val idParser = if (serializer.withoutId) { new IdParser() } else { new WithIdParser() }
36   private val delegateV3 = new KryoBufferV3(serializer, input)
37 
38   @volatile
39   private var seenV2 = false
40   private lazy val delegateV2 = {
41     seenV2 = true
42     val buf = new KryoBufferV2(serializer, input)
43     if (transforms != null) {
44       buf.setTransforms(transformSchema, transformDefinitions)
45     }
46     buf
47   }
48 
49   private lazy val geomIndex = serializer.out.getGeomIndex
50 
51   private var delegate: KryoBufferDelegate = _
52 
53   private var transforms: String = _
54   private var transformSchema: SimpleFeatureType = _
55   private var transformDefinitions: Array[Transform] = _
56 
57   private var userData: java.util.Map[AnyRef, AnyRef] = _
58 
59   /**
60     * Creates a new feature for later use - does not copy attribute bytes
61     *
62     * @return
63     */
64   def copy(): KryoBufferSimpleFeature = {
65     val sf = new KryoBufferSimpleFeature(serializer)
66     sf.setIdParser(idParser.parse)
67     if (transforms != null) {
68       sf.setTransforms(transforms, transformSchema)
69     }
70     sf
71   }
72 
73   /**
74     * Sets the parser for reading feature ids out of the id buffer
75     *
76     * @param parse parse method
77     */
78   def setIdParser(parse: (Array[Byte], Int, Int) => String): Unit = idParser.parse = parse
79 
80   /**
81     * Sets the transform to be applied to this feature when calling `transform`
82     *
83     * @param transforms transform definition, per geotools format
84     * @param transformSchema schema that results from applying the transform
85     */
86   def setTransforms(transforms: String, transformSchema: SimpleFeatureType): Unit = {
87     this.transforms = transforms
88     this.transformSchema = transformSchema
89     this.transformDefinitions = Transforms(serializer.out, transforms).toArray
90     delegateV3.setTransforms(transformSchema, transformDefinitions)
91     if (seenV2) {
92       delegateV2.setTransforms(transformSchema, transformDefinitions)
93     }
94   }
95 
96   /**
97     * Gets any transforms applied to this feature
98     *
99     * @return
100     */
101   def getTransform: Option[(String, SimpleFeatureType)] =
102     for { t <- Option(transforms); s <- Option(transformSchema) } yield { (t, s) }
103 
104   /**
105     * Set the serialized bytes to use for reading attributes
106     *
107     * @param bytes serialized byte array
108     */
109   def setBuffer(bytes: Array[Byte]): Unit = setBuffer(bytes, 0, bytes.length)
110 
111   /**
112     * Set the serialized bytes to use for reading attributes
113     *
114     * @param bytes serialized byte array
115     * @param offset offset into the byte array of valid bytes
116     * @param length number of valid bytes to read from the byte array
117     */
118   def setBuffer(bytes: Array[Byte], offset: Int, length: Int): Unit = {
119     input.setBuffer(bytes, offset, length)
120     delegate = input.readByte() match {
121       case KryoFeatureSerializer.Version3 => delegateV3
122       case KryoFeatureSerializer.Version2 => delegateV2
123       case b => throw new IllegalArgumentException(s"Can't process features serialized with version: $b")
124     }
125     delegate.reset()
126     userData = null
127   }
128 
129   /**
130     * Sets the serialized bytes containing the feature ID (i.e. the row key)
131     *
132     * @param bytes bytes
133     */
134   def setIdBuffer(bytes: Array[Byte]): Unit = setIdBuffer(bytes, 0, bytes.length)
135 
136   /**
137     * Sets the serialized bytes containing the feature ID (i.e. the row key)
138     *
139     * @param bytes bytes
140     * @param offset offset into the byte array of valid bytes
141     * @param length number of valid bytes to read from the byte array
142     */
143   def setIdBuffer(bytes: Array[Byte], offset: Int, length: Int): Unit = {
144     idParser.buffer = bytes
145     idParser.offset = offset
146     idParser.length = length
147   }
148 
149   /**
150     * Transform the feature into a serialized byte array
151     *
152     * @return
153     */
154   def transform(): Array[Byte] = delegate.transform(this)
155 
156   /**
157     * Get a date attribute as a raw long
158     *
159     * @param index attribute index
160     * @return
161     */
162   def getDateAsLong(index: Int): Long = delegate.getDateAsLong(index)
163 
164   /**
165     * Get the underlying kryo input, positioned to read the attribute at the given index
166     *
167     * @param index attribute index
168     * @return input, if the attribute is not null
169     */
170   def getInput(index: Int): Option[Input] = delegate.getInput(index)
171 
172   override def getAttribute(index: Int): AnyRef = delegate.getAttribute(index)
173 
174   override def getType: SimpleFeatureType = serializer.out
175   override def getFeatureType: SimpleFeatureType = serializer.out
176   override def getName: Name = serializer.out.getName
177 
178   override def getID: String = idParser.id()
179   override def getIdentifier: FeatureId = new ImmutableFeatureId(idParser.id())
180 
181   override def getAttribute(name: Name): AnyRef = getAttribute(name.getLocalPart)
182   override def getAttribute(name: String): Object = {
183     val index = serializer.out.indexOf(name)
184     if (index == -1) { null } else { getAttribute(index) }
185   }
186 
187   override def getDefaultGeometry: AnyRef = if (geomIndex == -1) { null } else { getAttribute(geomIndex) }
188   override def getAttributeCount: Int = serializer.out.getAttributeCount
189 
190   override def getBounds: BoundingBox = getDefaultGeometry match {
191     case g: Geometry => new ReferencedEnvelope(g.getEnvelopeInternal, serializer.out.getCoordinateReferenceSystem)
192     case _           => new ReferencedEnvelope(serializer.out.getCoordinateReferenceSystem)
193   }
194 
195   override def getAttributes: java.util.List[AnyRef] = {
196     val attributes = new java.util.ArrayList[AnyRef](serializer.out.getAttributeCount)
197     var i = 0
198     while (i < serializer.out.getAttributeCount) {
199       attributes.add(getAttribute(i))
200       i += 1
201     }
202     attributes
203   }
204 
205   override def getUserData: java.util.Map[AnyRef, AnyRef] = {
206     if (userData == null) {
207       userData = if (serializer.withoutUserData) { new java.util.HashMap(1) } else { delegate.getUserData }
208     }
209     userData
210   }
211 
212   override def getDefaultGeometryProperty: GeometryAttribute = throw new UnsupportedOperationException()
213   override def getProperties: java.util.Collection[Property] = throw new UnsupportedOperationException()
214   override def getProperties(name: Name): java.util.Collection[Property] = throw new UnsupportedOperationException()
215   override def getProperties(name: String): java.util.Collection[Property] = throw new UnsupportedOperationException()
216   override def getProperty(name: Name): Property = throw new UnsupportedOperationException()
217   override def getProperty(name: String): Property = throw new UnsupportedOperationException()
218   override def getValue: java.util.Collection[_ <: Property] = throw new UnsupportedOperationException()
219   override def getDescriptor: AttributeDescriptor = throw new UnsupportedOperationException()
220 
221   override def setAttribute(name: Name, value: Object): Unit = throw new UnsupportedOperationException()
222   override def setAttribute(name: String, value: Object): Unit = throw new UnsupportedOperationException()
223   override def setAttribute(index: Int, value: Object): Unit = throw new UnsupportedOperationException()
224   override def setAttributes(vals: java.util.List[Object]): Unit = throw new UnsupportedOperationException()
225   override def setAttributes(vals: Array[Object]): Unit = throw new UnsupportedOperationException()
226   override def setDefaultGeometry(geo: Object): Unit = throw new UnsupportedOperationException()
227   override def setDefaultGeometryProperty(geoAttr: GeometryAttribute): Unit = throw new UnsupportedOperationException()
228   override def setValue(newValue: Object): Unit = throw new UnsupportedOperationException()
229   override def setValue(values: java.util.Collection[Property]): Unit = throw new UnsupportedOperationException()
230 
231   override def isNillable: Boolean = true
232   override def validate(): Unit = throw new UnsupportedOperationException()
233 
234   override def toString: String = s"KryoBufferSimpleFeature:$getID"
235 
236   private class WithIdParser extends IdParser {
237     override def id(): String = delegate.id()
238   }
239 }
240 
241 object KryoBufferSimpleFeature {
242 
243   private class IdParser {
244     var parse: (Array[Byte], Int, Int) => String = (_, _, _) => null
245     var buffer: Array[Byte] = _
246     var offset: Int = 0
247     var length: Int = 0
248 
249     def id(): String = parse(buffer, offset, length)
250   }
251 
252   /**
253     * Common interface for handling serialization versions
254     */
255   private sealed trait KryoBufferDelegate extends Transformer {
256 
257     /**
258       * Invoked after the underlying kryo buffer has been updated with a new serialized feature
259       */
260     def reset(): Unit
261 
262     /**
263       * Sets the transform to be applied to the feature
264       *
265       * @param schema schema that results from applying the transform
266       * @param transforms transform definitions
267       */
268     def setTransforms(schema: SimpleFeatureType, transforms: Array[Transform]): Unit
269 
270     /**
271       * Deserialize the feature ID (assuming options.withId)
272       *
273       * @return
274       */
275     def id(): String
276 
277     /**
278       * Deserialize an attribute
279       *
280       * @param index attribute number
281       * @return
282       */
283     def getAttribute(index: Int): AnyRef
284 
285     /**
286       * Deserialize user data
287       *
288       * @return
289       */
290     def getUserData: java.util.Map[AnyRef, AnyRef]
291 
292     /**
293       * Optimized method to get a date attribute as millis without creating a new Date object
294       *
295       * @param index attribute number
296       * @return
297       */
298     def getDateAsLong(index: Int): Long
299 
300     /**
301       * Gets the input, positioned to read the given attribute
302       *
303       * @param index attribute index
304       * @return
305       */
306     def getInput(index: Int): Option[Input]
307   }
308 
309   /**
310     * Abstraction over transformer impls
311     */
312   private sealed trait Transformer {
313 
314     /**
315       * Transform a feature, based on previously set transform schema
316       *
317       * @param original original feature being transformed
318       * @return
319       */
320     def transform(original: SimpleFeature): Array[Byte]
321   }
322 
323   /**
324     * Serialization version 3 delegate
325     *
326     * @param serializer serializer
327     * @param input input
328     */
329   private class KryoBufferV3(serializer: KryoFeatureDeserialization, input: Input) extends KryoBufferDelegate {
330 
331     private var metadata: Metadata = _
332     private var transformer: Transformer = _
333 
334     override def reset(): Unit = metadata = Metadata(input) // reads count, size, nulls, etc
335 
336     override def setTransforms(schema: SimpleFeatureType, transforms: Array[Transform]): Unit = {
337       val indices = transforms.map {
338         case t: PropertyTransform => t.i
339         case t: RenameTransform => t.i
340         case _ => -1
341       }
342       if (indices.contains(-1)) {
343         transformer = new ReserializeTransformer(schema, transforms, serializer.options)
344       } else {
345         transformer = new BinaryTransformer(indices)
346       }
347     }
348 
349     override def id(): String = {
350       metadata.setIdPosition()
351       input.readString()
352     }
353 
354     override def getAttribute(index: Int): AnyRef = {
355       if (index >= metadata.count || metadata.nulls.contains(index)) { null } else {
356         metadata.setPosition(index)
357         serializer.readers(index).apply(input)
358       }
359     }
360 
361     override def getUserData: java.util.Map[AnyRef, AnyRef] = {
362       metadata.setUserDataPosition()
363       KryoUserDataSerialization.deserialize(input)
364     }
365 
366     override def getDateAsLong(index: Int): Long = {
367       if (index >= metadata.count || metadata.nulls.contains(index)) { 0L } else {
368         metadata.setPosition(index)
369         KryoLongReader.apply(input)
370       }
371     }
372 
373     override def getInput(index: Int): Option[Input] = {
374       if (index >= metadata.count || metadata.nulls.contains(index)) { None } else {
375         metadata.setPosition(index)
376         Some(input)
377       }
378     }
379 
380     override def transform(original: SimpleFeature): Array[Byte] = transformer.transform(original)
381 
382     // if we are just returning a subset of attributes, we can copy the bytes directly
383     private class BinaryTransformer(indices: Array[Int]) extends Transformer {
384 
385       private val positionsAndLengths = Array.ofDim[(Int, Int)](indices.length)
386 
387       override def transform(original: SimpleFeature): Array[Byte] = {
388         // +1 for version, +2 for count, +1 for size
389         var length = (metadata.size * (indices.length + 1)) + (IntBitSet.size(indices.length) * 4) + 4
390         val id = if (serializer.withoutId) { null } else {
391           val pos = metadata.setIdPosition() + metadata.offset
392           val id = input.readString()
393           length += input.position() - pos
394           id
395         }
396 
397         val nulls = IntBitSet(indices.length)
398 
399         // track the write position for copying attributes in the transformed array
400         var resultCursor = length
401 
402         var i = 0
403         while (i < indices.length) {
404           val index = indices(i) // the index of the attribute in the original feature
405           if (index >= metadata.count || metadata.nulls.contains(index)) { nulls.add(i) } else {
406             // read the offset and the subsequent offset to get the length
407             val pos = metadata.setPosition(index)
408             val len = metadata.setPosition(index + 1) - pos
409             length += len
410             positionsAndLengths(i) = (metadata.offset + pos, len)
411           }
412           i += 1
413         }
414 
415         // note: the input buffer is the raw buffer. we need to ensure that we use the
416         // offset into the raw buffer rather than the raw buffer directly
417         val buf = input.getBuffer
418 
419         val result = Array.ofDim[Byte](length)
420         val output = new Output(result, length)
421         output.writeByte(KryoFeatureSerializer.Version3)
422         output.writeShort(indices.length) // track the number of attributes
423         output.write(metadata.size)
424         i = 0
425         while (i < positionsAndLengths.length) {
426           // note: offset is always 4 here since we're using the full result array
427           if (metadata.size == 2) {
428             output.writeShort(resultCursor - 4)
429           } else {
430             output.writeInt(resultCursor - 4)
431           }
432           if (!nulls.contains(i)) {
433             val (pos, len) = positionsAndLengths(i)
434             System.arraycopy(buf, pos, result, resultCursor, len)
435             resultCursor += len
436           }
437           i += 1
438         }
439         // user data offset
440         if (metadata.size == 2) {
441           output.writeShort(resultCursor)
442         } else {
443           output.writeInt(resultCursor)
444         }
445 
446         // TODO user data?
447 
448         // write out nulls
449         nulls.serialize(output)
450 
451         // we're already positioned to write out the feature id
452         if (id != null) {
453           output.writeString(id)
454         }
455 
456         result
457       }
458     }
459   }
460 
461   /**
462     * Serialiation version 2 delegate
463     *
464     * @param serializer serializer
465     * @param input input
466     */
467   private class KryoBufferV2(serializer: KryoFeatureDeserialization, input: Input) extends KryoBufferDelegate {
468 
469     private val offsets = Array.ofDim[Int](serializer.out.getAttributeCount)
470     private var offset: Int = -1
471     private var startOfOffsets: Int = -1
472     private var missingAttributes: Boolean = false
473     private var userDataOffset: Int = -1
474 
475     private var reserializeTransform: Transformer = _
476     private var binaryTransform: Transformer = _
477 
478     override def reset(): Unit = {
479       // reset our offsets
480       offset = input.position() - 1 // we've already read the version byte
481       startOfOffsets = offset + input.readInt()
482       input.setPosition(startOfOffsets) // set to offsets start
483       var i = 0
484       while (i < offsets.length && input.position < input.limit) {
485         offsets(i) = offset + input.readInt(true)
486         i += 1
487       }
488       if (i < offsets.length) {
489         // attributes have been added to the sft since this feature was serialized
490         missingAttributes = true
491         while ({{ offsets(i) = -1; i += 1 }; i < offsets.length })()
492       } else {
493         missingAttributes = false
494       }
495       userDataOffset = input.position()
496     }
497 
498     override def setTransforms(schema: SimpleFeatureType, transforms: Array[Transform]): Unit = {
499       val indices = transforms.map {
500         case t: PropertyTransform => t.i
501         case t: RenameTransform => t.i
502         case _ => -1
503       }
504 
505       // transforms by evaluating the transform expressions and then serializing the resulting feature
506       // we use this for transform expressions and for data that was written using an old schema
507       reserializeTransform = new ReserializeTransformer(schema, transforms, serializer.options)
508       // if we are just returning a subset of attributes, we can copy the bytes directly
509       // and avoid creating new objects, reserializing, etc
510       binaryTransform =
511           if (indices.contains(-1)) { reserializeTransform } else { new BinaryTransformerV2(input, indices, offsets) }
512     }
513 
514     override def id(): String = {
515       input.setPosition(5)
516       input.readString()
517     }
518 
519     override def getAttribute(index: Int): AnyRef = {
520       val offset = offsets(index)
521       if (offset == -1) { null } else {
522         input.setPosition(offset)
523         serializer.readersV2(index)(input)
524       }
525     }
526 
527     override def getUserData: java.util.Map[AnyRef, AnyRef] = {
528       input.setPosition(userDataOffset)
529       KryoUserDataSerialization.deserialize(input)
530     }
531 
532     override def getDateAsLong(index: Int): Long = {
533       val offset = offsets(index)
534       if (offset == -1) { 0L } else {
535         input.setPosition(offset)
536         KryoFeatureDeserializationV2.LongReader.apply(input)
537       }
538     }
539 
540     override def getInput(index: Int): Option[Input] = {
541       val offset = offsets(index)
542       if (offset == -1) { None } else {
543         input.setPosition(offset)
544         Some(input)
545       }
546     }
547 
548     override def transform(original: SimpleFeature): Array[Byte] = {
549       // if attributes have been added to the sft, we have to reserialize to get the null serialized values
550       val transformer = if (missingAttributes) { reserializeTransform } else { binaryTransform }
551       transformer.transform(original)
552     }
553 
554     /**
555       * Serialization version 2 binary transformer. Copies serialized attribute bytes directly without
556       * deserializing them
557       *
558       * @param input input
559       * @param indices transform indices
560       * @param offsets attribute offsets
561       */
562     private class BinaryTransformerV2(input: Input, indices: Array[Int], offsets: Array[Int]) extends Transformer {
563 
564       private val mutableOffsetsAndLength = Array.ofDim[(Int,Int)](indices.length)
565 
566       override def transform(original: SimpleFeature): Array[Byte] = {
567         // NOTE: the input buffer is the raw buffer. we need to ensure that we use the
568         // offset into the raw buffer rather than the raw buffer directly
569         val buf = input.getBuffer
570         var length = offsets(0) - offset // space for version, offset block and ID
571         var idx = 0
572         while (idx < mutableOffsetsAndLength.length) {
573           val i = indices(idx)
574           val el = (if (i < offsets.length - 1) { offsets(i + 1) } else { startOfOffsets }) - offsets(i)
575           length += el
576           mutableOffsetsAndLength(idx) = (offsets(i), el)
577           idx += 1
578         }
579 
580         val dst = Array.ofDim[Byte](length)
581         // copy the version, offset block and id
582         var dstPos = offsets(0) - offset
583         System.arraycopy(buf, offset, dst, 0, dstPos)
584         mutableOffsetsAndLength.foreach { case (o, l) =>
585           System.arraycopy(buf, o, dst, dstPos, l)
586           dstPos += l
587         }
588         // note that the offset block is incorrect - we couldn't use this in another lazy feature
589         // but the normal serializer doesn't care
590         dst
591       }
592     }
593   }
594 
595   /**
596     * For non-attribute expressions, we have evaluate them, then serialize the resulting feature
597     *
598     * @param schema transform schema
599     * @param transforms transform definitions
600     * @param options serialization options
601     */
602   private class ReserializeTransformer(
603       schema: SimpleFeatureType,
604       transforms: Array[Transform],
605       options: Set[SerializationOption]
606     ) extends Transformer {
607 
608     private val serializer = KryoFeatureSerializer(schema, options)
609     private val sf = new ScalaSimpleFeature(schema, "")
610 
611     override def transform(original: SimpleFeature): Array[Byte] = {
612       sf.setId(original.getID)
613       var i = 0
614       while (i < transforms.length) {
615         sf.setAttribute(i, transforms(i).evaluate(original))
616         i += 1
617       }
618       serializer.serialize(sf)
619     }
620   }
621 }