stub: physical graph for workflow submission + column lineage

Author: nikhil-zlai

api/py/ai/chronon/cli/entrypoint.py

@@ -1,7 +1,15 @@
+from dataclasses import dataclass
+from typing import Dict, List
 import click
 from datetime import datetime, timedelta
 
-from ai.chronon.cli.compile.compiler import Compiler
+from ai.chronon.api.common.ttypes import ConfigType
+from ai.chronon.cli.plan.physical_index import (
+    PhysicalIndex,
+    get_backfill_physical_graph,
+    submit_physical_graph,
+)
+from ai.chronon.cli.compile.compiler import CompileResult, Compiler
 from ai.chronon.cli.compile.compile_context import CompileContext
 from ai.chronon.cli.git_utils import get_current_branch
 
@@ -12,14 +20,11 @@ def cli():
     pass
 
 
-@cli.command()
-@click.option("--branch", default=get_current_branch, help="Branch to sync")
-def sync(branch):
-    """Sync data for the specified branch"""
-    click.echo(f"\nSyncing data for branch \u001b[32m{branch}\u001b[0m")
+def compile() -> Dict[ConfigType, CompileResult]:
     compile_context = CompileContext()
     compiler = Compiler(compile_context)
-    compiler.compile(compile_context)
+    # TODO(orc): add column lineage to objects
+    return compiler.compile(compile_context)
 
 
 @cli.command()
@@ -34,17 +39,15 @@ def sync(branch):
     default=datetime.now().strftime("%Y-%m-%d"),
     help="End date for backfill (YYYY-MM-DD)",
 )
-@click.option(
-    "--scope",
-    type=click.Choice(["upstream", "self", "downstream"]),
-    default="upstream",
-    help="Scope of configs to backfill",
-)
-def backfill(conf: str, start_date: str, end_date: str, scope: str):
-    """Backfill data between start and end dates"""
-    click.echo(
-        f"Backfilling with scope {scope} for config {conf} from {start_date} to {end_date}"
+@click.option("--branch", default=get_current_branch, help="Branch to sync")
+def backfill(conf: str, start_date: str, end_date: str):
+    # compile
+    compile_result = compile()
+    physical_index = PhysicalIndex.from_compiled_obj(compile_result)
+    physical_graph = get_backfill_physical_graph(
+        conf, physical_index, start_date, end_date
     )
+    submit_physical_graph(physical_graph)
 
 
 @cli.command()

api/py/ai/chronon/cli/plan/physical_index.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+from typing import Dict, List
+from ai.chronon.api.common.ttypes import ConfigType
+from ai.chronon.cli.compile.compile_context import CompiledObj
+from ai.chronon.cli.compile.compiler import CompileResult
+from ai.chronon.lineage.ttypes import Column, ColumnLineage
+from ai.chronon.types import TableDependency
+
+
+@dataclass
+class PhysicalNode:
+    name: str
+    node_type: str  # NodeType
+    conf: CompiledObj
+    table_dependencies: List[TableDependency]
+    output_columns: List[str]
+    output_table: str
+    # TODO: online deps
+
+    def construct_lineage(self):
+        pass
+
+
+def get_physical_nodes(obj: CompiledObj) -> List[PhysicalNode]:
+    pass
+
+
+@dataclass
+class PhysicalGraph:
+    node: PhysicalNode
+    dependencies: List["PhysicalGraph"]
+    start_date: str
+    end_date: str
+
+
+@dataclass
+class PhysicalIndex:
+    table_to_physical: Dict[str, PhysicalNode]
+
+    # TODO incorporate stage column lineage
+    column_lineage: Dict[Column, ColumnLineage]
+
+    @classmethod
+    def from_compiled_obj(
+        cls, compiled_obj: Dict[ConfigType, CompileResult]
+    ) -> "PhysicalIndex":
+        pass
+
+
+def get_backfill_physical_graph(
+    conf: str, index: PhysicalIndex, start_date: str, end_date: str
+) -> PhysicalGraph:
+    pass
+
+
+def get_deploy_physical_graph(conf: str, index: PhysicalIndex) -> PhysicalGraph:
+    pass
+
+
+def submit_physical_graph(physical_graph: PhysicalGraph) -> str:
+    # 1. get physical nodes to submit
+    # 2. find missing nodes from server
+    # 3.
+    pass

api/thrift/api.thrift

@@ -3,6 +3,7 @@ namespace java ai.chronon.api
 
 include "common.thrift"
 include "observability.thrift"
+include "lineage.thrift"
 
 // cd /path/to/chronon
 // thrift --gen py -out api/py/ api/thrift/api.thrift
@@ -257,6 +258,8 @@ struct MetaData {
     // column -> tag_key -> tag_value
     7: optional map<string, map<string, string>> columnTags
 
+    8: optional list<lineage.StageColumnLineage> stageColumnLineages
+
     // marking this as true means that the conf can be served online
     // once marked online, a conf cannot be changed - compiling the conf won't be allowed
     100: optional bool online

api/thrift/lineage.thrift

@@ -0,0 +1,54 @@
+namespace py ai.chronon.lineage
+namespace java ai.chronon.lineage
+
+struct Column {
+    1: optional string logicalNodeName
+    2: optional string stageName
+    3: optional string columnName
+}
+
+enum ColumnLineageStage {
+    SELECT,
+    AGG,
+    DERIVE,
+    JOIN
+}
+
+
+/**
+* NOTE: Staging Query will only be purely output columns.
+*       stage = SELECT, output = col_name, expression = null, input_cols = null, type = DIRECT
+**/
+enum ColumnLineageType {
+    WHERE,
+    JOIN_KEY,
+    GROUP_BY_KEY,
+    DERIVE,
+    TIMESTAMP,
+    DIRECT
+
+    /**
+    * select x+y as a, c, json_extract(payload, "$.meta.user") as user where b = 10
+    * group_by user
+    * agg - count(a, window=7d, bucket=c)
+    *
+    * Stage = SELECT, output = a,  expression = x + y, input_columns = [x, y], type = DIRECT
+    *                              expression = "b = 10", input_columns = [b], type = WHERE
+    *                 output=user, expression = 'json_extract(payload, "$.meta.user")', input_columns = [payload], type = GROUP_BY_KEY
+    *
+    * Stage = AGG,    output = a_count_7d_by_c, expression = 'count(a, window=7d, bucket=c)', input_columns = [a, c], type = DIRECT
+    *                 output=user, input_columns = [user], type = GROUP_BY_KEY
+    **/
+}
+
+struct ColumnLineage {
+    1: optional list<Column> inputColumns
+    2: optional string expression
+    3: optional Column outputColumn
+    4: optional ColumnLineageType lineageType // not present means direct
+}
+
+struct StageColumnLineage {
+    1: optional ColumnLineageStage stage
+    2: optional list<ColumnLineage> lineage
+}