feat: use spark bq connector v1 (#664)

## Summary

- We need to bring back the v1 version of Datasource for spark bigquery
connector, since it supports partition pushdown. And alternative
project_id's. The catalog version in the spark bigquery connector does
not support that.

## 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

## Summary by CodeRabbit

- **New Features**
- Enhanced table reading capabilities with support for applying
partition filters and combining multiple predicates for more flexible
data queries.
- **Refactor**
- Improved internal handling of predicate filters and table loading
logic for more consistent and maintainable data access.
- Refined data filtering by explicitly incorporating partition column
information for more precise queries.
- **Chores**
- Updated script to ensure temporary files are cleaned up more reliably
during installation processes.
<!-- 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

api/python/ai/chronon/resources/gcp/zipline-cli-install.sh

@@ -49,8 +49,8 @@ done
 
 gcloud storage cp "${ARTIFACT_PREFIX%/}/release/$VERSION/wheels/zipline_ai-$VERSION-py3-none-any.whl" .
 
+trap 'rm -f ./zipline_ai-$VERSION-py3-none-any.whl' EXIT
+
 pip3 uninstall zipline-ai
 
 pip3 install ./zipline_ai-$VERSION-py3-none-any.whl
-
-trap 'rm -f ./zipline_ai-$VERSION-py3-none-any.whl' EXIT

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

@@ -4,14 +4,30 @@ import ai.chronon.spark.TableUtils
 import ai.chronon.spark.format.Format
 import com.google.cloud.bigquery.BigQueryOptions
 import com.google.cloud.spark.bigquery.v2.Spark35BigQueryTableProvider
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.functions.{col, date_format, to_date}
 
 case object BigQueryNative extends Format {
 
   private val bqFormat = classOf[Spark35BigQueryTableProvider].getName
   private lazy val bqOptions = BigQueryOptions.getDefaultInstance
 
+  override def table(tableName: String, partitionFilters: String)(implicit sparkSession: SparkSession): DataFrame = {
+    val bqTableId = SparkBQUtils.toTableId(tableName)
+    val bqFriendlyName = scala.Option(bqTableId.getProject) match {
+      case Some(project) => f"${project}.${bqTableId.getDataset}.${bqTableId.getTable}"
+      case None          => f"${bqTableId.getDataset}.${bqTableId.getTable}"
+    }
+    val dfw = sparkSession.read.format(bqFormat)
+    if (partitionFilters.isEmpty) {
+      dfw.load(bqFriendlyName)
+    } else {
+      dfw
+        .option("filter", partitionFilters.trim.stripPrefix("(").stripSuffix(")"))
+        .load(bqFriendlyName)
+    }
+  }
+
   override def primaryPartitions(tableName: String, partitionColumn: String, subPartitionsFilter: Map[String, String])(
       implicit sparkSession: SparkSession): List[String] =
     super.primaryPartitions(tableName, partitionColumn, subPartitionsFilter)

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

@@ -655,7 +655,7 @@ object GroupBy {
          |""".stripMargin)
     metaColumns ++= timeMapping
 
-    val partitionConditions = tableUtils.whereClauses(intersectedRange)
+    val partitionConditions = tableUtils.whereClauses(intersectedRange, source.partitionColumn(tableUtils))
 
     logger.info(s"""
          |Rendering source query:

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

@@ -115,8 +115,12 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     }
   }
 
-  def loadTable(tableName: String): DataFrame = {
-    sparkSession.read.table(tableName)
+  def loadTable(tableName: String, rangeWheres: Seq[String] = List.empty[String]): DataFrame = {
+    tableFormatProvider
+      .readFormat(tableName)
+      .map(_.table(tableName, andPredicates(rangeWheres))(sparkSession))
+      .getOrElse(
+        throw new RuntimeException(s"Could not load table: ${tableName} with partition filter: ${rangeWheres}"))
   }
 
   def createDatabase(database: String): Boolean = {
@@ -564,6 +568,12 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     }
   }
 
+  private def andPredicates(predicates: Seq[String]): String = {
+    val whereStr = predicates.map(p => s"($p)").mkString(" AND ")
+    logger.info(s"""Where str: $whereStr""")
+    whereStr
+  }
+
   def scanDfBase(selectMap: Map[String, String],
                  table: String,
                  wheres: Seq[String],
@@ -582,14 +592,12 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
                    |    ${rangeWheres.mkString(",\n    ").green}
                    |""".stripMargin)
 
-    var df = loadTable(table)
+    var df = loadTable(table, rangeWheres)
 
     if (selects.nonEmpty) df = df.selectExpr(selects: _*)
 
-    val allWheres = wheres ++ rangeWheres
-    if (allWheres.nonEmpty) {
-      val whereStr = allWheres.map(w => s"($w)").mkString(" AND ")
-      logger.info(s"""Where str: $whereStr""")
+    if (wheres.nonEmpty) {
+      val whereStr = andPredicates(wheres)
       df = df.where(whereStr)
     }
 

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

@@ -2,10 +2,21 @@ package ai.chronon.spark.format
 
 import org.apache.spark.sql.SparkSession
 import org.slf4j.{Logger, LoggerFactory}
+import org.apache.spark.sql.DataFrame
 
 trait Format {
+
   @transient protected lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
+  def table(tableName: String, partitionFilters: String)(implicit sparkSession: SparkSession): DataFrame = {
+    val df = sparkSession.read.table(tableName)
+    if (partitionFilters.isEmpty) {
+      df
+    } else {
+      df.where(partitionFilters)
+    }
+  }
+
   // Return the primary partitions (based on the 'partitionColumn') filtered down by sub-partition filters if provided
   // If subpartition filters are supplied and the format doesn't support it, we throw an error
   def primaryPartitions(tableName: String,

spark/src/test/scala/ai/chronon/spark/test/ResultValidationAbilityTest.scala

@@ -69,7 +69,7 @@ class ResultValidationAbilityTest extends AnyFlatSpec with BeforeAndAfter {
     val rdd = args.sparkSession.sparkContext.parallelize(leftData)
     val df = args.sparkSession.createDataFrame(rdd).toDF(columns: _*)
 
-    when(mockTableUtils.loadTable(any())).thenReturn(df)
+    when(mockTableUtils.loadTable(any(), any())).thenReturn(df)
 
     assertTrue(args.validateResult(df, Seq("keyId", "ds"), mockTableUtils))
   }
@@ -85,7 +85,7 @@ class ResultValidationAbilityTest extends AnyFlatSpec with BeforeAndAfter {
     val rightRdd = args.sparkSession.sparkContext.parallelize(rightData)
     val rightDf = args.sparkSession.createDataFrame(rightRdd).toDF(columns: _*)
 
-    when(mockTableUtils.loadTable(any())).thenReturn(rightDf)
+    when(mockTableUtils.loadTable(any(), any())).thenReturn(rightDf)
 
     assertFalse(args.validateResult(leftDf, Seq("keyId", "ds"), mockTableUtils))
   }