direct write and rebase

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

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

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

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

Author: tchow-zlai

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

@@ -1,17 +1,16 @@
 package ai.chronon.integrations.cloud_gcp
-
-import ai.chronon.spark.TableUtils
 import ai.chronon.spark.format.Format
 import ai.chronon.spark.format.FormatProvider
 import ai.chronon.spark.format.Hive
-import com.google.cloud.bigquery.BigQuery
-import com.google.cloud.bigquery.BigQueryOptions
-import com.google.cloud.bigquery.ExternalTableDefinition
-import com.google.cloud.bigquery.FormatOptions
-import com.google.cloud.bigquery.StandardTableDefinition
-import com.google.cloud.bigquery.Table
-import com.google.cloud.bigquery.TableDefinition
 import com.google.cloud.bigquery.connector.common.BigQueryUtil
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.BigQuery
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.BigQueryOptions
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.ExternalTableDefinition
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.FormatOptions
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.JobInfo
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.StandardTableDefinition
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.Table
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.TableDefinition
 import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.TableId
 import org.apache.spark.sql.SparkSession
 
@@ -43,17 +42,10 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
     assert(scala.Option(tableId.getProject).isDefined, s"project required for ${table}")
     assert(scala.Option(tableId.getDataset).isDefined, s"dataset required for ${table}")
 
-    val tu = TableUtils(sparkSession)
-    val partitionColumnOption =
-      if (tu.tableReachable(table)) Map.empty else Map("partitionField" -> tu.partitionColumn)
-
     val sparkOptions: Map[String, String] = Map(
-      // todo(tchow): No longer needed after https://github.com/GoogleCloudDataproc/spark-bigquery-connector/pull/1320
-      "temporaryGcsBucket" -> sparkSession.conf.get("spark.chronon.table.gcs.temporary_gcs_bucket"),
-      "writeMethod" -> "indirect", // writeMethod direct does not output partitioned tables. keep as indirect.
-      "materializationProject" -> tableId.getProject,
-      "materializationDataset" -> tableId.getDataset
-    ) ++ partitionColumnOption
+      "writeMethod" -> "direct",
+      "createDisposition" -> JobInfo.CreateDisposition.CREATE_NEVER.name
+    )
 
     BigQueryFormat(tableId.getProject, bigQueryClient, sparkOptions)
   }

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/GcpFormatProviderTest.scala

@@ -1,7 +1,7 @@
 package ai.chronon.integrations.cloud_gcp
 
 import ai.chronon.spark.SparkSessionBuilder
-import com.google.cloud.bigquery._
+import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery._
 import org.apache.spark.sql.SparkSession
 import org.mockito.Mockito.when
 import org.scalatest.flatspec.AnyFlatSpec

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

@@ -284,12 +284,12 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
                   tableProperties: Map[String, String] = null,
                   fileFormat: String,
                   autoExpand: Boolean = false): Unit = {
+    val writeFormat = tableFormatProvider.writeFormat(tableName)
 
     if (!tableReachable(tableName)) {
 
       try {
 
-        val writeFormat = tableFormatProvider.writeFormat(tableName)
         val createTableOperation =
           writeFormat.createTable(df, tableName, partitionColumns, tableProperties, fileFormat)
 
@@ -309,11 +309,13 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
 
     // TODO: we need to also allow for bigquery tables to have their table properties (or tags) to be persisted too.
     //  https://app.asana.com/0/1208949807589885/1209111629687568/f
-    if (tableProperties != null && tableProperties.nonEmpty) {
-      sql(alterTablePropertiesSql(tableName, tableProperties))
-    }
-    if (autoExpand) {
-      expandTable(tableName, df.schema)
+    if (writeFormat.name.toUpperCase != "BIGQUERY") {
+      if (tableProperties != null && tableProperties.nonEmpty) {
+        sql(alterTablePropertiesSql(tableName, tableProperties))
+      }
+      if (autoExpand) {
+        expandTable(tableName, df.schema)
+      }
     }
   }
 
@@ -328,12 +330,12 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
                        stats: Option[DfStats] = None,
                        sortByCols: Seq[String] = Seq.empty): Unit = {
     // partitions to the last
-    val dfRearranged: DataFrame = if (!df.columns.endsWith(partitionColumns)) {
-      val colOrder = df.columns.diff(partitionColumns) ++ partitionColumns
-      df.select(colOrder.map(df.col): _*)
-    } else {
-      df
-    }
+    val colOrder = df.columns.diff(partitionColumns) ++ partitionColumns
+    val dfRearranged: DataFrame = df.select(colOrder.map {
+      case c if c == partitionColumn =>
+        to_date(df.col(c), partitionFormat).as(partitionColumn)
+      case c => df.col(c)
+    }: _*)
 
     createTable(dfRearranged, tableName, partitionColumns, tableProperties, fileFormat, autoExpand)
 
@@ -526,11 +528,6 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
       val dataPointer = DataPointer.from(tableName, sparkSession)
 
       saltedDf
-        .select(saltedDf.columns.map {
-          case c if c == partitionColumn && dataPointer.writeFormat.map(_.toUpperCase).exists("BIGQUERY".equals) =>
-            to_date(saltedDf.col(c), partitionFormat).as(partitionColumn)
-          case c => saltedDf.col(c)
-        }.toList: _*)
         .repartition(shuffleParallelism, repartitionCols.map(saltedDf.col): _*)
         .drop(saltCol)
         .sortWithinPartitions(partitionSortCols.map(col): _*)