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.

## Checklist
- [ ] 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 checkout 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 track
the status of stacks 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 checkouts_external
-WHERE ds BETWEEN '{{ start_date }}' AND '{{ end_date }}'
+    SELECT
+        purchases.ds,
+        purchases.ts as purchase_ts,
+        purchases.user_id,
+        purchases.purchase_price,
+        checkouts.return_id,
+        checkouts.refund_amt,
+        checkouts.product_id,
+        checkouts.ts as checkout_ts
+    FROM data.purchases AS purchases 
+    LEFT OUTER JOIN data.checkouts_external AS checkouts
+    USING (user_id)
+    WHERE purchases.ds BETWEEN '{{ start_date }}' AND '{{ end_date }}'
 """
 
 staging_query = StagingQuery(
@@ -35,7 +39,4 @@
         name='checkouts_staging_query',
         outputNamespace="data"
     ),
-    setups=[
-        "CREATE OR REPLACE TEMPORARY VIEW checkouts_external USING parquet OPTIONS (path 'gs://zl-warehouse/data/checkouts_ds_not_in_parquet/')",
-    ],
 )

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

@@ -1,73 +1,36 @@
 package 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.repackaged.com.google.cloud.bigquery.BigQuery
-import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.StandardTableDefinition
-import com.google.cloud.spark.bigquery.repackaged.com.google.cloud.bigquery.TableInfo
-import com.google.cloud.spark.bigquery.repackaged.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 @@
 package 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.repackaged.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.repackaged.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 bucket.")
-
-      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 @@
 package 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.repackaged.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 backed 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 backed 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)