pr feedback

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

Author: tchow-zlai

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

@@ -11,6 +11,7 @@ import com.google.cloud.bigquery.Table
 import com.google.cloud.bigquery.connector.common.BigQueryUtil
 import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.TableId
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.functions.{col, to_date}
 
 import scala.collection.JavaConverters._
 
@@ -87,7 +88,6 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
       * case None             => throw new IllegalStateException("Dataproc federation service must be available.")
       *
       * }
-      * }
       *
       * case Some("hive") | None => Hive
       * }
@@ -151,6 +151,13 @@ case class BQuery(project: String) extends Format {
         .option("project", project)
         .option("query", partValsSql)
         .load()
+        .select(
+          to_date(col("partition_id"),
+                  "yyyyMMdd"
+          ) // Note: this "yyyyMMdd" format is hardcoded but we need to change it to be something else.
+            .as("partition_id"))
+        .na // Should filter out '__NULL__' and '__UNPARTITIONED__'. See: https://cloud.google.com/bigquery/docs/partitioned-tables#date_timestamp_partitioned_tables
+        .drop()
         .as[String]
         .collect
         .toList

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

@@ -7,10 +7,9 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.FileSourceScanExec
 import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex
+case class GCS(project: String, sourceUri: String, fileFormat: String) extends Format {
 
-case class GCS(project: String, sourceUri: String, format: String) extends Format {
-
-  override def name: String = format
+  override def name: String = fileFormat
 
   override def primaryPartitions(tableName: String, partitionColumn: String, subPartitionsFilter: Map[String, String])(
       implicit sparkSession: SparkSession): Seq[String] =
@@ -39,7 +38,8 @@ case class GCS(project: String, sourceUri: String, format: String) extends Forma
       *
       */
     val partitionSpec = sparkSession.read
-      .parquet(sourceUri)
+      .format(fileFormat)
+      .load(sourceUri)
       .queryExecution
       .sparkPlan
       .asInstanceOf[FileSourceScanExec]
@@ -52,7 +52,12 @@ case class GCS(project: String, sourceUri: String, format: String) extends Forma
     val partitions = partitionSpec.partitions.map(_.values)
 
     val deserializer =
-      Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
+      try {
+        Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
+      } catch {
+        case e: Exception =>
+          throw new RuntimeException(s"Failed to create deserializer for partition columns: ${e.getMessage}", e)
+      }
 
     val roundTripped = sparkSession
       .createDataFrame(sparkSession.sparkContext.parallelize(partitions.map(deserializer)), partitionColumns)

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

@@ -278,7 +278,7 @@ object Driver {
       val join = new Join(
         args.joinConf,
         args.endDate(),
-        args.buildTableUtils(),
+        tableUtils,
         !args.runFirstHole(),
         selectedJoinParts = args.selectedJoinParts.toOption
       )

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

@@ -309,6 +309,7 @@ object Extensions {
               dfw
                 .format("bigquery")
                 .options(dataPointer.options)
+                .option("writeMethod", "direct")
                 .save(dataPointer.tableOrPath)
             case "snowflake" | "sf" =>
               dfw

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

@@ -233,7 +233,11 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     try {
       // retrieve one row from the table
       val partitionFilter = lastAvailablePartition(tableName).getOrElse(fallbackPartition)
-      sparkSession.sql(s"SELECT * FROM $tableName where $partitionColumn='$partitionFilter' LIMIT 1").collect()
+      sparkSession.read
+        .load(DataPointer(tableName, sparkSession))
+        .where(s"$partitionColumn='$partitionFilter'")
+        .limit(1)
+        .collect()
       true
     } catch {
       case e: SparkException =>