Vz/add test case for different partition formats (#753)

Author: nikhil-zlai

.github/workflows/test_scala_2_12_spark.yaml

@@ -71,7 +71,7 @@ jobs:
             //spark:batch_test
 
   fetcher_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:
@@ -97,7 +97,7 @@ jobs:
             //spark:fetcher_test
 
   join_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:
@@ -123,7 +123,7 @@ jobs:
             //spark:join_test
 
   groupby_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:

.github/workflows/test_scala_2_13_spark.yaml

@@ -74,7 +74,7 @@ jobs:
             //spark:batch_test
 
   fetcher_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:
@@ -101,7 +101,7 @@ jobs:
             //spark:fetcher_test
 
   join_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:
@@ -128,7 +128,7 @@ jobs:
             //spark:join_test
 
   groupby_tests:
-    runs-on: ubuntu-8_cores-32_gb
+    runs-on: ubuntu_32_core_128gb
     container:
       image: ghcr.io/${{ github.repository }}-ci:latest
       credentials:

.gitignore

@@ -110,3 +110,5 @@ releases
 MODULE.bazel*
 
 /plugins/vscode/node_modules
+
+**/.claude/settings.local.json

aggregator/src/main/scala/ai/chronon/aggregator/base/MinHeap.scala

@@ -74,7 +74,7 @@ class MinHeap[T](maxSize: Int, ordering: Ordering[T]) {
     arr
   }
 
-  //mutating arr1 / arr2 intact
+  // mutating arr1 / arr2 intact
   def merge(
       arr1: ContainerType[T],
       arr2: ContainerType[T]

aggregator/src/main/scala/ai/chronon/aggregator/base/TimedAggregators.scala

@@ -72,7 +72,7 @@ abstract class TimeOrdered(inputType: DataType) extends TimedAggregator[Any, Tim
 }
 
 class First(inputType: DataType) extends TimeOrdered(inputType) {
-  //mutating
+  // mutating
   override def update(
       ir: util.ArrayList[Any],
       input: Any,
@@ -92,7 +92,7 @@ class First(inputType: DataType) extends TimeOrdered(inputType) {
 }
 
 class Last(inputType: DataType) extends TimeOrdered(inputType) {
-  //mutating
+  // mutating
   override def update(
       ir: util.ArrayList[Any],
       input: Any,

api/python/test/sample/scripts/data-loader.scala

@@ -3,7 +3,6 @@
 import java.io.File
 import org.apache.spark.sql.types._
 
-
 spark.sql("CREATE DATABASE IF NOT EXISTS data;")
 // Directory containing CSV files
 val folderPath = "/srv/chronon/data/"
@@ -27,9 +26,9 @@ files.foreach { file =>
   val customSchema = StructType(
     columns.map { columnName =>
       val dataType = columnName match {
-        case "ts" => LongType
+        case "ts"                                                     => LongType
         case name if name.endsWith("_price") || name.endsWith("_amt") => LongType
-        case _ => StringType
+        case _                                                        => StringType
       }
       StructField(columnName, dataType, nullable = true)
     }

api/src/main/scala/ai/chronon/api/DataRange.scala

@@ -88,7 +88,7 @@ case class PartitionRange(start: String, end: String)(implicit val partitionSpec
   }
 
   def partitions: Seq[String] = {
-    assert(wellDefined, s"Invalid partition range $this")
+    require(wellDefined, s"Invalid partition range $this")
     Stream
       .iterate(start)(partitionSpec.after)
       .takeWhile(_ <= end)
@@ -151,6 +151,15 @@ case class PartitionRange(start: String, end: String)(implicit val partitionSpec
       compareDate(this.end, that.end)
     }
   }
+
+  def translate(otherSpec: PartitionSpec): PartitionRange = {
+
+    val newStart = Option(start).map(d => partitionSpec.translate(d, otherSpec)).orNull
+    val newEnd = Option(end).map(d => partitionSpec.translate(d, otherSpec)).orNull
+
+    PartitionRange(newStart, newEnd)(otherSpec)
+  }
+
   override def toString: String = s"[$start...$end]"
 }
 

api/src/main/scala/ai/chronon/api/PartitionSpec.scala

@@ -80,6 +80,19 @@ case class PartitionSpec(column: String, format: String, spanMillis: Long) {
   def now: String = at(System.currentTimeMillis())
 
   def shiftBaour clientsFromNow(days: Int): String = shift(now, 0 - days)
+
+  def intervalWindow: Window = {
+    if (spanMillis == WindowUtils.Day.millis) WindowUtils.Day
+    else if (spanMillis == WindowUtils.Hour.millis) WindowUtils.Hour
+    else
+      throw new UnsupportedOperationException(
+        s"Partition Intervals should be either hour or day - found ${spanMillis / 60 * 1000} minutes")
+  }
+
+  def translate(date: String, targetSpec: PartitionSpec): String = {
+    val millis = epochMillis(date)
+    targetSpec.at(millis)
+  }
 }
 
 object PartitionSpec {

cloud_aws/src/test/scala/ai/chronon/integrations/aws/HudiTableUtilsTest.scala

@@ -18,12 +18,13 @@ class HudiTableUtilsTest extends AnyFlatSpec {
           "spark.sql.catalog.spark_catalog" -> "org.apache.spark.sql.hudi.catalog.HoodieCatalog",
           "spark.sql.extensions" -> "org.apache.spark.sql.hudi.HoodieSparkSessionExtension",
           "spark.chronon.table_write.format" -> "hudi",
-          "spark.kryo.registrator" -> classOf[ChrononHudiKryoRegistrator].getName,
+          "spark.kryo.registrator" -> classOf[ChrononHudiKryoRegistrator].getName
         )
-    ))
+      )
+    )
   private val tableUtils = TableUtils(spark)
 
-  //todo(tchow): Fix once we support HUDI
+  // todo(tchow): Fix once we support HUDI
   it should "create a hudi table and read the hudi table" ignore {
     import spark.implicits._
     val tableName = "db.test_create_table"

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/DelegatingBigQueryMetastoreCatalog.scala

@@ -123,7 +123,7 @@ class DelegatingBigQueryMetastoreCatalog extends TableCatalog with SupportsNames
               fileBasedTable
             }
             case _: StandardTableDefinition => {
-              //todo(tchow): Support partitioning
+              // todo(tchow): Support partitioning
 
               // Haour clients because there's a bug in the BigQueryCatalog where they ignore the projectId.
               // See: https://github.com/GoogleCloudDataproc/spark-bigquery-connector/pull/1340