chore: slim down tableutils (#458)

## Summary

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

- **Refactor**
- Updated the data saving process for enhanced consistency by replacing
the legacy unpartitioned saving functionality with a unified method that
explicitly handles partition columns.
- Removed the functionality to save unpartitioned DataFrames, ensuring
all saves now require partition column specifications.
- **Bug Fixes**
- Removed unnecessary partition checks in tests, streamlining the
validation process without impacting overall functionality.
- **Tests**
- Updated method calls in tests to reflect changes in how table formats
are accessed, ensuring accurate validation of expected outcomes.
<!-- 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

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/BigQueryCatalogTest.scala

@@ -66,8 +66,6 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
     val nativeTable = "data.sample_native"
     val table = tableUtils.loadTable(nativeTable)
     table.show
-    val partitioned = tableUtils.isPartitioned(nativeTable)
-    println(partitioned)
     // val database = tableUtils.createDatabase("test_database")
     val allParts = tableUtils.allPartitions(nativeTable)
     println(allParts)
@@ -80,8 +78,6 @@ class BigQueryCatalogTest extends AnyFlatSpec with MockitoSugar {
     println(bs)
     val table = tableUtils.loadTable(externalTable)
     table.show
-    val partitioned = tableUtils.isPartitioned(externalTable)
-    println(partitioned)
     // val database = tableUtils.createDatabase("test_database")
     val allParts = tableUtils.allPartitions(externalTable)
     println(allParts)

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

@@ -162,10 +162,6 @@ object Extensions {
                                                    sortByCols = sortByCols)
     }
 
-    def saveUnPartitioned(tableName: String, tableProperties: Map[String, String] = null): Unit = {
-      TableUtils(df.sparkSession).insertUnPartitioned(df, tableName, tableProperties)
-    }
-
     def prefixColumnNames(prefix: String, columns: Seq[String]): DataFrame = {
       columns.foldLeft(df) { (renamedDf, key) =>
         renamedDf.withColumnRenamed(key, s"${prefix}_$key")

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

@@ -258,7 +258,7 @@ object GroupByUpload {
     kvDf
       .union(metaDf)
       .withColumn("ds", lit(endDs))
-      .saveUnPartitioned(groupByConf.metaData.uploadTable, groupByConf.metaData.tableProps)
+      .save(groupByConf.metaData.uploadTable, groupByConf.metaData.tableProps, partitionColumns = List.empty)
 
     val kvDfReloaded = tableUtils
       .loadTable(groupByConf.metaData.uploadTable)

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

@@ -51,7 +51,7 @@ class StagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tab
     }
     // the input table is not partitioned, usually for data testing or for kaggle demos
     if (stagingQueryConf.startPartition == null) {
-      tableUtils.sql(stagingQueryConf.query).saveUnPartitioned(outputTable)
+      tableUtils.sql(stagingQueryConf.query).save(outputTable, partitionColumns = List.empty)
     } else {
       val overrideStart = overrideStartPartition.getOrElse(stagingQueryConf.startPartition)
       val unfilledRanges =

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

@@ -30,7 +30,7 @@ import ai.chronon.spark.TableUtils.{
   TableCreationStatus
 }
 import ai.chronon.spark.format.CreationUtils.alterTablePropertiesSql
-import ai.chronon.spark.format.{DefaultFormatProvider, Format, FormatProvider}
+import ai.chronon.spark.format.{DefaultFormatProvider, FormatProvider}
 import org.apache.hadoop.hive.metastore.api.AlreadyExistsException
 import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, Project}
@@ -111,7 +111,6 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
   private val aggregationParallelism: Int = sparkSession.conf.get("spark.chronon.group_by.parallelism", "1000").toInt
 
   sparkSession.sparkContext.setLogLevel("ERROR")
-  // converts String-s like "a=b/c=d" to Map("a" -> "b", "c" -> "d")
 
   def preAggRepartition(df: DataFrame): DataFrame =
     if (df.rdd.getNumPartitions < aggregationParallelism) {
@@ -122,7 +121,7 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
 
   def tableReachable(tableName: String): Boolean = {
     try {
-      tableReadFormat(tableName).isDefined
+      tableFormatProvider.readFormat(tableName).isDefined
     } catch {
       case ex: Exception =>
         logger.info(s"""Couldn't reach $tableName. Error: ${ex.getMessage.red}
@@ -137,12 +136,6 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     sparkSession.read.load(DataPointer.from(tableName, sparkSession))
   }
 
-  def isPartitioned(tableName: String): Boolean = {
-    // TODO: use proper way to detect if a table is partitioned or not
-    val schema = getSchemaFromTable(tableName)
-    schema.fieldNames.contains(partitionColumn)
-  }
-
   // Needs provider
   def createDatabase(database: String): Boolean = {
     try {
@@ -159,17 +152,16 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     }
   }
 
-  def tableReadFormat(tableName: String): Option[Format] = tableFormatProvider.readFormat(tableName)
-
-  // Needs provider
   // return all specified partition columns in a table in format of Map[partitionName, PartitionValue]
   def allPartitions(tableName: String, partitionColumnsFilter: Seq[String] = Seq.empty): Seq[Map[String, String]] = {
 
     if (!tableReachable(tableName)) return Seq.empty[Map[String, String]]
 
-    val format = tableReadFormat(tableName).getOrElse(
-      throw new IllegalStateException(
-        s"Could not determine read format of table ${tableName}. It is no longer reachable."))
+    val format = tableFormatProvider
+      .readFormat(tableName)
+      .getOrElse(
+        throw new IllegalStateException(
+          s"Could not determine read format of table ${tableName}. It is no longer reachable."))
     val partitionSeq = format.partitions(tableName)(sparkSession)
 
     if (partitionColumnsFilter.isEmpty) {
@@ -189,7 +181,8 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
                  subPartitionsFilter: Map[String, String] = Map.empty,
                  partitionColumnName: String = partitionColumn): Seq[String] = {
 
-    tableReadFormat(tableName)
+    tableFormatProvider
+      .readFormat(tableName)
       .map((format) => {
         val partitions = format.primaryPartitions(tableName, partitionColumnName, subPartitionsFilter)(sparkSession)
 
@@ -385,23 +378,6 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
     }
   }
 
-  // Needs provider
-  def insertUnPartitioned(df: DataFrame,
-                          tableName: String,
-                          tableProperties: Map[String, String] = null,
-                          saveMode: SaveMode = SaveMode.Overwrite,
-                          fileFormat: String = "PARQUET"): Unit = {
-
-    val creationStatus = createTable(df, tableName, Seq.empty[String], tableProperties, fileFormat)
-
-    creationStatus match {
-      case TableUtils.TableCreatedWithoutInitialData | TableUtils.TableAlreadyExists =>
-        repartitionAndWrite(df, tableName, saveMode, None, partitionColumns = Seq.empty)
-      case TableUtils.TableCreatedWithInitialData =>
-    }
-
-  }
-
   def columnSizeEstimator(dataType: DataType): Long = {
     dataType match {
       // TODO: improve upon this very basic estimate approach

spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala

@@ -79,13 +79,13 @@ class CompareJob(
     logger.info("Saving comparison output..")
     logger.info(
       s"Comparison schema ${compareDf.schema.fields.map(sb => (sb.name, sb.dataType)).toMap.mkString("\n - ")}")
-    compareDf.saveUnPartitioned(comparisonTableName, tableProps)
+    compareDf.save(comparisonTableName, tableProps, partitionColumns = List.empty)
 
     // Save the metrics table
     logger.info("Saving metrics output..")
     val metricsDf = metricsTimedKvRdd.toFlatDf
     logger.info(s"Metrics schema ${metricsDf.schema.fields.map(sb => (sb.name, sb.dataType)).toMap.mkString("\n - ")}")
-    metricsDf.saveUnPartitioned(metricsTableName, tableProps)
+    metricsDf.save(metricsTableName, tableProps, partitionColumns = List.empty)
 
     logger.info("Printing basic comparison results..")
     logger.info("(Note: This is just an estimation and not a detailed analysis of results)")

spark/src/main/scala/ai/chronon/spark/utils/PartitionRunner.scala

@@ -119,7 +119,7 @@ class PartitionRunner[T](verb: String,
       if (outputDf.columns.contains(tu.partitionColumn)) {
         outputDf.save(outputTable)
       } else {
-        outputDf.saveUnPartitioned(outputTable)
+        outputDf.save(outputTable, partitionColumns = List.empty)
       }
       println(s"""
            |Finished computing range ${i + 1}/$n

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

@@ -147,7 +147,7 @@ class TableUtilsFormatTest extends AnyFlatSpec {
   it should "return empty read format if table doesn't exist" in {
     val dbName = s"db_${System.currentTimeMillis()}"
     val tableName = s"$dbName.test_table_nonexistent_$format"
-    assertTrue(tableUtils.tableReadFormat(tableName).isEmpty)
+    assertTrue(tableUtils.tableFormatProvider.readFormat(tableName).isEmpty)
     assertFalse(tableUtils.tableReachable(tableName))
   }
 }
@@ -188,7 +188,7 @@ object TableUtilsFormatTest {
     tableUtils.insertPartitions(df2, tableName, autoExpand = true)
 
     // check that we wrote out a table in the right format
-    val readTableFormat = tableUtils.tableReadFormat(tableName).get.toString
+    val readTableFormat = tableUtils.tableFormatProvider.readFormat(tableName).get.toString
     assertTrue(s"Mismatch in table format: $readTableFormat; expected: $format", readTableFormat.toLowerCase == format)
 
     // check we have all the partitions written