join part output alignment

Author: nikhil-zlai

spark/src/main/scala/ai/chronon/spark/batch/JoinPartJob.scala

@@ -25,7 +25,9 @@ case class JoinPartJobContext(leftDf: Option[DfWithStats],
                               tableProps: Map[String, String],
                               runSmallMode: Boolean)
 
-class JoinPartJob(node: JoinPartNode, metaData: MetaData, range: DateRange, showDf: Boolean = false)(implicit
+// alignOutput forces the job to produce the partitions specified by range.
+// legacy behavior was to not align, but that prevents from partition aware orchestration
+class JoinPartJob(node: JoinPartNode, metaData: MetaData, range: DateRange, showDf: Boolean = false, alignOutput: Boolean = false)(implicit
     tableUtils: TableUtils) {
   @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
   implicit val partitionSpec: PartitionSpec = tableUtils.partitionSpec
@@ -195,7 +197,7 @@ class JoinPartJob(node: JoinPartNode, metaData: MetaData, range: DateRange, show
       skewFilteredLeft.select(columns: _*)
     }
 
-    lazy val shiftedPartitionRange = unfilledPartitionRange.shift(-1)
+    lazy val shiftedPartitionRange = if(alignOutput) unfilledPartitionRange else unfilledPartitionRange.shift(-1)
 
     val renamedLeftDf = renamedLeftRawDf.select(renamedLeftRawDf.columns.map {
       case c if c == tableUtils.partitionColumn =>

spark/src/main/scala/ai/chronon/spark/batch/ModularMonolith.scala

@@ -1,9 +1,9 @@
 package ai.chronon.spark.batch
 
 import ai.chronon.api
-import ai.chronon.api.Extensions.{JoinOps, MetadataOps, SourceOps}
+import ai.chronon.api.Extensions.{GroupByOps, JoinOps, MetadataOps, SourceOps}
 import ai.chronon.api.ScalaJavaConversions.IterableOps
-import ai.chronon.api.{DateRange, MetaData, PartitionRange}
+import ai.chronon.api.{Accuracy, DataModel, DateRange, MetaData, PartitionRange, PartitionSpec}
 import ai.chronon.planner.{JoinBootstrapNode, JoinDerivationNode, JoinMergeNode, JoinPartNode, SourceWithFilterNode}
 import ai.chronon.spark.JoinUtils
 import ai.chronon.spark.catalog.TableUtils
@@ -124,7 +124,15 @@ class ModularMonolith(join: api.Join, dateRange: DateRange)(implicit tableUtils:
         .setSkewKeys(join.skewKeys)
 
       StepRunner(dateRange, partMetaData) { stepRange =>
-        val joinPartJob = new JoinPartJob(joinPartNode, partMetaData, stepRange)
+        val shiftedRange = if (join.left.dataModel == DataModel.EVENTS && joinPart.groupBy.inferredAccuracy == Accuracy.SNAPSHOT) {
+          val spec = PartitionSpec.daily
+          new DateRange()
+            .setStartDate(spec.before(stepRange.startDate))
+            .setEndDate(spec.before(stepRange.endDate))
+        } else {
+          stepRange
+        }
+        val joinPartJob = new JoinPartJob(joinPartNode, partMetaData, shiftedRange, alignOutput = true)
         joinPartJob.run(None) // Run without context for now
       }
       logger.info(s"JoinPartJob completed for: $joinPartGroupByName, output table: $partFullTableName")