Add a check to only upload each partition once.

Author: chewys1024

spark/src/main/scala/ai/chronon/spark/stats/drift/SummaryUploader.scala

@@ -2,21 +2,45 @@ package ai.chronon.spark.stats.drift
 
 import ai.chronon.online.KVStore
 import ai.chronon.online.KVStore.PutRequest
+import ai.chronon.spark.TableUtils
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.types
+
+import java.io.FileNotFoundException
+
+class SummaryUploader(summaryDF: DataFrame, kvStore: KVStore)(implicit tu: TableUtils) {
+  val completed_schema: types.StructType = types.StructType(
+    Seq(
+      types.StructField(tu.partitionColumn, types.StringType, nullable = false)
+    )
+  )
 
-class SummaryUploader(summaryDF: DataFrame, kvStore: KVStore) {
   def run(): Unit = {
     summaryDF.rdd.foreachPartition(rows => {
       var putRequests: List[PutRequest] = Nil
+      var newPartitions: List[String] = Nil
       for (row <- rows) {
+        val partition = row.getAs[String](tu.partitionColumn)
 
-        putRequests = putRequests :+ PutRequest(
-          if (!row.isNullAt(row.fieldIndex("keyBytes"))) row.getAs[Array[Byte]]("keyBytes") else Array.empty[Byte],
-          if (!row.isNullAt(row.fieldIndex("valueBytes"))) row.getAs[Array[Byte]]("valueBytes") else Array.empty[Byte],
-          "drift_statistics"
-        )
+        try {
+          tu.sparkSession.read.parquet(s"${partition}_completed")
+        } catch {
+          case _: FileNotFoundException => {
+            putRequests = putRequests :+ PutRequest(
+              if (!row.isNullAt(row.fieldIndex("keyBytes"))) row.getAs[Array[Byte]]("keyBytes") else Array.empty[Byte],
+              if (!row.isNullAt(row.fieldIndex("valueBytes"))) row.getAs[Array[Byte]]("valueBytes")
+              else Array.empty[Byte],
+              "drift_statistics"
+            )
+            newPartitions = newPartitions :+ partition
+          }
+        }
       }
       kvStore.multiPut(putRequests)
+      for (partition <- newPartitions) {
+        val df = tu.sparkSession.emptyDataFrame
+        df.write.parquet(s"${partition}_completed")
+      }
     })
 
   }