Add delta lake integration (#51)

## Summary
Port of our OSS delta lake PR -
airbnb/chronon#869. Largely the same aside from
delta lake versions. We don't need this immediately atm but we'll need
this if we have other users come along that need delta lake (or we need
to add support for formats like hudi)

## Checklist
- [X] Added Unit Tests
- [X] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added support for Delta Lake operations with new dependencies and
configurations.
- Introduced new traits and case objects for handling different table
formats, enhancing data management capabilities.
- Added a new job in the CI workflow for testing Delta Lake format
functionality.

- **Bug Fixes**
	- Improved error handling in class registration processes.

- **Tests**
- Implemented a suite of unit tests for the `TableUtils` class to
validate partitioned data insertions with schema modifications.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: piyush-zlai

.github/workflows/test_scala_spark.yaml

@@ -59,6 +59,27 @@ jobs:
           export SBT_OPTS="-Xmx8G -Xms2G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
           sbt "spark/testOnly -- -n jointest"
 
+  table_utils_delta_format_spark_tests:
+    runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/${{ github.repository }}-ci:latest
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+    defaults:
+      run:
+        working-directory: ${{ github.workspace }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Run table utils format test for Delta Lake
+        run: |
+          export SBT_OPTS="-Xmx8G -Xms2G --add-opens=java.base/sun.nio.ch=ALL-UNNAMED"
+          sbt "spark/testOnly ai.chronon.spark.test.TableUtilsFormatTest"
+        env:
+          format_test: deltalake
+  
   mutation_spark_tests:
     runs-on: ubuntu-latest
     container:

build.sbt

@@ -29,6 +29,7 @@ lazy val flink_1_17 = "1.17.0"
 lazy val jackson_2_15 = "2.15.2"
 lazy val avro_1_11 = "1.11.2"
 lazy val circeVersion = "0.14.9"
+lazy val deltaVersion = "3.2.0"
 
 // skip tests on assembly - uncomment if builds become slow
 // ThisBuild / assembly / test := {}
@@ -61,6 +62,10 @@ val spark_sql = Seq(
 ).map(_ % spark_3_5)
 val spark_sql_provided = spark_sql.map(_ % "provided")
 
+val delta = Seq(
+  "io.delta" %% "delta-spark"
+).map(_ % deltaVersion)
+
 val spark_all = Seq(
   "org.apache.spark" %% "spark-sql",
   "org.apache.spark" %% "spark-hive",
@@ -172,7 +177,8 @@ lazy val spark = project
     libraryDependencies ++= spark_all.map(_ % "test"),
     libraryDependencies += "jakarta.servlet" % "jakarta.servlet-api" % "4.0.3",
     libraryDependencies += "com.google.guava" % "guava" % "33.3.1-jre",
-    libraryDependencies ++= log4j2
+    libraryDependencies ++= log4j2,
+    libraryDependencies ++= delta.map(_ % "provided")
   )
 
 lazy val flink = project

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

@@ -170,19 +170,32 @@ class ChrononKryoRegistrator extends KryoRegistrator {
       "org.apache.datasketches.kll.KllSketch$SketchStructure",
       "org.apache.datasketches.kll.KllSketch$SketchType"
     )
-    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

@@ -29,6 +29,7 @@ object SparkSessionBuilder {
 
   private val warehouseId = java.util.UUID.randomUUID().toString.takeRight(6)
   private val DefaultWarehouseDir = new File("/tmp/chronon/spark-warehouse_" + warehouseId)
+  val FormatTestEnvVar: String = "format_test"
 
   def expandUser(path: String): String = path.replaceFirst("~", System.getProperty("user.home"))
   // we would want to share locally generated warehouse during CI testing
@@ -38,6 +39,25 @@ object SparkSessionBuilder {
             localWarehouseLocation: Option[String] = None,
             additionalConfig: Option[Map[String, String]] = None,
             enforceKryoSerializer: Boolean = true): SparkSession = {
+
+    // 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, kryoRegistrator) = sys.env.get(FormatTestEnvVar) match {
+      case Some("deltalake") =>
+        logger.info("Using the delta lake table format + kryo registrators")
+        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"
+        )
+        (configMap, "ai.chronon.spark.ChrononDeltaLakeKryoRegistrator")
+      case _ => (Map.empty, "ai.chronon.spark.ChrononKryoRegistrator")
+    }
+
+    // tack on format configs with additional configs
+    val mergedConfigs = additionalConfig.getOrElse(Map.empty) ++ formatConfigs
+
     val userName = Properties.userName
     val warehouseDir = localWarehouseLocation.map(expandUser).getOrElse(DefaultWarehouseDir.getAbsolutePath)
     println(s"Using warehouse dir: $warehouseDir")
@@ -62,13 +82,12 @@ 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")
     }
-    additionalConfig.foreach { configMap =>
-      configMap.foreach { config => baseBuilder = baseBuilder.config(config._1, config._2) }
-    }
+
+    mergedConfigs.foreach { config => baseBuilder = baseBuilder.config(config._1, config._2) }
 
     if (SPARK_VERSION.startsWith("2")) {
       // Otherwise files left from deleting the table with the same name result in test failures