feat: Support iceberg reads and writes (#424)

## Summary

- Modifying TableUtils to use spark's built-in catalog wherever
applicable. This will both improve performance and allow us to access
iceberg tables through the bigquery catalog, since its built on
#393.

## Cheour clientslist
- [ ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced external cheour clientsout queries now deliver additional purchase
details for richer data insights.

- **Refactor**
- Updated cloud integrations to enforce supported table formats and
provide clearer error messaging.
- Streamlined table creation and partitioning processes in Spark for
improved performance and maintainability.
- Standardized collection types for partitioning and sorting parameters
in Dataframe operations.

- **Tests**
- Refined test cases to validate dynamic provider handling and verify
consistent table operations.
- Removed outdated test case for field name retrieval, reflecting
changes in validation focus.
  - Updated assertions to utilize the new format provider access method.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

<!-- av pr metadata
This information is embedded by the av CLI when creating PRs to traour clients
the status of staour clientss when using Aviator. Please do not delete or edit
this section of the PR.
```
{"parent":"main","parentHead":"","trunk":"main"}
```
-->

---------

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

Author: tchow-zlai

.bazelrc

@@ -16,4 +16,4 @@ build --java_language_version=11
 build --java_runtime_version=11
 build --remote_cache=https://storage.googleapis.com/zipline-bazel-cache
 test --test_output=errors
-test --test_timeout=900
+test --test_timeout=1200

api/py/test/sample/staging_queries/quickstart/checkouts_external.py

@@ -17,15 +17,19 @@
 
 
 query = """
-SELECT
-    ts,
-    ds,
-    return_id,
-    user_id,
-    product_id,
-    refund_amt
-FROM cheour clientsouts_external
-WHERE ds BETWEEN '{{ start_date }}' AND '{{ end_date }}'
+    SELECT
+        purchases.ds,
+        purchases.ts as purchase_ts,
+        purchases.user_id,
+        purchases.purchase_price,
+        cheour clientsouts.return_id,
+        cheour clientsouts.refund_amt,
+        cheour clientsouts.product_id,
+        cheour clientsouts.ts as cheour clientsout_ts
+    FROM data.purchases AS purchases 
+    LEFT OUTER JOIN data.cheour clientsouts_external AS cheour clientsouts
+    USING (user_id)
+    WHERE purchases.ds BETWEEN '{{ start_date }}' AND '{{ end_date }}'
 """
 
 staging_query = StagingQuery(
@@ -35,7 +39,4 @@
         name='cheour clientsouts_staging_query',
         outputNamespace="data"
     ),
-    setups=[
-        "CREATE OR REPLACE TEMPORARY VIEW cheour clientsouts_external USING parquet OPTIONS (path 'gs://zl-warehouse/data/cheour clientsouts_ds_not_in_parquet/')",
-    ],
 )

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

@@ -1,73 +1,36 @@
 paour clientsage ai.chronon.integrations.cloud_gcp
 
 import ai.chronon.spark.TableUtils
-import ai.chronon.spark.TableUtils.TableCreationStatus
 import ai.chronon.spark.format.Format
+import com.google.cloud.bigquery.BigQuery
 import com.google.cloud.bigquery.connector.common.BigQueryUtil
-import com.google.cloud.spark.bigquery.SchemaConverters
-import com.google.cloud.spark.bigquery.SchemaConvertersConfiguration
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery.BigQuery
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery.StandardTableDefinition
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery.TableInfo
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery.TimePartitioning
-import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.functions.date_format
-import org.apache.spark.sql.functions.to_date
+import com.google.cloud.spark.bigquery.v2.Spark35BigQueryTableProvider
+import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.sql.functions.{col, date_format, to_date}
 
 case class BigQueryFormat(project: String, bqClient: BigQuery, override val options: Map[String, String])
     extends Format {
   override def name: String = "bigquery"
 
+  private val bqFormat = classOf[Spark35BigQueryTableProvider].getName
+
   override def alterTableProperties(tableName: String,
                                     tableProperties: Map[String, String]): (String => Unit) => Unit = {
     throw new NotImplementedError("alterTableProperties not yet supported for BigQuery")
   }
 
   override def primaryPartitions(tableName: String, partitionColumn: String, subPartitionsFilter: Map[String, String])(
-      implicit sparkSession: SparkSession): Seq[String] =
+      implicit sparkSession: SparkSession): List[String] =
     super.primaryPartitions(tableName, partitionColumn, subPartitionsFilter)
   override def generateTableBuilder(df: DataFrame,
                                     tableName: String,
-                                    partitionColumns: Seq[String],
+                                    partitionColumns: List[String],
                                     tableProperties: Map[String, String],
-                                    fileFormat: String): (String => Unit) => TableCreationStatus = {
-
-    def inner(df: DataFrame, tableName: String, partitionColumns: Seq[String])(
-        sqlEvaluator: String => Unit): TableCreationStatus = {
-
-      // See: https://cloud.google.com/bigquery/docs/partitioned-tables#limitations
-      // "BigQuery does not support partitioning by multiple columns. Only one column can be used to partition a table."
-      assert(partitionColumns.size < 2,
-             s"BigQuery only supports at most one partition column, incoming spec: ${partitionColumns}")
-      val shadedTableId = BigQueryUtil.parseTableId(tableName)
-
-      val shadedBqSchema =
-        SchemaConverters.from(SchemaConvertersConfiguration.createDefault()).toBigQuerySchema(df.schema)
-
-      val baseTableDef = StandardTableDefinition.newBuilder
-        .setSchema(shadedBqSchema)
-
-      val tableDefinition = partitionColumns.headOption
-        .map((col) => {
-          val timePartitioning = TimePartitioning.newBuilder(TimePartitioning.Type.DAY).setField(col)
-          baseTableDef
-            .setTimePartitioning(timePartitioning.build())
-        })
-        .getOrElse(baseTableDef)
-
-      val tableInfoBuilder = TableInfo.newBuilder(shadedTableId, tableDefinition.build)
-
-      val tableInfo = tableInfoBuilder.build
-      bqClient.create(tableInfo)
-      TableUtils.TableCreatedWithoutInitialData
-    }
-
-    inner(df, tableName, partitionColumns)
+                                    fileFormat: String): (String => Unit) => Unit = {
+    throw new UnsupportedOperationException("generateTableBuilder not supported for BigQuery")
   }
 
-  override def partitions(tableName: String)(implicit sparkSession: SparkSession): Seq[Map[String, String]] = {
+  override def partitions(tableName: String)(implicit sparkSession: SparkSession): List[Map[String, String]] = {
     import sparkSession.implicits._
     val tableIdentifier = BigQueryUtil.parseTableId(tableName)
     val table = tableIdentifier.getTable
@@ -83,7 +46,7 @@ case class BigQueryFormat(project: String, bqClient: BigQuery, override val opti
            |""".stripMargin
 
     val partitionCol = sparkSession.read
-      .format("bigquery")
+      .format(bqFormat)
       .option("project", project)
       // See: https://github.com/GoogleCloudDataproc/spark-bigquery-connector/issues/434#issuecomment-886156191
       // and: https://cloud.google.com/bigquery/docs/information-schema-intro#limitations
@@ -109,7 +72,7 @@ case class BigQueryFormat(project: String, bqClient: BigQuery, override val opti
     val partitionFormat = TableUtils(sparkSession).partitionFormat
 
     val partitionInfoDf = sparkSession.read
-      .format("bigquery")
+      .format(bqFormat)
       .option("project", project)
       // See: https://github.com/GoogleCloudDataproc/spark-bigquery-connector/issues/434#issuecomment-886156191
       // and: https://cloud.google.com/bigquery/docs/information-schema-intro#limitations

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

@@ -167,6 +167,10 @@ class DelegatingBigQueryMetastoreCatalog extends CatalogExtension {
                            schema: StructType,
                            partitions: Array[Transform],
                            properties: util.Map[String, String]): Table = {
+    val provider = properties.get(TableCatalog.PROP_PROVIDER)
+    if (provider.toUpperCase != "ICEBERG") {
+      throw new UnsupportedOperationException("Only creating iceberg tables supported.")
+    }
     icebergCatalog.createTable(ident, schema, partitions, properties)
   }
 

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

@@ -1,24 +1,11 @@
 paour clientsage ai.chronon.integrations.cloud_gcp
 
-import ai.chronon.api.Extensions.StringOps
-import ai.chronon.api.ScalaJavaConversions.JListOps
-import ai.chronon.spark.TableUtils
-import ai.chronon.spark.TableUtils.{TableCreatedWithInitialData, TableCreationStatus}
 import ai.chronon.spark.format.Format
-import com.google.cloud.bigquery.connector.common.BigQueryUtil
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery.{
-  BigQuery,
-  BigQueryOptions,
-  ExternalTableDefinition,
-  FormatOptions,
-  HivePartitioningOptions,
-  TableInfo
-}
-import com.google.cloud.spark.bigquery.{SchemaConverters, SchemaConvertersConfiguration}
-import org.apache.spark.sql.{DataFrame, Encoders, Row, SparkSession}
+import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery._
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.FileSourceScanExec
 import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex
+import org.apache.spark.sql.{DataFrame, Encoders, Row, SparkSession}
 import org.slf4j.LoggerFactory
 
 case class GCS(sourceUri: String, fileFormat: String) extends Format {
@@ -31,10 +18,10 @@ case class GCS(sourceUri: String, fileFormat: String) extends Format {
   override def name: String = fileFormat
 
   override def primaryPartitions(tableName: String, partitionColumn: String, subPartitionsFilter: Map[String, String])(
-      implicit sparkSession: SparkSession): Seq[String] =
+      implicit sparkSession: SparkSession): List[String] =
     super.primaryPartitions(tableName, partitionColumn, subPartitionsFilter)
 
-  override def partitions(tableName: String)(implicit sparkSession: SparkSession): Seq[Map[String, String]] = {
+  override def partitions(tableName: String)(implicit sparkSession: SparkSession): List[Map[String, String]] = {
 
     /** Given:
       *  hdfs://<host>:<port>/ path/ to/ partition/ a=1/ b=hello/ c=3.14
@@ -88,62 +75,15 @@ case class GCS(sourceUri: String, fileFormat: String) extends Format {
           fieldName -> fieldValue.toString // Just going to cast this as a string.
 
         }.toMap)
+      .toList
   }
 
   override def generateTableBuilder(df: DataFrame,
                                     tableName: String,
-                                    partitionColumns: Seq[String],
+                                    partitionColumns: List[String],
                                     tableProperties: Map[String, String],
-                                    fileFormat: String): (String => Unit) => TableCreationStatus = {
-
-    def inner(df: DataFrame, tableName: String, partitionColumns: Seq[String])(sqlEvaluator: String => Unit) = {
-
-      // See: https://cloud.google.com/bigquery/docs/partitioned-tables#limitations
-      // "BigQuery does not support partitioning by multiple columns. Only one column can be used to partition a table."
-      require(partitionColumns.size < 2,
-              s"BigQuery only supports at most one partition column, incoming spec: ${partitionColumns}")
-
-      val shadedTableId = BigQueryUtil.parseTableId(tableName)
-
-      val writePrefix = TableUtils(df.sparkSession).writePrefix
-      require(writePrefix.nonEmpty, "Please set conf 'spark.chronon.table_write.prefix' pointing to a data buour clientset.")
-
-      val path = writePrefix.get + tableName.sanitize + "/" //split("/").map(_.sanitize).mkString("/")
-      val dataGlob = path + "*"
-
-      logger.info(s"""
-           |table source uri: $dataGlob
-           |partition uri: $path
-           |""".stripMargin)
-
-      df.write
-        .partitionBy(partitionColumns: _*)
-        .mode("overwrite") // or "append" based on your needs
-        .parquet(path)
-
-      val baseTableDef = ExternalTableDefinition
-        .newBuilder(dataGlob, FormatOptions.parquet())
-        .setAutodetect(true)
-
-      if (partitionColumns.nonEmpty) {
-        val timePartitioning = HivePartitioningOptions
-          .newBuilder()
-          .setFields(partitionColumns.toJava)
-          .setSourceUriPrefix(path)
-          .setMode("STRINGS")
-          .build()
-        baseTableDef.setHivePartitioningOptions(timePartitioning)
-      }
-
-      val tableInfo = TableInfo.newBuilder(shadedTableId, baseTableDef.build).build()
-      val createdTable = bigQueryClient.create(tableInfo)
-
-      println(s"Created external table ${createdTable.getTableId}")
-
-      TableCreatedWithInitialData
-    }
-
-    inner(df, tableName, partitionColumns)
+                                    fileFormat: String): (String => Unit) => Unit = {
+    throw new UnsupportedOperationException("generateTableBuilder not supported for GCS")
   }
 
   def createTableTypeString: String = throw new UnsupportedOperationException("GCS does not support create table")

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

@@ -1,11 +1,15 @@
 paour clientsage ai.chronon.integrations.cloud_gcp
 import ai.chronon.api.Extensions.StringOps
 import ai.chronon.spark.TableUtils
-import ai.chronon.spark.format.{Format, FormatProvider}
+import ai.chronon.spark.format.{Format, FormatProvider, Iceberg}
+import com.google.cloud.bigquery._
 import com.google.cloud.bigquery.connector.common.BigQueryUtil
-import com.google.cloud.spark.bigquery.repaour clientsaged.com.google.cloud.bigquery._
+import com.google.cloud.iceberg.bigquery.relocated.com.google.api.services.bigquery.model.TableReference
+import org.apache.iceberg.exceptions.NoSuchIcebergTableException
+import org.apache.iceberg.gcp.bigquery.{BigQueryClient, BigQueryClientImpl}
 import org.apache.spark.sql.SparkSession
 
+import scala.util.Try
 import scala.jdk.CollectionConverters._
 
 case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider {
@@ -18,7 +22,8 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
     * - No default project: An error will occur if no project ID is available.
     */
   private lazy val bqOptions = BigQueryOptions.getDefaultInstance
-  lazy val bigQueryClient: BigQuery = bqOptions.getService
+  private lazy val bigQueryClient: BigQuery = bqOptions.getService
+  private lazy val icebergClient: BigQueryClient = new BigQueryClientImpl()
 
   override def resolveTableName(tableName: String): String =
     format(tableName)
@@ -39,33 +44,49 @@ case class GcpFormatProvider(sparkSession: SparkSession) extends FormatProvider
     GCS(path, "PARQUET")
   }
 
-  private[cloud_gcp] def getFormat(table: Table): Format =
+  private[cloud_gcp] def getFormat(table: Table): Format = {
     table.getDefinition.asInstanceOf[TableDefinition] match {
-
       case definition: ExternalTableDefinition =>
-        val formatOptions = definition.getFormatOptions
-          .asInstanceOf[FormatOptions]
-        val externalTable = table.getDefinition.asInstanceOf[ExternalTableDefinition]
-        val uri = scala
-          .Option(externalTable.getHivePartitioningOptions)
-          .map(_.getSourceUriPrefix)
-          .getOrElse {
-            val uris = externalTable.getSourceUris.asScala
-            require(uris.size == 1, s"External table ${table} can be baour clientsed by only one URI.")
-            uris.head.replaceAll("/\\*\\.parquet$", "")
-          }
+        Try {
+          val tableRef = new TableReference()
+            .setProjectId(table.getTableId.getProject)
+            .setDatasetId(table.getTableId.getDataset)
+            .setTableId(table.getTableId.getTable)
+
+          icebergClient.getTable(tableRef) // Just try to load it. It'll fail if it's not an iceberg table.
+          Iceberg
+        }.recover {
+          case canHandle: NoSuchIcebergTableException =>
+            val formatOptions = definition.getFormatOptions.asInstanceOf[FormatOptions]
+            val externalTable = table.getDefinition.asInstanceOf[ExternalTableDefinition]
+
+            val uri = scala
+              .Option(externalTable.getHivePartitioningOptions)
+              .map(_.getSourceUriPrefix)
+              .getOrElse {
+                val uris = externalTable.getSourceUris.asScala
+                require(uris.size == 1, s"External table ${table} can be baour clientsed by only one URI.")
+                uris.head.replaceAll("/\\*\\.parquet$", "")
+              }
 
-        GCS(uri, formatOptions.getType)
+            GCS(uri, formatOptions.getType)
+          case e: Exception => throw e
+        }.get
 
       case _: StandardTableDefinition =>
         BigQueryFormat(table.getTableId.getProject, bigQueryClient, Map.empty)
 
-      case _ => throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")
+      case _ =>
+        throw new IllegalStateException(s"Cannot support table of type: ${table.getFriendlyName}")
     }
+  }
 
   private def format(tableName: String): scala.Option[Format] = {
-
-    val btTableIdentifier: TableId = BigQueryUtil.parseTableId(tableName)
+    val shadedTid = BigQueryUtil.parseTableId(tableName)
+    val btTableIdentifier: TableId = scala
+      .Option(shadedTid.getProject)
+      .map(TableId.of(_, shadedTid.getDataset, shadedTid.getTable))
+      .getOrElse(TableId.of(shadedTid.getDataset, shadedTid.getTable))
     val table = scala.Option(bigQueryClient.getTable(btTableIdentifier.getDataset, btTableIdentifier.getTable))
     table
       .map(getFormat)