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.arrow.vector
10 
11 import org.apache.arrow.vector.FieldVector
12 import org.apache.arrow.vector.dictionary.Dictionary
13 import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding}
14 import org.geotools.api.feature.`type`.AttributeDescriptor
15 import org.locationtech.geomesa.arrow.ArrowAllocator
16 import org.locationtech.geomesa.arrow.vector.SimpleFeatureVector.SimpleFeatureEncoding
17 import org.locationtech.geomesa.utils.geotools.ObjectType
18 import org.locationtech.geomesa.utils.geotools.ObjectType.ObjectType
19 import org.locationtech.geomesa.utils.io.CloseWithLogging
20 
21 import java.io.Closeable
22 import java.util.concurrent.atomic.AtomicInteger
23 import scala.reflect.ClassTag
24 
25 /**
26   * Holder for dictionary values
27   */
28 sealed trait ArrowDictionary extends Closeable {
29 
30   lazy private val map = {
31     val builder = scala.collection.mutable.Map.newBuilder[AnyRef, Int]
32     builder.sizeHint(length)
33     var i = 0
34     foreach { value =>
35       builder += ((value, i))
36       i += 1
37     }
38     builder.result()
39   }
40 
41   def encoding: DictionaryEncoding
42   def id: Long = encoding.getId
43   def length: Int
44 
45   /**
46     * Decode a dictionary int to a value. Note: may not be thread safe
47     *
48     * @param i dictionary encoded int
49     * @return value
50     */
51   def lookup(i: Int): AnyRef
52 
53   /**
54     * Dictionary encode a value to an int
55     *
56     * @param value value to encode
57     * @return dictionary encoded int
58     */
59   def index(value: AnyRef): Int = map.getOrElse(value, length)
60 
61   /**
62     * Apply a function to each value in the dictionary
63     *
64     * @param f function
65     * @tparam U function return type
66     */
67   def foreach[U](f: AnyRef => U): Unit = iterator.foreach(f)
68 
69   /**
70     * Create an iterator over the values in this dictionary
71     *
72     * @return
73     */
74   def iterator: Iterator[AnyRef] = new Iterator[AnyRef] {
75     private var i = 0
76     override def hasNext: Boolean = i < ArrowDictionary.this.length
77     override def next(): AnyRef = try { lookup(i) } finally { i += 1 }
78   }
79 
80   /**
81    * Create an arrow dictionary vector
82    *
83    * @param precision precision
84    * @return
85    */
86   def toDictionary(precision: SimpleFeatureEncoding): Dictionary with Closeable
87 }
88 
89 object ArrowDictionary {
90 
91   /**
92     * Create a dictionary based off a sequence of values. Encoding will be smallest that will fit all values.
93     *
94     * @param id dictionary id
95     * @param values dictionary values
96     * @return dictionary
97     */
98   def create[T <: AnyRef](typename: String, id: Long, values: Array[T])(implicit ct: ClassTag[T]): ArrowDictionary =
99     create(typename, id, values, values.length)
100 
101   /**
102     * Create a dictionary based on a subset of a value array
103     *
104     * @param id dictionary id
105     * @param values array of dictionary values
106     * @param length number of valid entries in the values array, starting at position 0
107     * @return
108     */
109   def create[T <: AnyRef](typename: String, id: Long, values: Array[T], length: Int)(implicit ct: ClassTag[T]): ArrowDictionary =
110     new ArrowDictionaryArray[T](typename, createEncoding(id, length), values, length, ct.runtimeClass.asInstanceOf[Class[T]])
111 
112   /**
113     * Create a dictionary based on wrapping an arrow vector
114     *
115     * @param encoding dictionary id and metadata
116     * @param values dictionary vector
117     * @param descriptor attribute descriptor for the dictionary, used to read values from the underlying vector
118     * @param precision simple feature encoding used on the dictionary values
119     * @return
120     */
121   def create(
122       encoding: DictionaryEncoding,
123       values: FieldVector,
124       descriptor: AttributeDescriptor,
125       precision: SimpleFeatureEncoding): ArrowDictionary = {
126     new ArrowDictionaryVector(encoding, values, ObjectType.selectType(descriptor), precision)
127   }
128 
129   /**
130    * Create a dictionary based on wrapping an arrow vector
131    *
132    * @param encoding dictionary id and metadata
133    * @param values dictionary vector
134    * @param bindings attribute descriptor bindings, used for reading values from the arrow vector
135    * @param precision simple feature encoding used on the dictionary values
136    * @return
137    */
138   def create(
139       encoding: DictionaryEncoding,
140       values: FieldVector,
141       bindings: Seq[ObjectType],
142       precision: SimpleFeatureEncoding): ArrowDictionary = {
143     new ArrowDictionaryVector(encoding, values, bindings, precision)
144   }
145 
146   /**
147     * Holder for dictionary values
148     *
149     * @param values dictionary values. When encoded, values are replaced with their index in the seq
150     * @param encoding dictionary id and int width, id must be unique per arrow file
151     */
152   class ArrowDictionaryArray[T <: AnyRef](
153       typename: String,
154       val encoding: DictionaryEncoding,
155       values: Array[T],
156       val length: Int,
157       binding: Class[T]
158     ) extends ArrowDictionary {
159 
160     override def lookup(i: Int): AnyRef = if (i < length) { values(i) } else { "[other]" }
161 
162     override def toDictionary(precision: SimpleFeatureEncoding): Dictionary with Closeable = {
163       val allocator = ArrowAllocator(s"dictionary-array:$typename")
164       val name = s"dictionary-$id"
165       val bindings = ObjectType.selectType(binding)
166       val writer = ArrowAttributeWriter(name, bindings, None, Map.empty, precision, VectorFactory(allocator))
167       var i = 0
168       while (i < length) {
169         writer.apply(i, values(i))
170         i += 1
171       }
172       writer.setValueCount(length)
173 
174       new Dictionary(writer.vector, encoding) with Closeable {
175         override def close(): Unit = CloseWithLogging.raise(writer.vector, allocator)
176       }
177     }
178 
179     override def close(): Unit = {}
180   }
181 
182   /**
183     * Dictionary that wraps an arrow vector
184     *
185     * @param encoding dictionary id and metadata
186     * @param vector arrow vector
187     * @param bindings attribute descriptor bindings, used for reading values from the arrow vector
188     * @param precision simple feature encoding used for the arrow vector
189     */
190   class ArrowDictionaryVector(
191       val encoding: DictionaryEncoding,
192       vector: FieldVector,
193       bindings: Seq[ObjectType],
194       precision: SimpleFeatureEncoding
195     ) extends ArrowDictionary {
196 
197     // we use an attribute reader to get the right type conversion
198     private val reader = ArrowAttributeReader(bindings, vector, None, precision)
199     private var references = 1
200 
201     override val length: Int = vector.getValueCount
202 
203     override def lookup(i: Int): AnyRef = if (i < length) { reader.apply(i) } else { "[other]" }
204 
205     override def toDictionary(precision: SimpleFeatureEncoding): Dictionary with Closeable = synchronized {
206       if (references < 1) {
207         throw new IllegalStateException("Trying to create a dictionary from a closed vector")
208       } else if (precision != this.precision) {
209         throw new IllegalArgumentException("Wrapped dictionary vectors can't be re-encoded with a different precision")
210       }
211       references += 1
212       new Dictionary(vector, encoding) with Closeable {
213         override def close(): Unit = ArrowDictionaryVector.this.close()
214       }
215     }
216 
217     override def close(): Unit = synchronized {
218       references -= 1
219       if (references < 1) {
220         vector.close()
221       }
222     }
223   }
224 
225   /**
226    * Tracks values seen and writes dictionary encoded ints instead
227    */
228   class ArrowDictionaryBuilder(val encoding: DictionaryEncoding) extends ArrowDictionary {
229 
230     // next dictionary index to use
231     private val counter = new AtomicInteger(0)
232 
233     // values that we have seen, and their dictionary index
234     private val values = scala.collection.mutable.LinkedHashMap.empty[AnyRef, Int]
235 
236     override def lookup(i: Int): AnyRef = values.find(_._2 == i).map(_._1).getOrElse("[other]")
237 
238     override def index(value: AnyRef): Int = values.getOrElseUpdate(value, counter.getAndIncrement())
239 
240     override def length: Int = values.size
241 
242     // note: iterator will return in insert order
243     // we need to keep it ordered so that dictionary values match up with their index
244     override def iterator: Iterator[AnyRef] = values.keys.iterator
245 
246     override def toDictionary(precision: SimpleFeatureEncoding): Dictionary with Closeable =
247       throw new NotImplementedError()
248 
249     override def close(): Unit = {}
250 
251     def clear(): Unit = {
252       counter.set(0)
253       values.clear()
254     }
255   }
256 
257   // use the smallest int type possible to minimize bytes used
258   private def createEncoding(id: Long, count: Int): DictionaryEncoding = {
259     // we check `MaxValue - 1` to allow for the fallback 'other'
260     if (count < Byte.MaxValue - 1) {
261       new DictionaryEncoding(id, false, new ArrowType.Int(8, true))
262     } else if (count < Short.MaxValue - 1) {
263       new DictionaryEncoding(id, false, new ArrowType.Int(16, true))
264     } else {
265       new DictionaryEncoding(id, false, new ArrowType.Int(32, true))
266     }
267   }
268 }