feat: StagingQuery param (#406)

## Summary

- Supporting StagingQueries for configurable compute engines. To support
BigQuery, the simplest way is to just write bigquery sql and run it on
bq to create the final table. Let's first make the API change.

## 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**
- Added an option for users to specify the compute engine when
processing queries, offering choices such as Spark and BigQuery.
- Introduced validation to ensure that queries run only with the
designated engine.

- **Style**
  - Streamlined code organization for enhanced readability.
  - Consolidated and reordered import statements for improved clarity.
<!-- 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/src/main/scala/ai/chronon/api/Builders.scala

@@ -308,13 +308,15 @@ object Builders {
         metaData: MetaData = null,
         startPartition: String = null,
         setups: Seq[String] = null,
-        partitionColumn: String = null
+        partitionColumn: String = null,
+        engineType: EngineType = EngineType.SPARK
     ): StagingQuery = {
       val stagingQuery = new StagingQuery()
       stagingQuery.setQuery(query)
       stagingQuery.setMetaData(metaData)
       stagingQuery.setStartPartition(startPartition)
       stagingQuery.setPartitionColumn(partitionColumn)
+      stagingQuery.setEngineType(engineType)
       if (setups != null) stagingQuery.setSetups(setups.toJava)
       stagingQuery
     }

api/thrift/api.thrift

@@ -54,6 +54,11 @@ struct StagingQuery {
     * Only needed for `max_date` template
     **/
     5: optional string partitionColumn
+
+    /**
+    * By default, spark is the compute engine. You can specify an override (eg. bigquery, etc.)
+    **/
+    6: optional EngineType engineType
 }
 
 struct EventSource {
@@ -220,6 +225,12 @@ enum Accuracy {
     SNAPSHOT = 1
 }
 
+enum EngineType {
+    SPARK = 0,
+    BIGQUERY = 1
+
+}
+
 struct MetaData {
     1: optional string name
     // marking this as true means that the conf can be served online

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

@@ -18,12 +18,11 @@ package ai.chronon.spark
 
 import ai.chronon.api
 import ai.chronon.api.Extensions._
-import ai.chronon.api.ParametricMacro
+import ai.chronon.api.{EngineType, ParametricMacro}
 import ai.chronon.api.ScalaJavaConversions._
 import ai.chronon.online.PartitionRange
 import ai.chronon.spark.Extensions._
-import org.slf4j.Logger
-import org.slf4j.LoggerFactory
+import org.slf4j.{Logger, LoggerFactory}
 
 import scala.collection.mutable
 
@@ -46,57 +45,64 @@ class StagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tab
                           enableAutoExpand: Option[Boolean] = Some(true),
                           overrideStartPartition: Option[String] = None,
                           skipFirstHole: Boolean = true): Unit = {
-    Option(stagingQueryConf.setups).foreach(_.toScala.foreach(tableUtils.sql))
+    if (stagingQueryConf.getEngineType != EngineType.SPARK) {
+      throw new UnsupportedOperationException(
+        s"Engine type ${stagingQueryConf.getEngineType} is not supported for Staging Query")
+    }
     // the input table is not partitioned, usually for data testing or for kaggle demos
     if (stagingQueryConf.startPartition == null) {
-      tableUtils.sql(stagingQueryConf.query).save(outputTable)
-      return
-    }
-    val overrideStart = overrideStartPartition.getOrElse(stagingQueryConf.startPartition)
-    val unfilledRanges =
-      tableUtils.unfilledRanges(outputTable,
-                                PartitionRange(overrideStart, endPartition)(tableUtils.partitionSpec),
-                                skipFirstHole = skipFirstHole)
+      tableUtils.sql(stagingQueryConf.query).saveUnPartitioned(outputTable)
+    } else {
+      val overrideStart = overrideStartPartition.getOrElse(stagingQueryConf.startPartition)
+      val unfilledRanges =
+        tableUtils.unfilledRanges(outputTable,
+                                  PartitionRange(overrideStart, endPartition)(tableUtils.partitionSpec),
+                                  skipFirstHole = skipFirstHole)
 
-    if (unfilledRanges.isEmpty) {
-      logger.info(s"""No unfilled range for $outputTable given
-           |start partition of ${stagingQueryConf.startPartition}
-           |override start partition of $overrideStart
-           |end partition of $endPartition
-           |""".stripMargin)
-      return
-    }
-    val stagingQueryUnfilledRanges = unfilledRanges.get
-    logger.info(s"Staging Query unfilled ranges: $stagingQueryUnfilledRanges")
-    val exceptions = mutable.Buffer.empty[String]
-    stagingQueryUnfilledRanges.foreach { stagingQueryUnfilledRange =>
-      try {
-        val stepRanges = stepDays.map(stagingQueryUnfilledRange.steps).getOrElse(Seq(stagingQueryUnfilledRange))
-        logger.info(s"Staging query ranges to compute: ${stepRanges.map { _.toString }.pretty}")
-        stepRanges.zipWithIndex.foreach { case (range, index) =>
-          val progress = s"| [${index + 1}/${stepRanges.size}]"
-          logger.info(s"Computing staging query for range: $range  $progress")
-          val renderedQuery =
-            StagingQuery.substitute(tableUtils, stagingQueryConf.query, range.start, range.end, endPartition)
-          logger.info(s"Rendered Staging Query to run is:\n$renderedQuery")
-          val df = tableUtils.sql(renderedQuery)
-          df.save(outputTable, tableProps, partitionCols, autoExpand = enableAutoExpand.get)
-          logger.info(s"Wrote to table $outputTable, into partitions: $range $progress")
-        }
-        logger.info(s"Finished writing Staging Query data to $outputTable")
-      } catch {
-        case err: Throwable =>
-          exceptions.append(s"Error handling range $stagingQueryUnfilledRange : ${err.getMessage}\n${err.traceString}")
+      if (unfilledRanges.isEmpty) {
+        logger.info(s"""No unfilled range for $outputTable given
+             |start partition of ${stagingQueryConf.startPartition}
+             |override start partition of $overrideStart
+             |end partition of $endPartition
+             |""".stripMargin)
+        return
       }
-    }
-    if (exceptions.nonEmpty) {
-      val length = exceptions.length
-      val fullMessage = exceptions.zipWithIndex
-        .map { case (message, index) =>
-          s"[${index + 1}/${length} exceptions]\n${message}"
+      val stagingQueryUnfilledRanges = unfilledRanges.get
+      logger.info(s"Staging Query unfilled ranges: $stagingQueryUnfilledRanges")
+      Option(stagingQueryConf.setups).foreach(_.toScala.foreach(tableUtils.sql))
+      val exceptions = mutable.Buffer.empty[String]
+      stagingQueryUnfilledRanges.foreach { stagingQueryUnfilledRange =>
+        try {
+          val stepRanges = stepDays.map(stagingQueryUnfilledRange.steps).getOrElse(Seq(stagingQueryUnfilledRange))
+          logger.info(s"Staging query ranges to compute: ${stepRanges.map {
+            _.toString
+          }.pretty}")
+          stepRanges.zipWithIndex.foreach { case (range, index) =>
+            val progress = s"| [${index + 1}/${stepRanges.size}]"
+            logger.info(s"Computing staging query for range: $range  $progress")
+            val renderedQuery =
+              StagingQuery.substitute(tableUtils, stagingQueryConf.query, range.start, range.end, endPartition)
+            logger.info(s"Rendered Staging Query to run is:\n$renderedQuery")
+            val df = tableUtils.sql(renderedQuery)
+            df.save(outputTable, tableProps, partitionCols, autoExpand = enableAutoExpand.get)
+            logger.info(s"Wrote to table $outputTable, into partitions: $range $progress")
+          }
+          logger.info(s"Finished writing Staging Query data to $outputTable")
+        } catch {
+          case err: Throwable =>
+            exceptions.append(
+              s"Error handling range $stagingQueryUnfilledRange : ${err.getMessage}\n${err.traceString}")
         }
-        .mkString("\n")
-      throw new Exception(fullMessage)
+      }
+      if (exceptions.nonEmpty) {
+        val length = exceptions.length
+        val fullMessage = exceptions.zipWithIndex
+          .map { case (message, index) =>
+            s"[${index + 1}/${length} exceptions]\n${message}"
+          }
+          .mkString("\n")
+        throw new Exception(fullMessage)
+      }
     }
   }
 }

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

@@ -19,8 +19,8 @@ package ai.chronon.spark
 import ai.chronon.aggregator.windowing.TsUtils
 import ai.chronon.api.ColorPrinter.ColorString
 import ai.chronon.api.Extensions._
-import ai.chronon.api.{Constants, PartitionSpec, Query, QueryUtils}
 import ai.chronon.api.ScalaJavaConversions._
+import ai.chronon.api.{Constants, PartitionSpec, Query, QueryUtils}
 import ai.chronon.online.PartitionRange
 import ai.chronon.spark.Extensions._
 import ai.chronon.spark.TableUtils.{
@@ -32,17 +32,17 @@ import ai.chronon.spark.TableUtils.{
 import ai.chronon.spark.format.CreationUtils.alterTablePropertiesSql
 import ai.chronon.spark.format.{DefaultFormatProvider, Format, FormatProvider}
 import org.apache.hadoop.hive.metastore.api.AlreadyExistsException
-import org.apache.spark.sql.{AnalysisException, DataFrame, SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, Project}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.{AnalysisException, DataFrame, SaveMode, SparkSession}
 import org.apache.spark.storage.StorageLevel
 import org.slf4j.{Logger, LoggerFactory}
 
 import java.io.{PrintWriter, StringWriter}
-import java.time.{Instant, ZoneId}
 import java.time.format.DateTimeFormatter
+import java.time.{Instant, ZoneId}
 import scala.collection.{Seq, immutable, mutable}
 import scala.util.{Failure, Try}