chore: small cleanup of TableUtils

Author: tchow-zlai

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

@@ -669,68 +669,6 @@ case class TableUtils(sparkSession: SparkSession) {
     }
   }
 
-  @deprecated
-  def dropPartitionsAfterHole(inputTable: String,
-                              outputTable: String,
-                              partitionRange: PartitionRange,
-                              subPartitionFilters: Map[String, String] = Map.empty): Option[String] = {
-
-    def partitionsInRange(table: String, partitionFilter: Map[String, String] = Map.empty): Set[String] = {
-      val allParts = partitions(table, partitionFilter)
-      val startPrunedParts = Option(partitionRange.start).map(start => allParts.filter(_ >= start)).getOrElse(allParts)
-      Option(partitionRange.end).map(end => startPrunedParts.filter(_ <= end)).getOrElse(startPrunedParts).toSet
-    }
-
-    val inputPartitions = partitionsInRange(inputTable)
-    val outputPartitions = partitionsInRange(outputTable, subPartitionFilters)
-    val earliestHoleOpt = (inputPartitions -- outputPartitions).reduceLeftOption(Ordering[String].min)
-    earliestHoleOpt.foreach { hole =>
-      val toDrop = outputPartitions.filter(_ > hole)
-      logger.info(s"""
-                 |Earliest hole at $hole in output table $outputTable, relative to $inputTable
-                 |Input Parts   : ${inputPartitions.toArray.sorted.mkString("Array(", ", ", ")")}
-                 |Output Parts  : ${outputPartitions.toArray.sorted.mkString("Array(", ", ", ")")}
-                 |Dropping Parts: ${toDrop.toArray.sorted.mkString("Array(", ", ", ")")}
-                 |Sub Partitions: ${subPartitionFilters.map(kv => s"${kv._1}=${kv._2}").mkString("Array(", ", ", ")")}
-          """.stripMargin)
-      dropPartitions(outputTable, toDrop.toArray.sorted, partitionColumn, subPartitionFilters)
-    }
-    earliestHoleOpt
-  }
-
-  def dropPartitions(tableName: String,
-                     partitions: Seq[String],
-                     partitionColumn: String = partitionColumn,
-                     subPartitionFilters: Map[String, String] = Map.empty): Unit = {
-    if (partitions.nonEmpty && tableExists(tableName)) {
-      val partitionSpecs = partitions
-        .map { partition =>
-          val mainSpec = s"$partitionColumn='$partition'"
-          val specs = mainSpec +: subPartitionFilters.map {
-            case (key, value) => s"$key='$value'"
-          }.toSeq
-          specs.mkString("PARTITION (", ",", ")")
-        }
-        .mkString(",")
-      val dropSql = s"ALTER TABLE $tableName DROP IF EXISTS $partitionSpecs"
-      sql(dropSql)
-    } else {
-      logger.info(s"$tableName doesn't exist, please double check before drop partitions")
-    }
-  }
-
-  def dropPartitionRange(tableName: String,
-                         startDate: String,
-                         endDate: String,
-                         subPartitionFilters: Map[String, String] = Map.empty): Unit = {
-    if (tableExists(tableName)) {
-      val toDrop = Stream.iterate(startDate)(partitionSpec.after).takeWhile(_ <= endDate)
-      dropPartitions(tableName, toDrop, partitionColumn, subPartitionFilters)
-    } else {
-      logger.info(s"$tableName doesn't exist, please double check before drop partitions")
-    }
-  }
-
   /*
    * This method detects new columns that appear in newSchema but not in current table,
    * and append those new columns at the end of the existing table. This allows continuous evolution

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

@@ -57,7 +57,7 @@ object TestRow {
 class JoinTest extends AnyFunSuite with TaggedFilterSuite {
 
   val spark: SparkSession = SparkSessionBuilder.build("JoinTest", local = true)
-  private implicit val tableUtils = TableUtils(spark)
+  private implicit val tableUtils = TableTestUtils(spark)
 
   private val today = tableUtils.partitionSpec.at(System.currentTimeMillis())
   private val monthAgo = tableUtils.partitionSpec.minus(today, new Window(30, TimeUnit.DAYS))

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

@@ -37,7 +37,7 @@ class LabelJoinTest {
   private val namespace = "label_join"
   private val tableName = "test_label_join"
   private val labelDS = "2022-10-30"
-  private val tableUtils = TableUtils(spark)
+  private val tableUtils = TableTestUtils(spark)
   tableUtils.createDatabase(namespace)
 
   private val viewsGroupBy = TestUtils.createViewsGroupBy(namespace, spark)

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

@@ -0,0 +1,41 @@
+package ai.chronon.spark.test
+
+import ai.chronon.spark.TableUtils
+import org.apache.spark.sql.SparkSession
+
+case class TableTestUtils(override val sparkSession: SparkSession) extends TableUtils(sparkSession: SparkSession) {
+
+  def dropPartitions(tableName: String,
+                     partitions: Seq[String],
+                     partitionColumn: String = partitionColumn,
+                     subPartitionFilters: Map[String, String] = Map.empty): Unit = {
+    if (partitions.nonEmpty && tableExists(tableName)) {
+      val partitionSpecs = partitions
+        .map { partition =>
+          val mainSpec = s"$partitionColumn='$partition'"
+          val specs = mainSpec +: subPartitionFilters.map {
+            case (key, value) => s"$key='$value'"
+          }.toSeq
+          specs.mkString("PARTITION (", ",", ")")
+        }
+        .mkString(",")
+      val dropSql = s"ALTER TABLE $tableName DROP IF EXISTS $partitionSpecs"
+      sql(dropSql)
+    } else {
+      logger.info(s"$tableName doesn't exist, please double check before drop partitions")
+    }
+  }
+
+  def dropPartitionRange(tableName: String,
+                         startDate: String,
+                         endDate: String,
+                         subPartitionFilters: Map[String, String] = Map.empty): Unit = {
+    if (tableExists(tableName)) {
+      val toDrop = Stream.iterate(startDate)(partitionSpec.after).takeWhile(_ <= endDate)
+      dropPartitions(tableName, toDrop, partitionColumn, subPartitionFilters)
+    } else {
+      logger.info(s"$tableName doesn't exist, please double check before drop partitions")
+    }
+  }
+
+}

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

@@ -38,8 +38,6 @@ import org.junit.Test
 
 import scala.util.Try
 
-
-
 class SimpleAddUDF extends UDF {
   def evaluate(value: Int): Int = {
     value + 20
@@ -48,7 +46,7 @@ class SimpleAddUDF extends UDF {
 
 class TableUtilsTest {
   lazy val spark: SparkSession = SparkSessionBuilder.build("TableUtilsTest", local = true)
-  private val tableUtils = TableUtils(spark)
+  private val tableUtils = TableTestUtils(spark)
   private implicit val partitionSpec: PartitionSpec = tableUtils.partitionSpec
 
   @Test
@@ -83,10 +81,12 @@ class TableUtilsTest {
       Seq(
         types.StructField("name", types.StringType, nullable = true),
         types.StructField("age", types.IntegerType, nullable = false),
-        types.StructField("address", types.StructType(Seq(
-          types.StructField("street", types.StringType, nullable = true),
-          types.StructField("city", types.StringType, nullable = true)
-        )))
+        types.StructField("address",
+                          types.StructType(
+                            Seq(
+                              types.StructField("street", types.StringType, nullable = true),
+                              types.StructField("city", types.StringType, nullable = true)
+                            )))
       )
     )
     val expectedFieldNames = Seq("name", "age", "address", "address.street", "address.city")
@@ -344,8 +344,7 @@ class TableUtilsTest {
     // verify the latest label version
     val labels = JoinUtils.getLatestLabelMapping(tableName, tableUtils)
     assertEquals(labels("2022-11-09"),
-                 List(PartitionRange("2022-10-01", "2022-10-02"),
-                      PartitionRange("2022-10-05", "2022-10-05")))
+                 List(PartitionRange("2022-10-01", "2022-10-02"), PartitionRange("2022-10-05", "2022-10-05")))
   }
 
   private def prepareTestDataWithSubPartitions(tableName: String): Unit = {