CatalogAwareDataPointer

Author: tchow-zlai

api/src/main/scala/ai/chronon/api/DataPointer.scala

@@ -1,10 +1,20 @@
 package ai.chronon.api
 import scala.util.parsing.combinator._
 
-case class DataPointer(catalog: Option[String],
-                       tableOrPath: String,
-                       format: Option[String],
-                       options: Map[String, String])
+abstract class DataPointer {
+  def tableOrPath: String
+  def readFormat: Option[String]
+  def writeFormat: Option[String]
+  def options: Map[String, String]
+
+}
+
+case class URIDataPointer(
+    override val tableOrPath: String,
+    override val readFormat: Option[String],
+    override val writeFormat: Option[String],
+    override val options: Map[String, String]
+) extends DataPointer
 
 // parses string representations of data pointers
 // ex: namespace.table
@@ -27,21 +37,26 @@ object DataPointer extends RegexParsers {
     opt(catalogWithOptionalFormat ~ opt(options) ~ "://") ~ tableOrPath ^^ {
       // format is specified in the prefix s3+parquet://bucket/path/to/data/*/*/
       // note that if you have s3+parquet://bucket/path/to/data.csv, format is still parquet
-      case Some((ctl, Some(fmt)) ~ opts ~ _) ~ path =>
-        DataPointer(Some(ctl), path, Some(fmt), opts.getOrElse(Map.empty))
+      case Some((ctl, Some(fmt)) ~ opts ~ sep) ~ path =>
+        URIDataPointer(ctl + sep + path, Some(fmt), Some(fmt), opts.getOrElse(Map.empty))
 
       // format is extracted from the path for relevant sources
       // ex: s3://bucket/path/to/data.parquet
       // ex: file://path/to/data.csv
       // ex: hdfs://path/to/data.with.dots.parquet
       // for other sources like bigquery, snowflake, format is None
-      case Some((ctl, None) ~ opts ~ _) ~ path =>
-        val (pathWithoutFormat, fmt) = extractFormatFromPath(path, ctl)
-        DataPointer(Some(ctl), path, fmt, opts.getOrElse(Map.empty))
+      case Some((ctl, None) ~ opts ~ sep) ~ path =>
+        val (_, fmt) = extractFormatFromPath(path, ctl)
+
+        fmt match {
+          // Retain the full uri if it's a path.
+          case Some(ft) => URIDataPointer(ctl + sep + path, Some(ft), Some(ft), opts.getOrElse(Map.empty))
+          case None     => URIDataPointer(path, Some(ctl), Some(ctl), opts.getOrElse(Map.empty))
+        }
 
       case None ~ path =>
         // No prefix case (direct table reference)
-        DataPointer(None, path, None, Map.empty)
+        URIDataPointer(path, None, None, Map.empty)
     }
 
   private def catalogWithOptionalFormat: Parser[(String, Option[String])] =

api/src/test/scala/ai/chronon/api/test/DataPointerTest.scala

@@ -1,67 +1,70 @@
 package ai.chronon.api.test
 
 import ai.chronon.api.DataPointer
+import ai.chronon.api.URIDataPointer
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 
 class DataPointerTest extends AnyFlatSpec with Matchers {
 
   "DataPointer.apply" should "parse a simple s3 path" in {
     val result = DataPointer("s3://bucket/path/to/data.parquet")
-    result should be(DataPointer(Some("s3"), "bucket/path/to/data.parquet", Some("parquet"), Map.empty))
+    result should be(URIDataPointer("s3://bucket/path/to/data.parquet", Some("parquet"), Some("parquet"), Map.empty))
   }
 
   it should "parse a bigquery table with options" in {
     val result = DataPointer("bigquery(option1=value1,option2=value2)://project-id.dataset.table")
     result should be(
-      DataPointer(Some("bigquery"),
-                  "project-id.dataset.table",
-                  None,
-                  Map("option1" -> "value1", "option2" -> "value2")))
+      URIDataPointer("project-id.dataset.table",
+                     Some("bigquery"),
+                     Some("bigquery"),
+                     Map("option1" -> "value1", "option2" -> "value2")))
   }
 
   it should "parse a bigquery table without options" in {
     val result = DataPointer("bigquery://project-id.dataset.table")
-    result should be(DataPointer(Some("bigquery"), "project-id.dataset.table", None, Map.empty))
+    result should be(URIDataPointer("project-id.dataset.table", Some("bigquery"), Some("bigquery"), Map.empty))
   }
 
   it should "parse a kafka topic" in {
     val result = DataPointer("kafka://my-topic")
-    result should be(DataPointer(Some("kafka"), "my-topic", None, Map.empty))
+    result should be(URIDataPointer("my-topic", Some("kafka"), Some("kafka"), Map.empty))
   }
 
   it should "parse a file path with format" in {
     val result = DataPointer("file://path/to/data.csv")
-    result should be(DataPointer(Some("file"), "path/to/data.csv", Some("csv"), Map.empty))
+    result should be(URIDataPointer("file://path/to/data.csv", Some("csv"), Some("csv"), Map.empty))
   }
 
   it should "parse options with spaces" in {
     val result = DataPointer("hive(key1 = value1, key2 = value2)://database.table")
-    result should be(DataPointer(Some("hive"), "database.table", None, Map("key1" -> "value1", "key2" -> "value2")))
+    result should be(
+      URIDataPointer("database.table", Some("hive"), Some("hive"), Map("key1" -> "value1", "key2" -> "value2")))
   }
 
   it should "handle paths with dots" in {
     val result = DataPointer("hdfs://path/to/data.with.dots.parquet")
-    result should be(DataPointer(Some("hdfs"), "path/to/data.with.dots.parquet", Some("parquet"), Map.empty))
+    result should be(
+      URIDataPointer("hdfs://path/to/data.with.dots.parquet", Some("parquet"), Some("parquet"), Map.empty))
   }
 
   it should "handle paths with multiple dots and no format" in {
     val result = DataPointer("file://path/to/data.with.dots")
-    result should be(DataPointer(Some("file"), "path/to/data.with.dots", Some("dots"), Map.empty))
+    result should be(URIDataPointer("file://path/to/data.with.dots", Some("dots"), Some("dots"), Map.empty))
   }
 
   it should "handle paths with multiple dots and prefixed format" in {
     val result = DataPointer("file+csv://path/to/data.with.dots")
-    result should be(DataPointer(Some("file"), "path/to/data.with.dots", Some("csv"), Map.empty))
+    result should be(URIDataPointer("file://path/to/data.with.dots", Some("csv"), Some("csv"), Map.empty))
   }
 
   it should "handle paths with format and pointer to folder with glob matching" in {
     val result = DataPointer("s3+parquet://path/to/*/*/")
-    result should be(DataPointer(Some("s3"), "path/to/*/*/", Some("parquet"), Map.empty))
+    result should be(URIDataPointer("s3://path/to/*/*/", Some("parquet"), Some("parquet"), Map.empty))
   }
 
   it should "handle no catalog, just table" in {
     val result = DataPointer("namespace.table")
-    result should be(DataPointer(None, "namespace.table", None, Map.empty))
+    result should be(URIDataPointer("namespace.table", None, None, Map.empty))
   }
 }

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

