rebase and lint

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

api/python/ai/chronon/dag_builder.py

@@ -0,0 +1,59 @@
+from typing import Union
+
+import ai.chronon.utils as c_utils
+from ai.chronon.api.common.ttypes import TableDependency, TableInfo, TimeUnit, Window
+from ai.chronon.api.ttypes import BootstrapPart, GroupBy, JoinPart, Source
+
+"""
+Given a node that represents an upstream, turn it into a TableDependency.
+"""
+
+
+def to_dependency(node: Union[Source, GroupBy, JoinPart, BootstrapPart], lag: int = 0) -> TableDependency:
+    if isinstance(node, JoinPart):
+        groupby_dep = to_dependency(node.groupBy, lag=lag)
+        return groupby_dep
+    elif isinstance(node, GroupBy):
+        ti = node.metaData.executionInfo.outputTableInfo
+        return TableDependency(
+            tableInfo=ti,
+            startOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            endOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            startCutOff=None,
+            endCutOff=None,
+            forceCompute=False,
+        )
+    elif isinstance(node, BootstrapPart):
+        partition_column = node.query.partitionColumn if node.query is not None else None
+        table_info = TableInfo(
+            table=node.table,
+            partitionColumn=partition_column,
+        )
+        start_cutoff = node.query.startPartition if node.query is not None else None
+        end_cutoff = node.query.endPartition if node.query is not None else None
+        return TableDependency(
+            tableInfo=table_info,
+            startOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            endOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            startCutOff=start_cutoff,
+            endCutOff=end_cutoff,
+            forceCompute=False,
+        )
+    else:  # When type of node is a Source
+        table_name = c_utils.get_table(node)
+        query = c_utils.get_query(node)
+        partition_column = query.partitionColumn
+        table_info = TableInfo(
+            table=table_name,
+            partitionColumn=partition_column,
+        )
+        start_cutoff = query.startPartition if query is not None else None
+        end_cutoff = query.endPartition if query is not None else None
+        return TableDependency(
+            tableInfo=table_info,
+            startOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            endOffset=Window(timeUnit=TimeUnit.DAYS, length=0),
+            startCutOff=start_cutoff,
+            endCutOff=end_cutoff,
+            forceCompute=False,
+        )

api/python/ai/chronon/group_by.py

@@ -20,6 +20,7 @@
 
 import ai.chronon.api.common.ttypes as common
 import ai.chronon.api.ttypes as ttypes
+import ai.chronon.dag_builder as dag_builder
 import ai.chronon.utils as utils
 import ai.chronon.windows as window_utils
 
@@ -629,6 +630,7 @@ def _normalize_source(source):
         sources = [sources]
 
     sources = [_sanitize_columns(_normalize_source(source)) for source in sources]
+    table_deps = [dag_builder.to_dependency(s) for s in sources]
 
     # get caller's filename to assign team
     team = inspect.stack()[1].filename.split("/")[-2]
@@ -639,6 +641,7 @@ def _normalize_source(source):
         env=env_vars,
         stepDays=step_days,
         historicalBackfill=disable_historical_backfill,
+        tableDependencies=table_deps
     )
 
     column_tags = {}

api/python/ai/chronon/join.py

@@ -21,6 +21,7 @@
 
 import ai.chronon.api.common.ttypes as common
 import ai.chronon.api.ttypes as api
+import ai.chronon.dag_builder as dag_builder
 import ai.chronon.repo.extract_objects as eo
 import ai.chronon.utils as utils
 
@@ -510,6 +511,13 @@ def Join(
     """
     # create a deep copy for case: multiple LeftOuterJoin use the same left,
     # validation will fail after the first iteration
+
+    if bootstrap_parts is None:
+        bootstrap_parts = []
+    if online_external_parts is None:
+        online_external_parts = []
+    if right_parts is None:
+        right_parts = []
     updated_left = copy.deepcopy(left)
     if left.events and left.events.query.selects:
         assert "ts" not in left.events.query.selects.keys(), (
@@ -562,12 +570,17 @@ def Join(
             )
         ]
 
+    table_deps = [dag_builder.to_dependency(left)] \
+        + [dag_builder.to_dependency(right_part) for right_part in right_parts] \
+        + [dag_builder.to_dependency(b_part) for b_part in bootstrap_parts]
+
     exec_info = common.ExecutionInfo(
         scheduleCron=offline_schedule,
         conf=conf,
         env=env_vars,
         stepDays=step_days,
         historicalBackfill=historical_backfill,
+        tableDependencies=table_deps
     )
 
     metadata = api.MetaData(