Repository: TobikoData/sqlmesh
Pull request #2010 (2024-01-24 00:06:14)

fix: use snapshot names in scheduler dag generation for performance


Patch:
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
index 86d33665..b7c601fa 100644
--- a/sqlmesh/core/scheduler.py
+++ b/sqlmesh/core/scheduler.py
@@ -38,7 +38,9 @@ logger = logging.getLogger(__name__)
 Interval = t.Tuple[datetime, datetime]
 Batch = t.List[Interval]
 SnapshotToBatches = t.Dict[Snapshot, Batch]
-SchedulingUnit = t.Tuple[Snapshot, t.Tuple[Interval, int]]
+# we store snapshot name instead of snapshots/snapshotids because pydantic
+# is extremely slow to hash. snapshot names should be unique within a dag run
+SchedulingUnit = t.Tuple[str, t.Tuple[Interval, int]]
 
 
 class Scheduler:
@@ -250,23 +252,20 @@ class Scheduler:
 
         dag = self._dag(batches)
 
-        visited = set()
-        for snapshot, _ in dag.sorted:
-            if snapshot in visited:
-                continue
-            visited.add(snapshot)
-
         self.console.start_evaluation_progress(
             {snapshot: len(intervals) for snapshot, intervals in batches.items()},
             environment_naming_info,
             self.default_catalog,
         )
 
+        snapshots_by_name = {snapshot.name: snapshot for snapshot in self.snapshots.values()}
+
         def evaluate_node(node: SchedulingUnit) -> None:
             if circuit_breaker and circuit_breaker():
                 raise CircuitBreakerError()
 
-            snapshot, ((start, end), batch_idx) = node
+            snapshot_name, ((start, end), batch_idx) = node
+            snapshot = snapshots_by_name[snapshot_name]
 
             self.console.start_snapshot_evaluation_progress(snapshot)
 
@@ -334,7 +333,7 @@ class Scheduler:
             if not intervals:
                 continue
             upstream_dependencies = [
-                (self.snapshots[p_sid], (interval, i))
+                (p_sid.name, (interval, i))
                 for p_sid in snapshot.parents
                 if p_sid in self.snapshots
                 for i, interval in enumerate(
@@ -348,11 +347,10 @@ class Scheduler:
                 )
             ]
             for i, interval in enumerate(intervals):
-                dag.add((snapshot, (interval, i)), upstream_dependencies)
-                if snapshot.depends_on_past:
+                dag.add((snapshot.name, (interval, i)), upstream_dependencies)
+                if snapshot.depends_on_past and i > 0:
                     dag.add(
-                        (snapshot, (interval, i)),
-                        [(snapshot, (_interval, _i)) for _i, _interval in enumerate(intervals[:i])],
+                        (snapshot.name, (interval, i)), [(snapshot.name, (intervals[i - 1], i - 1))]
                     )
         return dag
 


Test patch:
diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py
index 8caf352a..3635a0c0 100644
--- a/tests/core/test_scheduler.py
+++ b/tests/core/test_scheduler.py
@@ -129,7 +129,7 @@ def test_incremental_by_unique_key_kind_dag(mocker: MockerFixture, make_snapshot
     assert dag.graph == {
         # Depends on no one
         (
-            unique_by_key_snapshot,
+            unique_by_key_snapshot.name,
             ((to_datetime("2023-01-01"), to_datetime("2023-01-07")), 0),
         ): set(),
     }
@@ -170,33 +170,33 @@ def test_incremental_time_self_reference_dag(mocker: MockerFixture, make_snapsho
     assert dag.graph == {
         # Only run one day at a time and each day relies on the previous days
         (
-            incremental_self_snapshot,
+            incremental_self_snapshot.name,
             ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
         ): set(),
-        (incremental_self_snapshot, ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1)): {
-            (incremental_self_snapshot, ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0))
-        },
-        (incremental_self_snapshot, ((to_datetime("2023-01-04"), to_datetime("2023-01-05")), 2)): {
+        (
+            incremental_self_snapshot.name,
+            ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1),
+        ): {
             (
-                incremental_self_snapshot,
+                incremental_self_snapshot.name,
                 ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
-            ),
-            (
-                incremental_self_snapshot,
-                ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1),
-            ),
+            )
         },
-        (incremental_self_snapshot, ((to_datetime("2023-01-06"), to_datetime("2023-01-07")), 3)): {
-            (
-                incremental_self_snapshot,
-                ((to_datetime("2023-01-01"), to_datetime("2023-01-02")), 0),
-            ),
+        (
+            incremental_self_snapshot.name,
+            ((to_datetime("2023-01-04"), to_datetime("2023-01-05")), 2),
+        ): {
             (
-                incremental_self_snapshot,
+                incremental_self_snapshot.name,
                 ((to_datetime("2023-01-03"), to_datetime("2023-01-04")), 1),
             ),
+        },
+        (
+            incremental_self_snapshot.name,
+            ((to_datetime("2023-01-06"), to_datetime("2023-01-07")), 3),
+        ): {
             (
-                incremental_self_snapshot,
+                incremental_self_snapshot.name,
                 ((to_datetime("2023-01-04"), to_datetime("2023-01-05")), 2),
             ),
         },
diff --git a/tests/integrations/github/cicd/test_integration.py b/tests/integrations/github/cicd/test_integration.py
index e0924ea7..b7ea6cfd 100644
--- a/tests/integrations/github/cicd/test_integration.py
+++ b/tests/integrations/github/cicd/test_integration.py
@@ -1781,7 +1781,8 @@ def test_capture_console_errors(
         """:x: Failed to create or update PR Environment `hello_world_2`.
 **Errors:**
 ```
-FAILED processing snapshot Snapshot<name: "memory"."sushi"."waiter_revenue_by_day","""
+FAILED processing snapshot "memory"."sushi"."waiter_revenue_by_day"
+"""
     )
 
     assert "SQLMesh - Prod Plan Preview" in controller._check_run_mapping