Add support for Delta lake tables in TableUtils

.circleci/config.yml

@@ -78,6 +78,36 @@ jobs:
                   destination: spark_warehouse.tar.gz
                   when: on_fail
 
+    # run these separately as we need a isolated JVM to not have Spark session settings interfere with other runs
+    # long term goal is to refactor the current testing spark session builder and avoid adding new single test to CI
+    "Scala 13 -- Delta Lake Format Tests":
+      executor: docker_baseimg_executor
+      steps:
+        - checkout
+        - run:
+            name: Run Scala 13 tests for Delta Lake format
+            environment:
+              format_test: deltalake
+            shell: /bin/bash -leuxo pipefail
+            command: |
+              conda activate chronon_py
+              # Increase if we see OOM.
+              export SBT_OPTS="-XX:+CMSClassUnloadingEnabled -XX:MaxPermSize=4G -Xmx4G -Xms2G"
+              sbt '++ 2.13.6' "testOnly ai.chronon.spark.test.TableUtilsFormatTest"
+        - store_test_results:
+            path: /chronon/spark/target/test-reports
+        - store_test_results:
+            path: /chronon/aggregator/target/test-reports
+        - run:
+            name: Compress spark-warehouse
+            command: |
+              cd /tmp/ && tar -czvf spark-warehouse.tar.gz chronon/spark-warehouse
+            when: on_fail
+        - store_artifacts:
+            path: /tmp/spark-warehouse.tar.gz
+            destination: spark_warehouse.tar.gz
+            when: on_fail
+
     "Scala 11 -- Compile":
       executor: docker_baseimg_executor
       steps:
@@ -147,6 +177,9 @@ workflows:
             - "Scala 13 -- Tests":
                   requires:
                       - "Pull Docker Image"
+            - "Scala 13 -- Delta Lake Format Tests":
+                  requires:
+                    - "Pull Docker Image"
             - "Scalafmt Check":
                   requires:
                       - "Pull Docker Image"

build.sbt

@@ -156,6 +156,14 @@ val VersionMatrix: Map[String, VersionDependency] = Map(
     None,
     Some("1.0.4")
   ),
+  "delta-core" -> VersionDependency(
+    Seq(
+      "io.delta" %% "delta-core"
+    ),
+    Some("0.6.1"),
+    Some("1.0.1"),
+    Some("2.0.2")
+  ),
   "jackson" -> VersionDependency(
     Seq(
       "com.fasterxml.jackson.core" % "jackson-core",
@@ -365,7 +373,7 @@ lazy val spark_uber = (project in file("spark"))
     sparkBaseSettings,
     version := git.versionProperty.value,
     crossScalaVersions := supportedVersions,
-    libraryDependencies ++= fromMatrix(scalaVersion.value, "jackson", "spark-all/provided")
+    libraryDependencies ++= fromMatrix(scalaVersion.value, "jackson", "spark-all/provided", "delta-core/provided")
   )
 
 lazy val spark_embedded = (project in file("spark"))
@@ -374,7 +382,7 @@ lazy val spark_embedded = (project in file("spark"))
     sparkBaseSettings,
     version := git.versionProperty.value,
     crossScalaVersions := supportedVersions,
-    libraryDependencies ++= fromMatrix(scalaVersion.value, "spark-all"),
+    libraryDependencies ++= fromMatrix(scalaVersion.value, "spark-all", "delta-core"),
     target := target.value.toPath.resolveSibling("target-embedded").toFile,
     Test / test := {}
   )

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

@@ -149,19 +149,32 @@ class ChrononKryoRegistrator extends KryoRegistrator {
       "scala.collection.immutable.ArraySeq$ofRef",
       "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/Driver.scala

@@ -81,6 +81,9 @@ object Driver {
                    default = Some(false),
                    descr = "Skip the first unfilled partition range if some future partitions have been populated.")
 
+    val useDeltaCatalog: ScallopOption[Boolean] =
+      opt[Boolean](required = false, default = Some(false), descr = "Enable the use of the delta lake catalog")
+
     val stepDays: ScallopOption[Int] =
       opt[Int](required = false,
                descr = "Runs offline backfill in steps, step-days at a time. Default is 30 days",
@@ -136,8 +139,22 @@ object Driver {
     def isLocal: Boolean = localTableMapping.nonEmpty || localDataPath.isDefined
 
     protected def buildSparkSession(): SparkSession = {
+      // use of the delta lake catalog requires a couple of additional spark config options
+      val extraDeltaConfigs = useDeltaCatalog.toOption match {
+        case Some(true) =>
+          Some(
+            Map(
+              "spark.sql.extensions" -> "io.delta.sql.DeltaSparkSessionExtension",
+              "spark.sql.catalog.spark_catalog" -> "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+            ))
+        case _ => None
+      }
+
       if (localTableMapping.nonEmpty) {
-        val localSession = SparkSessionBuilder.build(subcommandName(), local = true, localWarehouseLocation.toOption)
+        val localSession = SparkSessionBuilder.build(subcommandName(),
+                                                     local = true,
+                                                     localWarehouseLocation.toOption,
+                                                     additionalConfig = extraDeltaConfigs)
         localTableMapping.foreach {
           case (table, filePath) =>
             val file = new File(filePath)
@@ -150,13 +167,16 @@ object Driver {
         val localSession =
           SparkSessionBuilder.build(subcommandName(),
                                     local = true,
-                                    localWarehouseLocation = localWarehouseLocation.toOption)
+                                    localWarehouseLocation = localWarehouseLocation.toOption,
+                                    additionalConfig = extraDeltaConfigs)
         LocalDataLoader.loadDataRecursively(dir, localSession)
         localSession
       } else {
         // We use the KryoSerializer for group bys and joins since we serialize the IRs.
         // But since staging query is fairly freeform, it's better to stick to the java serializer.
-        SparkSessionBuilder.build(subcommandName(), enforceKryoSerializer = !subcommandName().contains("staging_query"))
+        SparkSessionBuilder.build(subcommandName(),
+                                  enforceKryoSerializer = !subcommandName().contains("staging_query"),
+                                  additionalConfig = extraDeltaConfigs)
       }
     }
 

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

@@ -22,14 +22,14 @@ import org.apache.spark.SPARK_VERSION
 
 import java.io.File
 import java.util.logging.Logger
-import scala.reflect.io.Path
 import scala.util.Properties
 
 object SparkSessionBuilder {
   @transient private lazy val logger = LoggerFactory.getLogger(getClass)
 
   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 +38,24 @@ 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") =>
+        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
+
     if (local) {
       //required to run spark locally with hive support enabled - for sbt test
       System.setSecurityManager(null)
@@ -61,13 +79,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