chore: refactor some existing methods

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)