Repository: TobikoData/sqlmesh
Pull request #2247 (2024-03-11 19:45:54)

Fix: List existing data objects concurrently when creating snapshot tables
Make the listing of existing snapshot tables faster by running queries concurrently per unique schema.


Patch:
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index 9a498459..02c9d51d 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -57,7 +57,10 @@ from sqlmesh.core.snapshot import (
     SnapshotTableCleanupTask,
 )
 from sqlmesh.utils import random_id
-from sqlmesh.utils.concurrency import concurrent_apply_to_snapshots
+from sqlmesh.utils.concurrency import (
+    concurrent_apply_to_snapshots,
+    concurrent_apply_to_values,
+)
 from sqlmesh.utils.date import TimeLike, now
 from sqlmesh.utils.errors import AuditError, ConfigError, SQLMeshError
 
@@ -251,11 +254,19 @@ class SnapshotEvaluator:
                 snapshots_with_table_names[snapshot].add(table.name)
                 tables_by_schema[d.schema_(table.db, catalog=table.catalog)].add(table.name)
 
-        existing_objects: t.Set[str] = set()
-        for schema, object_names in tables_by_schema.items():
+        def _get_data_objects(schema: exp.Table) -> t.Set[str]:
             logger.info("Listing data objects in schema %s", schema.sql())
-            objs = self.adapter.get_data_objects(schema, object_names)
-            existing_objects.update(obj.name for obj in objs)
+            objs = self.adapter.get_data_objects(schema, tables_by_schema[schema])
+            return {obj.name for obj in objs}
+
+        with self.concurrent_context():
+            existing_objects = {
+                obj
+                for objs in concurrent_apply_to_values(
+                    list(tables_by_schema), _get_data_objects, self.ddl_concurrent_tasks
+                )
+                for obj in objs
+            }
 
         snapshots_to_create = []
         for snapshot, table_names in snapshots_with_table_names.items():
diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py
index 3bdad505..9881bb5b 100644
--- a/sqlmesh/utils/concurrency.py
+++ b/sqlmesh/utils/concurrency.py
@@ -8,6 +8,8 @@ from sqlmesh.utils.errors import ConfigError, SQLMeshError
 
 H = t.TypeVar("H", bound=t.Hashable)
 S = t.TypeVar("S", bound=SnapshotInfoLike)
+A = t.TypeVar("A")
+R = t.TypeVar("R")
 
 
 class NodeExecutionFailedError(t.Generic[H], SQLMeshError):
@@ -234,3 +236,36 @@ def sequential_apply_to_dag(
             failed_or_skipped_nodes.add(node)
 
     return node_errors, skipped_nodes
+
+
+def concurrent_apply_to_values(
+    values: t.Sequence[A],
+    fn: t.Callable[[A], R],
+    tasks_num: int,
+) -> t.List[R]:
+    """Applies a function to the given collection of values concurrently.
+
+    Args:
+        values: Target values.
+        fn: The function that will be applied concurrently to each value.
+        tasks_num: The number of concurrent tasks.
+
+    Returns:
+        A list of results.
+    """
+    if tasks_num == 1:
+        return [fn(value) for value in values]
+
+    futures: t.List[Future] = [Future() for _ in values]
+
+    def _process_value(value: A, index: int) -> None:
+        try:
+            futures[index].set_result(fn(value))
+        except Exception as ex:
+            futures[index].set_exception(ex)
+
+    with ThreadPoolExecutor(max_workers=tasks_num) as pool:
+        for index, value in enumerate(values):
+            pool.submit(_process_value, value, index)
+
+    return [f.result() for f in futures]


Test patch:
diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py
index f24dad58..892c6ef4 100644
--- a/tests/utils/test_concurrency.py
+++ b/tests/utils/test_concurrency.py
@@ -5,6 +5,7 @@ from sqlmesh.core.snapshot import SnapshotId
 from sqlmesh.utils.concurrency import (
     NodeExecutionFailedError,
     concurrent_apply_to_snapshots,
+    concurrent_apply_to_values,
 )
 
 
@@ -139,3 +140,10 @@ def test_concurrent_apply_to_snapshots_skip_each_node_only_once(mocker: MockerFi
     assert errors[0].node == failed_snapshot.snapshot_id
 
     assert skipped == [snapshot_a.snapshot_id, snapshot_b.snapshot_id, snapshot_c.snapshot_id]
+
+
+@pytest.mark.parametrize("tasks_num", [1, 3])
+def test_concurrent_apply(tasks_num: int):
+    values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    results = concurrent_apply_to_values(values, lambda x: x * 2, tasks_num)
+    assert results == [x * 2 for x in values]