clean up

airbnb · Nov 10, 2023 · 49f80a0 · 49f80a0
1 parent 14f5adc
commit 49f80a0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/spark/src/main/scala/ai/chronon/spark/JoinBase.scala b/spark/src/main/scala/ai/chronon/spark/JoinBase.scala
@@ -301,15 +301,7 @@ abstract class JoinBase(joinConf: api.Join,
     }
 
     // First run command to archive tables that have changed semantically since the last run
-    val archivedAtTs = Instant.now()
-    val tablesAfterVersionCheck = tablesToRecompute(joinConf, outputTable, tableUtils, forceOverwriteMetadata)
-    tablesAfterVersionCheck.foreach { table =>
-      if (forceOverwriteMetadata) {
-        tableUtils.sql(tableUtils.unsetTablePropertiesSql(table, Constants.SemanticHashKey))
-      } else {
-        tableUtils.archiveOrDropTableIfExists(table, Some(archivedAtTs))
-      }
-    }
+    tablesToRecompute(joinConf, outputTable, tableUtils, forceOverwriteMetadata)
 
     // detect holes and chunks to fill
     // OverrideStartPartition is used to replace the start partition of the join config. This is useful when

diff --git a/spark/src/main/scala/ai/chronon/spark/JoinUtils.scala b/spark/src/main/scala/ai/chronon/spark/JoinUtils.scala
@@ -10,6 +10,7 @@ import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.functions.{coalesce, col, udf}
 import org.apache.spark.util.sketch.BloomFilter
 
+import java.time.Instant
 import scala.collection.Seq
 import scala.util.ScalaJavaConversions.MapOps
 
@@ -321,17 +322,27 @@ object JoinUtils {
                         outputTable: String,
                         tableUtils: TableUtils,
                         forceOverwriteMetadata: Boolean = false): collection.Seq[String] = {
-    if (forceOverwriteMetadata) {
-      return collection.Seq.empty
-    }
     val gson = new Gson()
-    (for (
+    val tablesAfterVersionCheck = (for (
       props <- tableUtils.getTableProperties(outputTable);
       oldSemanticJson <- props.get(Constants.SemanticHashKey);
       oldSemanticHash = gson.fromJson(oldSemanticJson, classOf[java.util.HashMap[String, String]]).toScala
     ) yield {
       println(s"Comparing Hashes:\nNew: ${joinConf.semanticHash},\nOld: $oldSemanticHash")
       joinConf.tablesToDrop(oldSemanticHash)
     }).getOrElse(collection.Seq.empty)
+
+    if (forceOverwriteMetadata) {
+      tablesAfterVersionCheck.foreach { table =>
+        tableUtils.sql(tableUtils.unsetTablePropertiesSql(table, Constants.SemanticHashKey))
+      }
+      collection.Seq.empty
+    } else {
+      val archivedAtTs = Instant.now()
+      tablesAfterVersionCheck.foreach { table =>
+        tableUtils.archiveOrDropTableIfExists(table, Some(archivedAtTs))
+      }
+      tablesAfterVersionCheck
+    }
   }
 }