some additional changes

Author: tchow-zlai

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

@@ -5,9 +5,9 @@ import ai.chronon.spark.FormatProvider
 import ai.chronon.spark.Hive
 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.connector.common.BigQueryUtil
-import com.google.cloud.bigquery.{TableId => BTableId}
 import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.TableId
 import org.apache.spark.sql.SparkSession
 
@@ -19,8 +19,8 @@ case class GCPFormatProvider(sparkSession: SparkSession) extends FormatProvider
 
   override def resolveTableName(tableName: String): String = {
     format(tableName: String) match {
-      case BQuery(_)      => tableName
       case GCS(_, uri, _) => uri
+      case _              => tableName
     }
   }
   override def readFormat(tableName: String): Format = format(tableName)
@@ -39,15 +39,16 @@ case class GCPFormatProvider(sparkSession: SparkSession) extends FormatProvider
     // Active project in the gcloud CLI configuration.
     // No default project: An error will occur if no project ID is available.
 
-    val unshadedTI: BTableId =
-      BTableId.of(bqOptions.getProjectId, btTableIdentifier.getDataset, btTableIdentifier.getTable)
-
-    val tableOpt = Option(bigQueryClient.getTable(unshadedTI))
+    val tableOpt = Option(bigQueryClient.getTable(btTableIdentifier.getDataset, btTableIdentifier.getTable))
     tableOpt match {
       case Some(table) => {
         if (table.getDefinition.isInstanceOf[ExternalTableDefinition]) {
-          import com.google.cloud.bigquery.FormatOptions
-          val uris = table.getDefinition.asInstanceOf[ExternalTableDefinition].getSourceUris.asScala.toList
+          val uris = table.getDefinition
+            .asInstanceOf[ExternalTableDefinition]
+            .getSourceUris
+            .asScala
+            .toList
+            .map((uri) => uri.stripSuffix("/*") + "/")
 
           assert(uris.length == 1, s"External table ${tableName} can be backed by only one URI.")
 
@@ -56,8 +57,9 @@ case class GCPFormatProvider(sparkSession: SparkSession) extends FormatProvider
             .getFormatOptions
             .asInstanceOf[FormatOptions]
             .getType
-          GCS(unshadedTI.getProject, uris.head, formatStr)
-        } else if (table.getDefinition.isInstanceOf[StandardTableDefinition]) BQuery(unshadedTI.getProject)
+
+          GCS(table.getTableId.getProject, uris.head, formatStr)
+        } else if (table.getDefinition.isInstanceOf[StandardTableDefinition]) BQuery(table.getTableId.getProject)
         else throw new IllegalStateException(s"Cannot support table of type: ${table.getDefinition}")
       }
       case None => Hive

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

@@ -4,9 +4,9 @@ import ai.chronon.spark.Format
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.execution.FileSourceScanExec
 import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex
-import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.functions.explode
-import org.apache.spark.sql.functions.url_decode
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.Row
 
 case class GCS(project: String, sourceUri: String, format: String) extends Format {
 
@@ -17,37 +17,6 @@ case class GCS(project: String, sourceUri: String, format: String) extends Forma
     super.primaryPartitions(tableName, partitionColumn, subPartitionsFilter)
 
   override def partitions(tableName: String)(implicit sparkSession: SparkSession): Seq[Map[String, String]] = {
-    import sparkSession.implicits._
-
-    val tableIdentifier = sparkSession.sessionState.sqlParser.parseTableIdentifier(tableName)
-    val table = tableIdentifier.table
-    val database = tableIdentifier.database.getOrElse(throw new IllegalArgumentException("database required!"))
-
-    // See: https://github.com/GoogleCloudDataproc/spark-bigquery-connector/issues/434#issuecomment-886156191
-    // and: https://cloud.google.com/bigquery/docs/information-schema-intro#limitations
-    sparkSession.conf.set("viewsEnabled", "true")
-    sparkSession.conf.set("materializationDataset", database)
-
-    // First, grab the URI location from BQ
-    val uriSQL =
-      s"""
-         |select JSON_EXTRACT_STRING_ARRAY(option_value) as option_values from `${project}.${database}.INFORMATION_SCHEMA.TABLE_OPTIONS`
-         |WHERE table_name = '${table}' and option_name = 'uris'
-         |
-         |""".stripMargin
-
-    val uris = sparkSession.read
-      .format("bigquery")
-      .option("project", project)
-      .option("query", uriSQL)
-      .load()
-      .select(explode(col("option_values")).as("option_value"))
-      .select(url_decode(col("option_value")))
-      .as[String]
-      .collect
-      .toList
-
-    assert(uris.length == 1, s"External table ${tableName} can be backed by only one URI.")
 
     /**
       * Given:
@@ -70,7 +39,7 @@ case class GCS(project: String, sourceUri: String, format: String) extends Forma
       *
       */
     val partitionSpec = sparkSession.read
-      .parquet(uris: _*)
+      .parquet(sourceUri)
       .queryExecution
       .sparkPlan
       .asInstanceOf[FileSourceScanExec]
@@ -82,16 +51,23 @@ case class GCS(project: String, sourceUri: String, format: String) extends Forma
     val partitionColumns = partitionSpec.partitionColumns
     val partitions = partitionSpec.partitions.map(_.values)
 
-    partitions
+    val deserializer =
+      Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
+
+    val roundTripped = sparkSession
+      .createDataFrame(sparkSession.sparkContext.parallelize(partitions.map(deserializer)), partitionColumns)
+      .collect
+      .toList
+
+    roundTripped
       .map((part) =>
         partitionColumns.fields.toList.zipWithIndex.map {
           case (field, idx) => {
             val fieldName = field.name
-            val fieldValue = part.get(idx, field.dataType)
+            val fieldValue = part.get(idx)
             fieldName -> fieldValue.toString // Just going to cast this as a string.
           }
         }.toMap)
-      .toList
   }
 
   def createTableTypeString: String = throw new UnsupportedOperationException("GCS does not support create table")

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

@@ -61,6 +61,7 @@ trait FormatProvider extends Serializable {
   def sparkSession: SparkSession
   def readFormat(tableName: String): Format
   def writeFormat(tableName: String): Format
+
   def resolveTableName(tableName: String) = tableName
 }