Shuffle around kryo registrator

Author: piyush-zlai

spark/src/main/scala/ai/chronon/spark/ChrononKryoRegistrator.scala

@@ -147,23 +147,34 @@ class ChrononKryoRegistrator extends KryoRegistrator {
       "org.apache.spark.sql.catalyst.InternalRow$$anonfun$getAccessor$8",
       "org.apache.spark.sql.catalyst.InternalRow$$anonfun$getAccessor$5",
       "scala.collection.immutable.ArraySeq$ofRef",
-      "org.apache.spark.sql.catalyst.expressions.GenericInternalRow",
-      "org.apache.spark.sql.delta.stats.DeltaFileStatistics",
-      "org.apache.spark.sql.delta.actions.AddFile"
+      "org.apache.spark.sql.catalyst.expressions.GenericInternalRow"
     )
-    names.foreach { name =>
-      try {
-        kryo.register(Class.forName(name))
-        kryo.register(Class.forName(s"[L$name;")) // represents array of a type to jvm
-      } catch {
-        case _: ClassNotFoundException => // do nothing
-      }
-    }
+    names.foreach(name => doRegister(name, kryo))
 
     kryo.register(classOf[Array[Array[Array[AnyRef]]]])
     kryo.register(classOf[Array[Array[AnyRef]]])
     kryo.register(classOf[CpcSketch], new CpcSketchKryoSerializer())
     kryo.register(classOf[Array[ItemSketchSerializable]])
     kryo.register(classOf[ItemsSketchIR[AnyRef]], new ItemsSketchKryoSerializer[AnyRef])
   }
+
+  def doRegister(name: String, kryo: Kryo): Unit = {
+    try {
+      kryo.register(Class.forName(name))
+      kryo.register(Class.forName(s"[L$name;")) // represents array of a type to jvm
+    } catch {
+      case _: ClassNotFoundException => // do nothing
+    }
+  }
+}
+
+class ChrononDeltaLakeKryoRegistrator extends ChrononKryoRegistrator {
+  override def registerClasses(kryo: Kryo): Unit = {
+    super.registerClasses(kryo)
+    val additionalDeltaNames = Seq(
+      "org.apache.spark.sql.delta.stats.DeltaFileStatistics",
+      "org.apache.spark.sql.delta.actions.AddFile"
+    )
+    additionalDeltaNames.foreach(name => doRegister(name, kryo))
+  }
 }

spark/src/main/scala/ai/chronon/spark/SparkSessionBuilder.scala

@@ -42,14 +42,15 @@ object SparkSessionBuilder {
     // allow us to override the format by specifying env vars. This allows us to not have to worry about interference
     // between Spark sessions created in existing chronon tests that need the hive format and some specific tests
     // that require a format override like delta lake.
-    val formatConfigs = sys.env.get(FormatTestEnvVar) match {
+    val (formatConfigs, kryoRegistrator) = sys.env.get(FormatTestEnvVar) match {
       case Some("deltalake") =>
-        Map(
+        val configMap = Map(
           "spark.sql.extensions" -> "io.delta.sql.DeltaSparkSessionExtension",
           "spark.sql.catalog.spark_catalog" -> "org.apache.spark.sql.delta.catalog.DeltaCatalog",
           "spark.chronon.table_write.format" -> "delta"
         )
-      case _ => Map.empty
+        (configMap, "ai.chronon.spark.ChrononDeltaLakeKryoRegistrator")
+      case _ => (Map.empty, "ai.chronon.spark.ChrononKryoRegistrator")
     }
 
     // tack on format configs with additional configs
@@ -78,7 +79,7 @@ object SparkSessionBuilder {
     if (enforceKryoSerializer) {
       baseBuilder
         .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
-        .config("spark.kryo.registrator", "ai.chronon.spark.ChrononKryoRegistrator")
+        .config("spark.kryo.registrator", kryoRegistrator)
         .config("spark.kryoserializer.buffer.max", "2000m")
         .config("spark.kryo.referenceTracking", "false")
     }