Change GpuKryoRegistrator to load the classes we want to register wit…

…h the ShimLoader (NVIDIA#6475) Signed-off-by: Thomas Graves <[email protected]> Signed-off-by: Thomas Graves <[email protected]>
razajafri · Sep 1, 2022 · 356d32b · 356d32b
1 parent 399438f
commit 356d32b
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuKryoRegistrator.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuKryoRegistrator.scala
@@ -20,13 +20,14 @@ import com.esotericsoftware.kryo.Kryo
 import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer}
 
 import org.apache.spark.serializer.KryoRegistrator
-import org.apache.spark.sql.rapids.execution.{SerializeBatchDeserializeHostBuffer, SerializeConcatHostBuffersDeserializeBatch}
 
 class GpuKryoRegistrator extends KryoRegistrator {
   override def registerClasses(kryo: Kryo): Unit = {
-    kryo.register(classOf[SerializeConcatHostBuffersDeserializeBatch],
-      new KryoJavaSerializer())
-    kryo.register(classOf[SerializeBatchDeserializeHostBuffer],
-      new KryoJavaSerializer())
+    val allClassesToRegister = Seq(
+      "org.apache.spark.sql.rapids.execution.SerializeConcatHostBuffersDeserializeBatch",
+      "org.apache.spark.sql.rapids.execution.SerializeBatchDeserializeHostBuffer")
+    allClassesToRegister.foreach { classToRegister =>
+      kryo.register(ShimLoader.loadClass(classToRegister), new KryoJavaSerializer())
+    }
   }
 }