Make Serializers thread local

Author: pomadchin

core/src/main/scala/org/locationtech/rasterframes/encoders/SerializersCache.scala

@@ -4,67 +4,58 @@ import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
 
-import scala.collection.concurrent.TrieMap
-
+import scala.collection.mutable
 import scala.reflect.runtime.universe.TypeTag
 
-object SerializersCache { self =>
+object SerializersCache {
   /**
-   * The point of {Serizalizer | Deserializer} wrappers to make application atomic.
-   * If that is the chain of encoders, i.e. T <=> InternalRow <=> Row the whole chain applciation should be atomic.
+   * Spark partitions are executed on a blocking thread pool.
+   * We can keep the cache of (De)Serializers (every serializer instance creation is pretty expensive),
+   * but the cache should be local per thread.
+   *
+   * When used from multiple threads (De)Serializers tend to corrupt data and / or fail at runtime.
+   * The alternative can be to use global locks or to use a separate executor per each (De)Serializer.
    */
-  case class DeserializerSynchronized[T](underlying: ExpressionEncoder.Deserializer[T]) {
-    def apply(i: InternalRow): T = self.synchronized(underlying.apply(i))
-  }
-
-  case class SerializerSynchronized[T](underlying: ExpressionEncoder.Serializer[T]) {
-    // copy should happen within the same lock, otherwise we're risking to loose the InternalRow
-    def apply(t: T): InternalRow = self.synchronized(underlying.apply(t).copy())
+  private class ThreadLocalHashMap[K, V] extends ThreadLocal[mutable.HashMap[K, V]] {
+    override def initialValue(): mutable.HashMap[K, V] = mutable.HashMap.empty
   }
-
-  case class DeserializerRowSynchronized[T](underlying: Row => T) extends AnyVal {
-    def apply(i: Row): T = self.synchronized(underlying(i))
+  private object ThreadLocalHashMap {
+    def empty[K, V]: ThreadLocalHashMap[K, V] = new ThreadLocalHashMap
   }
 
-  case class SerializerRowSynchronized[T](underlying: T => Row) extends AnyVal {
-    def apply(i: T): Row = self.synchronized(underlying(i))
+  /** SerializerSafe ensures that all Serializers from the pool call copy after application. */
+  case class SerializerSafe[T](underlying: ExpressionEncoder.Serializer[T]) {
+    def apply(t: T): InternalRow = underlying.apply(t).copy()
   }
 
-  private val cacheSerializer: TrieMap[TypeTag[_], SerializerSynchronized[_]] = TrieMap.empty
-  private val cacheSerializerRow: TrieMap[TypeTag[_], SerializerSynchronized[Row]] = TrieMap.empty
-  private val cacheDeserializer: TrieMap[TypeTag[_], DeserializerSynchronized[_]] = TrieMap.empty
-  private val cacheDeserializerRow: TrieMap[TypeTag[_], DeserializerSynchronized[Row]] = TrieMap.empty
-
-  /** Serializer is threadsafe.*/
-  def serializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): SerializerSynchronized[T] =
-    cacheSerializer
-      .getOrElseUpdate(tag, SerializerSynchronized(encoder.createSerializer()))
-      .asInstanceOf[SerializerSynchronized[T]]
-
-  def rowSerializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): SerializerSynchronized[Row] =
-    cacheSerializerRow.getOrElseUpdate(tag, SerializerSynchronized(RowEncoder(encoder.schema).createSerializer()))
+  // T => InternalRow
+  private val cacheSerializer: ThreadLocalHashMap[TypeTag[_], SerializerSafe[_]] = ThreadLocalHashMap.empty
+  // Row with Schema T => InternalRow
+  private val cacheSerializerRow: ThreadLocalHashMap[TypeTag[_], SerializerSafe[Row]] = ThreadLocalHashMap.empty
+  // InternalRow => T
+  private val cacheDeserializer: ThreadLocalHashMap[TypeTag[_], ExpressionEncoder.Deserializer[_]] = ThreadLocalHashMap.empty
+  // InternalRow => Row with Schema T
+  private val cacheDeserializerRow: ThreadLocalHashMap[TypeTag[_], ExpressionEncoder.Deserializer[Row]] = ThreadLocalHashMap.empty
 
-  /** Both Serializer and Deserializer are not thread safe, and expensive to derive.
-   * Per partition instance would give us no performance regressions,
-   * however would require a significant DynamicExtractors refactor. */
-  def deserializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): DeserializerSynchronized[T] =
-    cacheDeserializer
-      .getOrElseUpdate(tag, DeserializerSynchronized(encoder.resolveAndBind().createDeserializer()))
-      .asInstanceOf[DeserializerSynchronized[T]]
+  def serializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): SerializerSafe[T] =
+    cacheSerializer.get.getOrElseUpdate(tag, SerializerSafe(encoder.createSerializer())).asInstanceOf[SerializerSafe[T]]
 
+  def rowSerializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): SerializerSafe[Row] =
+    cacheSerializerRow.get.getOrElseUpdate(tag, SerializerSafe(RowEncoder(encoder.schema).createSerializer()))
 
-  def rowDeserializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): DeserializerSynchronized[Row] =
-    cacheDeserializerRow
-      .getOrElseUpdate(tag, DeserializerSynchronized(RowEncoder(encoder.schema).resolveAndBind().createDeserializer()))
+  def deserializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): ExpressionEncoder.Deserializer[T] =
+    cacheDeserializer.get.getOrElseUpdate(tag, encoder.resolveAndBind().createDeserializer()).asInstanceOf[ExpressionEncoder.Deserializer[T]]
 
+  def rowDeserializer[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): ExpressionEncoder.Deserializer[Row] =
+    cacheDeserializerRow.get.getOrElseUpdate(tag, RowEncoder(encoder.schema).resolveAndBind().createDeserializer())
 
   /**
    * https://jaceklaskowski.gitbooks.io/mastering-spark-sql/content/spark-sql-RowEncoder.html
    * https://github.com/apache/spark/blob/93cec49212fe82816fcadf69f429cebaec60e058/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala#L75-L86
    */
-  def rowDeserialize[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): DeserializerRowSynchronized[T] =
-    DeserializerRowSynchronized { row => deserializer[T](tag, encoder)(rowSerializer[T](tag, encoder)(row)) }
+  def rowDeserialize[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): Row => T =
+    { row => deserializer[T](tag, encoder)(rowSerializer[T](tag, encoder)(row)) }
 
-  def rowSerialize[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): SerializerRowSynchronized[T] =
-    SerializerRowSynchronized[T] ({ t => rowDeserializer[T](tag, encoder)(serializer[T](tag, encoder)(t)) })
+  def rowSerialize[T](implicit tag: TypeTag[T], encoder: ExpressionEncoder[T]): T => Row =
+    { t => rowDeserializer[T](tag, encoder)(serializer[T](tag, encoder)(t)) }
 }