chore: refactor some existing methods (#156)

## Summary


Don't rely on `TableUtils` directly as much as possible, use the
dataframe extensions.

## Checklist
- [ ] Added Unit Tests
- [x] Covered by existing CI
- [x] Integration tested
- [ ] Documentation update
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Refactor**
	- Updated DataFrame saving methods across multiple Spark job classes
- Replaced `tableUtils.insertPartitions` and
`tableUtils.insertUnPartitioned` with direct DataFrame `save` and
`saveUnPartitioned` methods
- Simplified data persistence mechanism in LogFlattenerJob,
StagingQuery, CompareJob, and ConsistencyJob
<!-- 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

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

@@ -22,6 +22,7 @@ import ai.chronon.api._
 import ai.chronon.online.OnlineDerivationUtil.timeFields
 import ai.chronon.online._
 import ai.chronon.spark.Extensions.StructTypeOps
+import ai.chronon.spark.Extensions._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.Dataset
@@ -224,11 +225,10 @@ class LogFlattenerJob(session: SparkSession,
       val schemaTblProps = buildTableProperties(schemaStringsMap)
       logger.info("======= Log table schema =======")
       logger.info(flattenedDf.schema.pretty)
-      tableUtils.insertPartitions(flattenedDf,
-                                  joinConf.metaData.loggedTable,
-                                  tableProperties =
-                                    joinTblProps ++ schemaTblProps ++ Map(Constants.ChrononLogTable -> true.toString),
-                                  autoExpand = true)
+
+      flattenedDf.save(joinConf.metaData.loggedTable,
+                       joinTblProps ++ schemaTblProps ++ Map(Constants.ChrononLogTable -> true.toString),
+                       autoExpand = true)
 
       val inputRowCount = rawDf.count()
       // read from output table to avoid recomputation

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

@@ -80,7 +80,7 @@ class StagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tab
               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)
-            tableUtils.insertPartitions(df, outputTable, tableProps, partitionCols, autoExpand = enableAutoExpand.get)
+            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")

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

@@ -25,12 +25,12 @@ import ai.chronon.online.DataMetrics
 import ai.chronon.online.PartitionRange
 import ai.chronon.online.SparkConversions
 import ai.chronon.spark.Analyzer
+import ai.chronon.spark.Extensions._
 import ai.chronon.spark.StagingQuery
 import ai.chronon.spark.TableUtils
 import ai.chronon.spark.TimedKvRdd
 import ai.chronon.spark.stats.CompareJob.getJoinKeys
 import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.SaveMode
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 
@@ -82,13 +82,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 - ")}")
-    tableUtils.insertUnPartitioned(compareDf, comparisonTableName, tableProps, saveMode = SaveMode.Overwrite)
+    compareDf.saveUnPartitioned(comparisonTableName, tableProps)
 
     // 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 - ")}")
-    tableUtils.insertUnPartitioned(metricsDf, metricsTableName, tableProps, saveMode = SaveMode.Overwrite)
+    metricsDf.saveUnPartitioned(metricsTableName, tableProps)
 
     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/stats/ConsistencyJob.scala

@@ -131,10 +131,7 @@ class ConsistencyJob(session: SparkSession, joinConf: Join, endDate: String) ext
       logger.info("Saving output.")
       val outputDf = metricsKvRdd.toFlatDf.withTimeBasedColumn("ds")
       logger.info(s"output schema ${outputDf.schema.fields.map(sb => (sb.name, sb.dataType)).toMap.mkString("\n - ")}")
-      tableUtils.insertPartitions(outputDf,
-                                  joinConf.metaData.consistencyTable,
-                                  tableProperties = tblProperties,
-                                  autoExpand = true)
+      outputDf.save(joinConf.metaData.consistencyTable, tableProperties = tblProperties, autoExpand = true)
       metricsKvRdd.toAvroDf
         .withTimeBasedColumn(tableUtils.partitionColumn)
         .save(joinConf.metaData.consistencyUploadTable, tblProperties)