Repository: TobikoData/sqlmesh
Pull request #3189 (2024-09-26 03:59:24)

fix: more perf by offloading mork work to cores


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 3fcf9b4b..4c669b75 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -453,8 +453,9 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         update_model_schemas(
             self.dag,
-            self._models,
-            self.path,
+            models=self._models,
+            audits=self._audits,
+            context_path=self.path,
         )
 
         if model.dialect:
@@ -2057,7 +2058,8 @@ class GenericContext(BaseContext, t.Generic[C]):
     def _new_selector(self) -> Selector:
         return Selector(
             self.state_reader,
-            self._models,
+            models=self._models,
+            audits=self._audits,
             context_path=self.path,
             default_catalog=self.default_catalog,
             dialect=self.default_dialect,
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
index 609c7ea8..e6977e6f 100644
--- a/sqlmesh/core/loader.py
+++ b/sqlmesh/core/loader.py
@@ -14,7 +14,7 @@ from sqlglot.errors import SchemaError, SqlglotError
 from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
-from sqlmesh.core.audit import Audit, load_multiple_audits
+from sqlmesh.core.audit import Audit, ModelAudit, load_multiple_audits
 from sqlmesh.core.dialect import parse
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.metric import Metric, MetricMeta, expand_metrics, load_metric_ddl
@@ -27,7 +27,7 @@ from sqlmesh.core.model import (
     create_external_model,
     load_sql_based_model,
 )
-from sqlmesh.core.model.cache import optimized_query_cache_pool, load_optimized_query_cache
+from sqlmesh.core.model.cache import load_optimized_query_and_mapping, optimized_query_cache_pool
 from sqlmesh.core.model import model as model_registry
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.dag import DAG
@@ -104,8 +104,9 @@ class Loader(abc.ABC):
         if update_schemas:
             update_model_schemas(
                 self._dag,
-                models,
-                self._context.path,
+                models=models,
+                audits={k: v for k, v in audits.items() if isinstance(v, ModelAudit)},
+                context_path=self._context.path,
             )
             for model in models.values():
                 # The model definition can be validated correctly only after the schema is set.
@@ -552,6 +553,7 @@ class SqlMeshLoader(Loader):
 def update_model_schemas(
     dag: DAG[str],
     models: UniqueKeyDict[str, Model],
+    audits: t.Dict[str, ModelAudit],
     context_path: Path,
 ) -> None:
     schema = MappingSchema(normalize=False)
@@ -560,14 +562,14 @@ def update_model_schemas(
     if c.MAX_FORK_WORKERS == 1:
         _update_model_schemas_sequential(dag, models, schema, optimized_query_cache)
     else:
-        _update_model_schemas_parallel(dag, models, schema, optimized_query_cache)
+        _update_model_schemas_parallel(dag, models, audits, schema, optimized_query_cache)
 
 
 def _update_schema_with_model(schema: MappingSchema, model: Model) -> None:
     columns_to_types = model.columns_to_types
     if columns_to_types:
         try:
-            schema.add_table(model.fqn, columns_to_types, dialect=model.dialect, normalize=False)
+            schema.add_table(model.fqn, columns_to_types, dialect=model.dialect)
         except SchemaError as e:
             if "nesting level:" in str(e):
                 logger.error(
@@ -597,6 +599,7 @@ def _update_model_schemas_sequential(
 def _update_model_schemas_parallel(
     dag: DAG[str],
     models: UniqueKeyDict[str, Model],
+    audits: t.Dict[str, ModelAudit],
     schema: MappingSchema,
     optimized_query_cache: OptimizedQueryCache,
 ) -> None:
@@ -617,19 +620,29 @@ def _update_model_schemas_parallel(
             if not deps:
                 del graph[name]
                 model = models[name]
-                model.update_schema(schema)
-                futures.add(executor.submit(load_optimized_query_cache, model))
+                futures.add(
+                    executor.submit(
+                        load_optimized_query_and_mapping,
+                        model,
+                        mapping={
+                            parent: models[parent].columns_to_types
+                            for parent in model.depends_on
+                            if parent in models
+                        },
+                    )
+                )
 
-    with optimized_query_cache_pool(optimized_query_cache) as executor:
+    with optimized_query_cache_pool(optimized_query_cache, audits) as executor:
         process_models()
 
         while futures:
             for future in as_completed(futures):
                 futures.remove(future)
-                fqn, entry_name = future.result()
+                fqn, entry_name, data_hash, metadata_hash, mapping_schema = future.result()
                 model = models[fqn]
-                if entry_name:
-                    optimized_query_cache.with_optimized_query(model, entry_name)
-
+                model._data_hash = data_hash
+                model._metadata_hash = metadata_hash
+                model.set_mapping_schema(mapping_schema)
+                optimized_query_cache.with_optimized_query(model, entry_name)
                 _update_schema_with_model(schema, model)
                 process_models(completed_model=model)
diff --git a/sqlmesh/core/model/cache.py b/sqlmesh/core/model/cache.py
index 4a79634a..7b31a988 100644
--- a/sqlmesh/core/model/cache.py
+++ b/sqlmesh/core/model/cache.py
@@ -8,6 +8,7 @@ from pathlib import Path
 
 from sqlglot import exp
 from sqlglot.optimizer.simplify import gen
+from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.model.definition import Model, SqlModel, _Model
@@ -18,7 +19,11 @@ from dataclasses import dataclass
 
 logger = logging.getLogger(__name__)
 
-T = t.TypeVar("T")
+if t.TYPE_CHECKING:
+    from sqlmesh.core.audit import ModelAudit
+    from sqlmesh.core.snapshot import SnapshotId
+
+    T = t.TypeVar("T")
 
 
 class ModelCache:
@@ -133,47 +138,68 @@ class OptimizedQueryCache:
         return f"{model.name}_{crc32(hash_data)}"
 
 
-def optimized_query_cache_pool(optimized_query_cache: OptimizedQueryCache) -> ProcessPoolExecutor:
+def optimized_query_cache_pool(
+    optimized_query_cache: OptimizedQueryCache,
+    audits: t.Optional[t.Dict[str, ModelAudit]] = None,
+) -> ProcessPoolExecutor:
     return ProcessPoolExecutor(
         mp_context=mp.get_context("fork"),
         initializer=_init_optimized_query_cache,
-        initargs=(optimized_query_cache,),
+        initargs=(optimized_query_cache, audits),
         max_workers=c.MAX_FORK_WORKERS,
     )
 
 
-@t.overload
-def load_optimized_query_cache(
-    model_or_tuple: t.Tuple[Model, T],
-) -> t.Tuple[T, t.Optional[str]]: ...
+_optimized_query_cache: t.Optional[OptimizedQueryCache] = None
+_audits: t.Optional[t.Dict[str, ModelAudit]] = None
 
 
-@t.overload
-def load_optimized_query_cache(model_or_tuple: Model) -> t.Tuple[str, t.Optional[str]]: ...
+def _init_optimized_query_cache(
+    optimized_query_cache: OptimizedQueryCache,
+    audits: t.Optional[t.Dict[str, ModelAudit]] = None,
+) -> None:
+    global _optimized_query_cache, _audits
+    _optimized_query_cache = optimized_query_cache
+    _audits = audits
 
 
-def load_optimized_query_cache(model_or_tuple):  # type: ignore
+def load_optimized_query(
+    model_snapshot_id: t.Tuple[Model, SnapshotId],
+) -> t.Tuple[SnapshotId, t.Optional[str]]:
     assert _optimized_query_cache
-
-    if isinstance(model_or_tuple, _Model):
-        model = model_or_tuple
-        key = None
-    else:
-        model, key = model_or_tuple
+    model, snapshot_id = model_snapshot_id
 
     if isinstance(model, SqlModel):
         entry_name = _optimized_query_cache.put(model)
     else:
         entry_name = None
-    return key or model.fqn, entry_name
+    return snapshot_id, entry_name
 
 
-_optimized_query_cache: t.Optional[OptimizedQueryCache] = None
+def load_optimized_query_and_mapping(
+    model: Model, mapping: t.Dict
+) -> t.Tuple[str, t.Optional[str], str, str, t.Dict]:
+    assert _optimized_query_cache
+    assert _audits is not None
 
+    schema = MappingSchema(normalize=False)
+    for parent, columns_to_types in mapping.items():
+        schema.add_table(parent, columns_to_types, dialect=model.dialect)
+    model.update_schema(schema)
 
-def _init_optimized_query_cache(optimized_query_cache: OptimizedQueryCache) -> None:
-    global _optimized_query_cache
-    _optimized_query_cache = optimized_query_cache
+    if isinstance(model, SqlModel):
+        entry_name = _optimized_query_cache._entry_name(model)
+        _optimized_query_cache.with_optimized_query(model, entry_name)
+    else:
+        entry_name = None
+
+    return (
+        model.fqn,
+        entry_name,
+        model.data_hash,
+        model.metadata_hash(_audits),
+        model.mapping_schema,
+    )
 
 
 def _mapping_schema_hash_data(schema: t.Dict[str, t.Any]) -> t.List[str]:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 38ba2939..592bb63f 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -627,10 +627,11 @@ class _Model(ModelMeta, frozen=True):
             return to_time_column(time, time_column_type, self.time_column.format)
         return exp.convert(time)
 
-    def update_schema(
-        self,
-        schema: MappingSchema,
-    ) -> None:
+    def set_mapping_schema(self, schema: t.Dict) -> None:
+        self.mapping_schema.clear()
+        self.mapping_schema.update(schema)
+
+    def update_schema(self, schema: MappingSchema) -> None:
         """Updates the schema for this model's dependencies based on the given mapping schema."""
         for dep in self.depends_on:
             table = exp.to_table(dep)
@@ -1109,11 +1110,15 @@ class SqlModel(_SqlBasedModel):
             if select.comments
         }
 
-    def update_schema(
-        self,
-        schema: MappingSchema,
-    ) -> None:
+    def set_mapping_schema(self, schema: t.Dict) -> None:
+        super().set_mapping_schema(schema)
+        self._on_mapping_schema_set()
+
+    def update_schema(self, schema: MappingSchema) -> None:
         super().update_schema(schema)
+        self._on_mapping_schema_set()
+
+    def _on_mapping_schema_set(self) -> None:
         self._columns_to_types = None
         self._query_renderer.update_schema(self.mapping_schema)
 
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
index 7176c7cf..a97bb660 100644
--- a/sqlmesh/core/selector.py
+++ b/sqlmesh/core/selector.py
@@ -11,8 +11,6 @@ from sqlglot import exp
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.loader import update_model_schemas
-from sqlmesh.core.model import Model
-from sqlmesh.core.state_sync import StateReader
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.git import GitClient
@@ -20,11 +18,18 @@ from sqlmesh.utils.git import GitClient
 logger = logging.getLogger(__name__)
 
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.audit import ModelAudit
+    from sqlmesh.core.model import Model
+    from sqlmesh.core.state_sync import StateReader
+
+
 class Selector:
     def __init__(
         self,
         state_reader: StateReader,
         models: UniqueKeyDict[str, Model],
+        audits: t.Dict[str, ModelAudit],
         context_path: Path = Path("."),
         dag: t.Optional[DAG[str]] = None,
         default_catalog: t.Optional[str] = None,
@@ -32,6 +37,7 @@ class Selector:
     ):
         self._state_reader = state_reader
         self._models = models
+        self._audits = audits
         self._context_path = context_path
         self._default_catalog = default_catalog
         self._dialect = dialect
@@ -144,7 +150,9 @@ class Selector:
             models[model.fqn] = model
 
         if needs_update:
-            update_model_schemas(dag, models, self._context_path)
+            update_model_schemas(
+                dag, models=models, audits=self._audits, context_path=self._context_path
+            )
 
         return models
 
diff --git a/sqlmesh/core/snapshot/cache.py b/sqlmesh/core/snapshot/cache.py
index 0ace6f14..f99dca64 100644
--- a/sqlmesh/core/snapshot/cache.py
+++ b/sqlmesh/core/snapshot/cache.py
@@ -6,7 +6,7 @@ from pathlib import Path
 from sqlmesh.core.model.cache import (
     OptimizedQueryCache,
     optimized_query_cache_pool,
-    load_optimized_query_cache,
+    load_optimized_query,
 )
 from sqlmesh.core import constants as c
 from sqlmesh.core.snapshot.definition import Snapshot, SnapshotId
@@ -54,7 +54,7 @@ class SnapshotCache:
         if c.MAX_FORK_WORKERS != 1:
             with optimized_query_cache_pool(self._optimized_query_cache) as executor:
                 for key, entry_name in executor.map(
-                    load_optimized_query_cache,
+                    load_optimized_query,
                     (
                         (snapshot.model, s_id)
                         for s_id, snapshot in snapshots.items()


Test patch:
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index 3ee4b3d6..f97264fa 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -87,7 +87,7 @@ def test_select_models(mocker: MockerFixture, make_snapshot, default_catalog: t.
     local_models[modified_model_v2.fqn] = modified_model_v2.copy(
         update={"mapping_schema": added_model_schema}
     )
-    selector = Selector(state_reader_mock, local_models, default_catalog=default_catalog)
+    selector = Selector(state_reader_mock, local_models, {}, default_catalog=default_catalog)
 
     _assert_models_equal(
         selector.select_models(["db.added_model"], env_name),
@@ -242,7 +242,7 @@ def test_select_models_expired_environment(mocker: MockerFixture, make_snapshot)
 
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[modified_model_v2.fqn] = modified_model_v2
-    selector = Selector(state_reader_mock, local_models)
+    selector = Selector(state_reader_mock, local_models, {})
 
     _assert_models_equal(
         selector.select_models(["*.modified_model"], env_name, fallback_env_name="prod"),
@@ -298,7 +298,7 @@ def test_select_change_schema(mocker: MockerFixture, make_snapshot):
     local_child = child.copy(update={"mapping_schema": {'"db"': {'"parent"': {"b": "INT"}}}})
     local_models[local_child.fqn] = local_child
 
-    selector = Selector(state_reader_mock, local_models)
+    selector = Selector(state_reader_mock, local_models, {})
 
     selected = selector.select_models(["db.parent"], env_name)
     assert selected[local_child.fqn].data_hash != child.data_hash
@@ -332,7 +332,7 @@ def test_select_models_missing_env(mocker: MockerFixture, make_snapshot):
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[model.fqn] = model
 
-    selector = Selector(state_reader_mock, local_models)
+    selector = Selector(state_reader_mock, local_models, {})
 
     assert selector.select_models([model.name], "missing_env").keys() == {model.fqn}
     assert not selector.select_models(["missing"], "missing_env")
@@ -508,7 +508,7 @@ def test_expand_model_selections(
         )
         models[model.fqn] = model
 
-    selector = Selector(mocker.Mock(), models)
+    selector = Selector(mocker.Mock(), models, {})
     assert selector.expand_model_selections(selections) == output
 
 
@@ -521,7 +521,7 @@ def test_model_selection_normalized(mocker: MockerFixture, make_snapshot):
         dialect="bigquery",
     )
     models[model.fqn] = model
-    selector = Selector(mocker.Mock(), models, dialect="bigquery")
+    selector = Selector(mocker.Mock(), models, {}, dialect="bigquery")
     assert selector.expand_model_selections(["db.test_Model"]) == {'"db"."test_Model"'}
 
 
@@ -569,7 +569,7 @@ def test_expand_git_selection(
     git_client_mock.list_uncommitted_changed_files.return_value = []
     git_client_mock.list_committed_changed_files.return_value = [model_a._path, model_c._path]
 
-    selector = Selector(mocker.Mock(), models)
+    selector = Selector(mocker.Mock(), models, {})
     selector._git_client = git_client_mock
 
     assert selector.expand_model_selections(expressions) == expected_fqns
@@ -603,7 +603,7 @@ def test_select_models_with_external_parent(mocker: MockerFixture):
     local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
     local_models[added_model.fqn] = added_model
 
-    selector = Selector(state_reader_mock, local_models, default_catalog=default_catalog)
+    selector = Selector(state_reader_mock, local_models, {}, default_catalog=default_catalog)
 
     expanded_selections = selector.expand_model_selections(["+*added_model*"])
     assert expanded_selections == {added_model.fqn}