@@ -0,0 +1,42 @@
+package ai.chronon.spark
+
+import ai.chronon.api.DataPointer
+import org.apache.spark.sql.SparkSession
+
+import scala.reflect.runtime.universe._
+
+case class CatalogAwareDataPointer(inputTableOrPath: String, formatProvider: FormatProvider) extends DataPointer {
+
+  override def tableOrPath: String = {
+    formatProvider.resolveTableName(inputTableOrPath)
+  }
+  override lazy val options: Map[String, String] = Map.empty
+
+  override lazy val readFormat: Option[String] = {
+    Option(formatProvider.readFormat(inputTableOrPath)).map(_.name)
+  }
+
+  override lazy val writeFormat: Option[String] = {
+    Option(formatProvider.writeFormat(inputTableOrPath)).map(_.name)
+  }
+
+}
+
+object DataPointer {
+
+  def apply(tableOrPath: String, sparkSession: SparkSession): DataPointer = {
+    val clazzName =
+      sparkSession.conf.get("spark.chronon.table.format_provider.class", classOf[DefaultFormatProvider].getName)
+    val mirror = runtimeMirror(getClass.getClassLoader)
+    val classSymbol = mirror.staticClass(clazzName)
+    val classMirror = mirror.reflectClass(classSymbol)
+    val constructor = classSymbol.primaryConstructor.asMethod
+    val constructorMirror = classMirror.reflectConstructor(constructor)
+    val reflected = constructorMirror(sparkSession)
+    val provider = reflected.asInstanceOf[FormatProvider]
+
+    CatalogAwareDataPointer(tableOrPath, provider)
+
+  }
+
+}

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

@@ -26,8 +26,9 @@ import ai.chronon.online.SparkConversions
 import ai.chronon.online.TimeRange
 import org.apache.avro.Schema
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.DataFrameReader
+import org.apache.spark.sql.DataFrameWriter
 import org.apache.spark.sql.Row
-import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.functions._
@@ -322,53 +323,80 @@ object Extensions {
     }
   }
 
