Chaining fix

api/py/ai/chronon/group_by.py

@@ -431,6 +431,7 @@ def _normalize_source(source):
         elif isinstance(source, ttypes.EntitySource):
             return ttypes.Source(entities=source)
         elif isinstance(source, ttypes.JoinSource):
+            source.join.metadata = output_namespace or source.metaData.outputNamespace
             return ttypes.Source(joinSource=source)
         elif isinstance(source, ttypes.Source):
             return source

api/py/ai/chronon/join.py

@@ -377,6 +377,7 @@ def Join(left: api.Source,
          consistency_sample_percent: float = 5.0,
          online_external_parts: List[api.ExternalPart] = None,
          offline_schedule: str = '@daily',
+         historical_backfill: bool = True,
          row_ids: List[str] = None,
          bootstrap_parts: List[api.BootstrapPart] = None,
          bootstrap_from_log: bool = False,
@@ -468,6 +469,10 @@ def Join(left: api.Source,
     :param tags:
         Additional metadata about the Join that you wish to track. Does not effect computation.
     :type tags: Dict[str, str]
+    :param historical_backfill:
+        Flag to indicate whether join backfill should backfill previous holes.
+        Setting to false will only backfill latest single partition
+    :type historical_backfill: bool
     :return:
         A join object that can be used to backfill or serve data. For ML use-cases this should map 1:1 to model.
     """
@@ -558,6 +563,7 @@ def Join(left: api.Source,
 
     bootstrap_dependencies = [] if dependencies is not None else utils.get_bootstrap_dependencies(bootstrap_parts)
 
+    print("$$ historical_backfill: {}".format(historical_backfill))
     metadata = api.MetaData(
         online=online,
         production=production,
@@ -568,7 +574,8 @@ def Join(left: api.Source,
         modeToEnvMap=env,
         samplePercent=sample_percent,
         offlineSchedule=offline_schedule,
-        consistencySamplePercent=consistency_sample_percent
+        consistencySamplePercent=consistency_sample_percent,
+        historicalBackfill=historical_backfill
     )
 
     return api.Join(

api/py/test/sample/group_bys/sample_team/sample_chaining_group_by.py

@@ -59,5 +59,5 @@
         "sample_config_json": """{"sample_key": "sample_value"}""",
         "description": "sample description"
     },
-    output_namespace="sample_namespace",
+    output_namespace="test_namespace",
 )

api/py/test/sample/joins/sample_team/sample_chaining_join.py

@@ -34,7 +34,8 @@
         ),
     ],
     online=True,
-    check_consistency=True
+    check_consistency=True,
+    historical_backfill=False,
 )
 
 chaining_group_by_v1 = GroupBy(

api/py/test/sample/production/group_bys/sample_team/sample_chaining_group_by

@@ -36,7 +36,8 @@
             "team": "sample_team",
             "samplePercent": 100.0,
             "offlineSchedule": "@daily",
-            "consistencySamplePercent": 5.0
+            "consistencySamplePercent": 5.0,
+            "historicalBackfill": 0
           },
           "left": {
             "events": {

api/py/test/sample/production/joins/sample_team/sample_backfill_mutation_join.v0

@@ -15,7 +15,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_chaining_join.parent_join

@@ -16,7 +16,8 @@
     "team": "sample_team",
     "samplePercent": 100.0,
     "offlineSchedule": "@daily",
-    "consistencySamplePercent": 5.0
+    "consistencySamplePercent": 5.0,
+    "historicalBackfill": 0
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_chaining_join.v1

@@ -15,7 +15,8 @@
     "team": "sample_team",
     "samplePercent": 100.0,
     "offlineSchedule": "@daily",
-    "consistencySamplePercent": 5.0
+    "consistencySamplePercent": 5.0,
+    "historicalBackfill": 1
   },
   "left": {
     "events": {
@@ -72,7 +73,8 @@
                   "team": "sample_team",
                   "samplePercent": 100.0,
                   "offlineSchedule": "@daily",
-                  "consistencySamplePercent": 5.0
+                  "consistencySamplePercent": 5.0,
+                  "historicalBackfill": 0
                 },
                 "left": {
                   "events": {

api/py/test/sample/production/joins/sample_team/sample_join.consistency_check

@@ -14,7 +14,8 @@
     "team": "sample_team",
     "samplePercent": 100.0,
     "offlineSchedule": "@daily",
-    "consistencySamplePercent": 5.0
+    "consistencySamplePercent": 5.0,
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join.never

@@ -13,7 +13,8 @@
     "outputNamespace": "sample_namespace",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@never"
+    "offlineSchedule": "@never",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join.no_log_flattener

@@ -13,7 +13,8 @@
     "outputNamespace": "sample_namespace",
     "team": "sample_team",
     "samplePercent": 0.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join.v1

@@ -18,7 +18,8 @@
       }
     },
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join_bootstrap.v1

@@ -17,7 +17,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_join_bootstrap.v2

@@ -19,7 +19,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_join_derivation.v1

@@ -15,7 +15,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_join_external_parts.v1

@@ -13,7 +13,8 @@
     "outputNamespace": "sample_namespace",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join_from_module.v1

@@ -17,7 +17,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_join_from_shorthand.v1

@@ -14,7 +14,8 @@
     "outputNamespace": "chronon_db",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "entities": {

api/py/test/sample/production/joins/sample_team/sample_label_join.v1

@@ -15,7 +15,8 @@
     "outputNamespace": "sample_namespace",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_label_join_with_agg.v1

@@ -15,7 +15,8 @@
     "outputNamespace": "sample_namespace",
     "team": "sample_team",
     "samplePercent": 100.0,
-    "offlineSchedule": "@daily"
+    "offlineSchedule": "@daily",
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

api/py/test/sample/production/joins/sample_team/sample_online_join.v1

@@ -18,7 +18,8 @@
     "team": "sample_team",
     "samplePercent": 100.0,
     "offlineSchedule": "@daily",
-    "consistencySamplePercent": 5.0
+    "consistencySamplePercent": 5.0,
+    "historicalBackfill": 1
   },
   "left": {
     "events": {

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

@@ -250,7 +250,8 @@ object Builders {
         team: String = null,
         samplePercent: Double = 100,
         consistencySamplePercent: Double = 5,
-        tableProperties: Map[String, String] = Map.empty
+        tableProperties: Map[String, String] = Map.empty,
+        historicalBackill: Boolean = true
     ): MetaData = {
       val result = new MetaData()
       result.setName(name)
@@ -259,6 +260,7 @@ object Builders {
       result.setCustomJson(customJson)
       result.setOutputNamespace(namespace)
       result.setTeam(Option(team).getOrElse("chronon"))
+      result.setHistoricalBackfill(historicalBackill)
       if (dependencies != null)
         result.setDependencies(dependencies.toSeq.toJava)
       if (samplePercent > 0)

api/thrift/api.thrift

@@ -253,6 +253,9 @@ struct MetaData {
     12: optional string offlineSchedule
     // percentage of online serving requests used to compute consistency metrics
     13: optional double consistencySamplePercent
+    // Flag to indicate whether join backfill should backfill previous holes.
+    // Setting to false will only backfill latest single partition
+    14: optional bool historicalBackfill
 }
 
 // Equivalent to a FeatureSet in chronon terms

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

@@ -247,6 +247,7 @@ class Analyzer(tableUtils: TableUtils,
     val dataAvailabilityErrors: ListBuffer[(String, String, String)] = ListBuffer.empty[(String, String, String)]
 
     val rangeToFill = JoinUtils.getRangesToFill(joinConf.left, tableUtils, endDate)
+//                                                historicalBackfill = joinConf.metaData.historicalBackfill)
     println(s"[Analyzer] Join range to fill $rangeToFill")
     val unfilledRanges = tableUtils
       .unfilledRanges(joinConf.metaData.outputTable, rangeToFill, Some(Seq(joinConf.left.table)))

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

@@ -209,6 +209,10 @@ object Driver {
         opt[Boolean](required = false,
                      default = Some(false),
                      descr = "Skip the first unfilled partition range if some future partitions have been populated.")
+      val startPartitionOverride: ScallopOption[String] =
+        opt[String](required = false,
+                    descr =
+                      "Start date to compute join backfill, this start date will override start partition in conf.")
       lazy val joinConf: api.Join = parseConf[api.Join](confPath())
       override def subcommandName() = s"join_${joinConf.metaData.name}"
     }
@@ -221,7 +225,7 @@ object Driver {
         args.buildTableUtils(),
         !args.runFirstHole()
       )
-      val df = join.computeJoin(args.stepDays.toOption)
+      val df = join.computeJoin(args.stepDays.toOption, args.startPartitionOverride.toOption)
 
       if (args.shouldExport()) {
         args.exportTableToLocal(args.joinConf.metaData.outputTable, tableUtils)

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

@@ -4,7 +4,7 @@ import ai.chronon.aggregator.base.TimeTuple
 import ai.chronon.aggregator.row.RowAggregator
 import ai.chronon.aggregator.windowing._
 import ai.chronon.api
-import ai.chronon.api.Constants
+import ai.chronon.api.{Accuracy, Constants, DataModel}
 import ai.chronon.api.DataModel.{Entities, Events}
 import ai.chronon.api.Extensions._
 import ai.chronon.online.{RowWrapper, SparkConversions}
@@ -388,8 +388,13 @@ object GroupBy {
         println("Join source detected. Materializing the join.")
         val joinSource = source.getJoinSource
         val joinConf = joinSource.join
-        // materialize the table
-        val join = new Join(joinConf, queryRange.end, tableUtils, mutationScan = false, showDf = showDf)
+        // materialize the table with the right end date. QueryRange.end could be shifted for temporal events
+        val beforeDs = tableUtils.partitionSpec.before(queryRange.end)
+        val isPreShifted =
+          groupByConf.dataModel == DataModel.Events && groupByConf.inferredAccuracy == Accuracy.TEMPORAL
+        val endDate = if (isPreShifted) beforeDs else queryRange.end
+
+        val join = new Join(joinConf, endDate, tableUtils, mutationScan = false, showDf = showDf)
         if (computeDependency) {
           val df = join.computeJoin()
           if (showDf) {

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

@@ -166,6 +166,8 @@ object GroupByUpload {
                                     showDf = showDf)
     lazy val groupByUpload = new GroupByUpload(endDs, groupBy)
     // for temporal accuracy - we don't need to scan mutations for upload
+    // when endDs = xxxx-01-02 the timestamp from airflow is more than (xxxx-01-03 00:00:00)
+    // we wait for event partitions of (xxxx-01-02) which contain data until (xxxx-01-02 23:59:59.999)
     lazy val shiftedGroupBy =
       GroupBy.from(groupByConf,
                    PartitionRange(endDs, endDs).shift(1),

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

@@ -270,7 +270,7 @@ abstract class JoinBase(joinConf: api.Join,
 
   def computeRange(leftDf: DataFrame, leftRange: PartitionRange, bootstrapInfo: BootstrapInfo): DataFrame
 
-  def computeJoin(stepDays: Option[Int] = None): DataFrame = {
+  def computeJoin(stepDays: Option[Int] = None, overrideStartPartition: Option[String] = None): DataFrame = {
 
     assert(Option(joinConf.metaData.team).nonEmpty,
            s"join.metaData.team needs to be set for join ${joinConf.metaData.name}")
@@ -303,7 +303,14 @@ abstract class JoinBase(joinConf: api.Join,
       tableUtils.archiveOrDropTableIfExists(_, Some(archivedAtTs)))
 
     // detect holes and chunks to fill
-    val rangeToFill = JoinUtils.getRangesToFill(joinConf.left, tableUtils, endPartition)
+    // OverrideStartPartition is used to replace the start partition of the join config. This is useful when
+    //  1 - User would like to test run with different start partition
+    //  2 - User has entity table which is accumulative and only want to run backfill for the latest partition
+    val rangeToFill = JoinUtils.getRangesToFill(joinConf.left,
+                                                tableUtils,
+                                                endPartition,
+                                                overrideStartPartition,
+                                                joinConf.metaData.historicalBackfill)
     println(s"Join range to fill $rangeToFill")
     val unfilledRanges = tableUtils
       .unfilledRanges(outputTable, rangeToFill, Some(Seq(joinConf.left.table)), skipFirstHole = skipFirstHole)

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

@@ -79,9 +79,18 @@ object JoinUtils {
     */
   def getRangesToFill(leftSource: ai.chronon.api.Source,
                       tableUtils: TableUtils,
-                      endPartition: String): PartitionRange = {
-    val leftStart = Option(leftSource.query.startPartition)
+                      endPartition: String,
+                      overrideStartPartition: Option[String] = None,
+                      historicalBackfill: Boolean = true): PartitionRange = {
+    val overrideStart = if (historicalBackfill) {
+      overrideStartPartition
+    } else {
+      println(s"Historical backfill is set to false. Backfill single partition only: $endPartition")
+      Some(endPartition)
+    }
+    lazy val defaultLeftStart = Option(leftSource.query.startPartition)
       .getOrElse(tableUtils.firstAvailablePartition(leftSource.table, leftSource.subPartitionFilters).get)
+    val leftStart = overrideStart.getOrElse(defaultLeftStart)
     val leftEnd = Option(leftSource.query.endPartition).getOrElse(endPartition)
     PartitionRange(leftStart, leftEnd)(tableUtils)
   }