Reuse stencil client between Spark Tasks (#58)

pyalex · web-flow · commit 5730cf4329f6 · 2021-04-19T11:27:53.000+08:00
* reuse stencil client between tasks

Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;

* format

Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/registry/proto/StencilProtoRegistry.scala b/spark/ingestion/src/main/scala/feast/ingestion/registry/proto/StencilProtoRegistry.scala
@@ -22,17 +22,21 @@ import com.gojek.de.stencil.StencilClientFactory
 import com.gojek.de.stencil.client.StencilClient
 
 class StencilProtoRegistry(val url: String) extends ProtoRegistry {
+  import StencilProtoRegistry.stencilClient
+
+  override def getProtoDescriptor(className: String): Descriptors.Descriptor = {
+    stencilClient(url).get(className)
+  }
+}
+
+object StencilProtoRegistry {
   @transient
   private var _stencilClient: StencilClient = _
 
-  def stencilClient: StencilClient = {
+  def stencilClient(url: String): StencilClient = {
     if (_stencilClient == null) {
       _stencilClient = StencilClientFactory.getClient(url, Collections.emptyMap[String, String])
     }
     _stencilClient
   }
-
-  override def getProtoDescriptor(className: String): Descriptors.Descriptor = {
-    stencilClient.get(className)
-  }
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/utils/ProtoReflection.scala b/spark/ingestion/src/main/scala/feast/ingestion/utils/ProtoReflection.scala
@@ -124,15 +124,13 @@ object ProtoReflection {
   }
 
   def createMessageParser(protoRegistry: ProtoRegistry, className: String): Array[Byte] => Row = {
-    // perform request to registry in driver, so serialized protoRegistry will have cached descriptor
-    protoRegistry.getProtoDescriptor(className)
+    bytes =>
+      {
+        val protoDescriptor = protoRegistry.getProtoDescriptor(className)
 
-    bytes => {
-      val protoDescriptor = protoRegistry.getProtoDescriptor(className)
-
-      Try { DynamicMessage.parseFrom(protoDescriptor, bytes) }
-        .map(messageToRow(protoDescriptor, _))
-        .getOrElse(null)
-    }
+        Try { DynamicMessage.parseFrom(protoDescriptor, bytes) }
+          .map(messageToRow(protoDescriptor, _))
+          .getOrElse(null)
+      }
   }
 }