-  implicit class DataPointerOps(dataPointer: DataPointer) {
-    def toDf(implicit sparkSession: SparkSession): DataFrame = {
+  implicit class DataPointerAwareDataFrameWriter[T](dfw: DataFrameWriter[T]) {
+
+    def save(dataPointer: DataPointer): Unit = {
+
+      dataPointer.writeFormat
+        .map((wf) => {
+          val normalized = wf.toLowerCase
+          normalized match {
+            case "bigquery" | "bq" =>
+              dfw
+                .format("bigquery")
+                .options(dataPointer.options)
+                .save(dataPointer.tableOrPath)
+            case "snowflake" | "sf" =>
+              dfw
+                .format("net.snowflake.spark.snowflake")
+                .options(dataPointer.options)
+                .option("dbtable", dataPointer.tableOrPath)
+                .save()
+            case "parquet" | "csv" =>
+              dfw
+                .format(normalized)
+                .options(dataPointer.options)
+                .save(dataPointer.tableOrPath)
+            case "hive" =>
+              dfw
+                .format("hive")
+                .saveAsTable(dataPointer.tableOrPath)
+            case _ =>
+              throw new UnsupportedOperationException(s"Unsupported write catalog: ${normalized}")
+          }
+        })
+        .getOrElse(
+          // None case is just table against default catalog
+          dfw
+            .format("hive")
+            .saveAsTable(dataPointer.tableOrPath))
+    }
+  }
+
+  implicit class DataPointerAwareDataFrameReader(dfr: DataFrameReader) {
+
+    def load(dataPointer: DataPointer): DataFrame = {
       val tableOrPath = dataPointer.tableOrPath
-      val format = dataPointer.format.getOrElse("parquet")
-      dataPointer.catalog.map(_.toLowerCase) match {
-        case Some("bigquery") | Some("bq") =>
-          // https://github.com/GoogleCloudDataproc/spark-bigquery-connector?tab=readme-ov-file#reading-data-from-a-bigquery-table
-          sparkSession.read
-            .format("bigquery")
-            .options(dataPointer.options)
-            .load(tableOrPath)
-
-        case Some("snowflake") | Some("sf") =>
-          // https://docs.snowflake.com/en/user-guide/spark-connector-use#moving-data-from-snowflake-to-spark
-          val sfOptions = dataPointer.options
-          sparkSession.read
-            .format("net.snowflake.spark.snowflake")
-            .options(sfOptions)
-            .option("dbtable", tableOrPath)
-            .load()
-
-        case Some("s3") | Some("s3a") | Some("s3n") =>
-          // https://sites.google.com/site/hellobenchen/home/wiki/big-data/spark/read-data-files-from-multiple-sub-folders
-          // "To get spark to read through all subfolders and subsubfolders, etc. simply use the wildcard *"
-          // "df= spark.read.parquet('/datafolder/*/*')"
-          //
-          // https://docs.aws.amazon.com/emr/latest/ManagementGuide/emr-plan-file-systems.html
-          // "Previously, Amazon EMR used the s3n and s3a file systems. While both still work, "
-          // "we recommend that you use the s3 URI scheme for the best performance, security, and reliability."
-          // TODO: figure out how to scan subfolders in a date range without reading the entire folder
-          sparkSession.read
-            .format(format)
-            .options(dataPointer.options)
-            .load("ș3://" + tableOrPath)
-
-        case Some("file") =>
-          sparkSession.read
-            .format(format)
-            .options(dataPointer.options)
-            .load(tableOrPath)
-
-        case Some("hive") | None =>
-          sparkSession.table(tableOrPath)
-
-        case _ =>
-          throw new UnsupportedOperationException(s"Unsupported catalog: ${dataPointer.catalog}")
-      }
+
+      dataPointer.readFormat
+        .map((fmt) => {
+          val normalized = fmt.toLowerCase
+          normalized match {
+            case "bigquery" | "bq" =>
+              dfr
+                .format("bigquery")
+                .options(dataPointer.options)
+                .load(tableOrPath)
+            case "snowflake" | "sf" =>
+              dfr
+                .format("net.snowflake.spark.snowflake")
+                .options(dataPointer.options)
+                .option("dbtable", tableOrPath)
+                .load()
+            case "parquet" | "csv" =>
+              dfr
+                .format(normalized)
+                .options(dataPointer.options)
+                .load(tableOrPath)
+            case "hive" => dfr.table(tableOrPath)
+            case _ =>
+              throw new UnsupportedOperationException(s"Unsupported read catalog: ${normalized}")
+          }
+        })
+        .getOrElse {
+          // None case is just table against default catalog
+          dfr.table(tableOrPath)
+        }
     }
   }
 }