fix: supply partition column only when needed

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

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

@@ -5,16 +5,22 @@ abstract class DataPointer {
   def tableOrPath: String
   def readFormat: Option[String]
   def writeFormat: Option[String]
-  def options: Map[String, String]
+
+  def readOptions: Map[String, String]
+  def writeOptions: 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
+    options: Map[String, String]
+) extends DataPointer {
+
+  override val readOptions: Map[String, String] = options
+  override val writeOptions: Map[String, String] = options
+}
 
 // parses string representations of data pointers
 // ex: namespace.table

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

@@ -5,8 +5,7 @@ import ai.chronon.api.URIDataPointer
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 
-class
-DataPointerTest extends AnyFlatSpec with 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")

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala

@@ -40,13 +40,17 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
   // Fixed to BigQuery for now.
   override def writeFormat(tableName: String): Format = {
 
+    val tu = TableUtils(sparkSession)
+    val partitionColumnOption =
+      if (tu.tableReachable(tableName)) Map.empty else Map("partitionField" -> tu.partitionColumn)
+
     val sparkOptions: Map[String, String] =
       Map(
         "temporaryGcsBucket" -> sparkSession.conf.get(
           "spark.chronon.table.gcs.temporary_gcs_bucket"
         ), // todo(tchow): No longer needed after https://github.com/GoogleCloudDataproc/spark-bigquery-connector/pull/1320
         "writeMethod" -> "indirect"
-      )
+      ) ++ partitionColumnOption
     BQuery(bqOptions.getProjectId, sparkOptions)
   }
 

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

@@ -10,10 +10,13 @@ case class CatalogAwareDataPointer(inputTableOrPath: String, formatProvider: For
   override def tableOrPath: String = {
     formatProvider.resolveTableName(inputTableOrPath)
   }
-  override lazy val options: Map[String, String] = {
-    // Hack for now, include both read and write options for the datapointer.
-    // todo(tchow): rework this abstraction. https://app.asana.com/0/1208785567265389/1209026103291854/f
-    formatProvider.readFormat(inputTableOrPath).options ++ formatProvider.writeFormat(inputTableOrPath).options
+
+  override lazy val readOptions: Map[String, String] = {
+    formatProvider.readFormat(inputTableOrPath).options
+  }
+
+  override lazy val writeOptions: Map[String, String] = {
+    formatProvider.writeFormat(inputTableOrPath).options
   }
 
   override lazy val readFormat: Option[String] = {
@@ -28,7 +31,7 @@ case class CatalogAwareDataPointer(inputTableOrPath: String, formatProvider: For
 
 object DataPointer {
 
-  def apply(tableOrPath: String, sparkSession: SparkSession): DataPointer = {
+  def from(tableOrPath: String, sparkSession: SparkSession): DataPointer = {
     val clazzName =
       sparkSession.conf.get("spark.chronon.table.format_provider.class", classOf[DefaultFormatProvider].getName)
     val mirror = runtimeMirror(getClass.getClassLoader)

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

@@ -304,28 +304,26 @@ object Extensions {
 
     def save(dataPointer: DataPointer): Unit = {
 
+      val optionDfw = dfw.options(dataPointer.writeOptions)
       dataPointer.writeFormat
         .map((wf) => {
           val normalized = wf.toLowerCase
           normalized match {
             case "bigquery" | "bq" =>
-              dfw
+              optionDfw
                 .format("bigquery")
-                .options(dataPointer.options)
                 .save(dataPointer.tableOrPath)
             case "snowflake" | "sf" =>
-              dfw
+              optionDfw
                 .format("net.snowflake.spark.snowflake")
-                .options(dataPointer.options)
                 .option("dbtable", dataPointer.tableOrPath)
                 .save()
             case "parquet" | "csv" =>
-              dfw
+              optionDfw
                 .format(normalized)
-                .options(dataPointer.options)
                 .save(dataPointer.tableOrPath)
             case "hive" | "delta" | "iceberg" =>
-              dfw
+              optionDfw
                 .format(normalized)
                 .insertInto(dataPointer.tableOrPath)
             case _ =>
@@ -334,7 +332,7 @@ object Extensions {
         })
         .getOrElse(
           // None case is just table against default catalog
-          dfw
+          optionDfw
             .format("hive")
             .insertInto(dataPointer.tableOrPath))
     }
@@ -345,25 +343,24 @@ object Extensions {
     def load(dataPointer: DataPointer): DataFrame = {
       val tableOrPath = dataPointer.tableOrPath
 
+      val optionDfr = dfr.options(dataPointer.readOptions)
+
       dataPointer.readFormat
         .map((fmt) => {
           val normalized = fmt.toLowerCase
           normalized match {
             case "bigquery" | "bq" =>
-              dfr
+              optionDfr
                 .format("bigquery")
-                .options(dataPointer.options)
                 .load(tableOrPath)
             case "snowflake" | "sf" =>
-              dfr
+              optionDfr
                 .format("net.snowflake.spark.snowflake")
-                .options(dataPointer.options)
                 .option("dbtable", tableOrPath)
                 .load()
             case "parquet" | "csv" =>
-              dfr
+              optionDfr
                 .format(normalized)
-                .options(dataPointer.options)
                 .load(tableOrPath)
             case "hive" | "delta" | "iceberg" => dfr.table(tableOrPath)
             case _ =>
@@ -372,7 +369,7 @@ object Extensions {
         })
         .getOrElse {
           // None case is just table against default catalog
-          dfr.table(tableOrPath)
+          optionDfr.table(tableOrPath)
         }
     }
   }

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

@@ -139,7 +139,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
 
   // Needs provider
   def loadTable(tableName: String): DataFrame = {
-    sparkSession.read.load(DataPointer(tableName, sparkSession))
+    sparkSession.read.load(DataPointer.from(tableName, sparkSession))
   }
 
   def isPartitioned(tableName: String): Boolean = {
@@ -243,7 +243,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
   }
 
   def getSchemaFromTable(tableName: String): StructType = {
-    sparkSession.read.load(DataPointer(tableName, sparkSession)).limit(1).schema
+    sparkSession.read.load(DataPointer.from(tableName, sparkSession)).limit(1).schema
   }
 
   // method to check if a user has access to a table
@@ -256,7 +256,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
       // retrieve one row from the table
       val partitionFilter = lastAvailablePartition(tableName).getOrElse(fallbackPartition)
       sparkSession.read
-        .load(DataPointer(tableName, sparkSession))
+        .load(DataPointer.from(tableName, sparkSession))
         .where(s"$partitionColumn='$partitionFilter'")
         .limit(1)
         .collect()
@@ -533,7 +533,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
           (Seq(partitionColumn, saltCol), Seq(partitionColumn) ++ sortByCols)
         } else { (Seq(saltCol), sortByCols) }
       logger.info(s"Sorting within partitions with cols: $partitionSortCols")
-      val dataPointer = DataPointer(tableName, sparkSession)
+      val dataPointer = DataPointer.from(tableName, sparkSession)
 
       val dfw = saltedDf
         .select(saltedDf.columns.map {
@@ -799,7 +799,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
                  wheres: Seq[String],
                  rangeWheres: Seq[String],
                  fallbackSelects: Option[Map[String, String]] = None): DataFrame = {
-    val dp = DataPointer(table, sparkSession)
+    val dp = DataPointer.from(table, sparkSession)
     var df = sparkSession.read.load(dp)
     val selects = QueryUtils.buildSelects(selectMap, fallbackSelects)
     logger.info(s""" Scanning data: