wip

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

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

Author: tchow-zlai

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

@@ -0,0 +1,11 @@
+package ai.chronon.integrations.cloud_gcp
+
+import ai.chronon.api
+import ai.chronon.spark.{StagingQuery, TableUtils}
+
+class BigQueryStagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tableUtils: TableUtils) extends StagingQuery(stagingQueryConf, endPartition, tableUtils) {
+  override def computeStagingQuery(stepDays: Option[Int] = None, enableAutoExpand: Option[Boolean] = Some(true), overrideStartPartition: Option[String] = None, skipFirstHole: Boolean = true): Unit = {
+    super.computeStagingQuery(stepDays, enableAutoExpand, overrideStartPartition, skipFirstHole)
+  }
+
+}

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

@@ -41,6 +41,15 @@ class StagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tab
         .asInstanceOf[java.util.ArrayList[String]]
         .toScala
 
+  def computeRenderedQuery(renderedQuery: String, enableAutoExpand: Option[Boolean]) = {
+    val df = tableUtils.sql(renderedQuery)
+    df.save(outputTable, tableProps, partitionCols, autoExpand = enableAutoExpand.get)
+  }
+
+  def prepare() = {
+    Option(stagingQueryConf.setups).foreach(_.toScala.foreach(tableUtils.sql))
+  }
+
   def computeStagingQuery(stepDays: Option[Int] = None,
                           enableAutoExpand: Option[Boolean] = Some(true),
                           overrideStartPartition: Option[String] = None,

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

@@ -23,20 +23,12 @@ 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.{
-  TableAlreadyExists,
-  TableCreatedWithInitialData,
-  TableCreatedWithoutInitialData,
-  TableCreationStatus
-}
+import ai.chronon.spark.TableUtils.{TableAlreadyExists, TableCreatedWithInitialData, TableCreatedWithoutInitialData, TableCreationStatus}
 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.catalyst.analysis.TableAlreadyExistsException
-import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
-import org.apache.spark.sql.catalyst.plans.logical.Filter
-import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-import org.apache.spark.sql.catalyst.plans.logical.Project
+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}