perf: reuse table mapping during snapshot promotion

sqlmesh/core/renderer.py

@@ -149,6 +149,7 @@ def _render(
         table_mapping: t.Optional[t.Dict[str, str]] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         runtime_stage: RuntimeStage = RuntimeStage.LOADING,
+        render_table_mapping: t.Optional[TableMapping] = None,
         **kwargs: t.Any,
     ) -> t.List[t.Optional[exp.Expr]]:
         """Renders a expression, expanding macros with provided kwargs
@@ -161,6 +162,7 @@ def _render(
             table_mapping: Table mapping of physical locations. Takes precedence over snapshot mappings.
             deployability_index: Determines snapshots that are deployable in the context of this evaluation.
             runtime_stage: Indicates the current runtime stage, for example if we're still loading the project, etc.
+            render_table_mapping: Shared snapshot mapping and overrides prepared for this promotion.
             kwargs: Additional kwargs to pass to the renderer.
 
         Returns:
@@ -232,6 +234,7 @@ def _resolve_table(table: str | exp.Table) -> str:
                 end=end,
                 execution_time=execution_time,
                 runtime_stage=runtime_stage,
+                render_table_mapping=render_table_mapping,
             ),
             snapshots=snapshots,
             default_catalog=self._default_catalog,
@@ -430,6 +433,7 @@ def _resolve_tables(
         table_mapping: t.Optional[t.Dict[str, str]] = None,
         expand: t.Iterable[str] = tuple(),
         deployability_index: t.Optional[DeployabilityIndex] = None,
+        render_table_mapping: t.Optional[TableMapping] = None,
         **kwargs: t.Any,
     ) -> E:
         if not snapshots and not table_mapping and not expand:
@@ -444,10 +448,14 @@ def _resolve_tables(
 
             snapshots = snapshots or {}
             table_mapping = table_mapping or {}
-            mapping = {
-                **self._to_table_mapping(snapshots.values(), deployability_index),
-                **table_mapping,
-            }
+            mapping: t.Dict[str, str] = (
+                render_table_mapping
+                if render_table_mapping is not None
+                else {
+                    **self._to_table_mapping(snapshots.values(), deployability_index),
+                    **table_mapping,
+                }
+            )
             expand = set(expand) | {
                 name for name, snapshot in snapshots.items() if snapshot.is_embedded
             }
@@ -477,6 +485,7 @@ def _expand(node: exp.Expr) -> exp.Expr:
                                 table_mapping=table_mapping,
                                 expand=expand,
                                 deployability_index=deployability_index,
+                                render_table_mapping=render_table_mapping,
                                 **kwargs,
                             )
                             if nested_query is not None:
@@ -489,6 +498,20 @@ def _expand(node: exp.Expr) -> exp.Expr:
 
                 expression = expression.transform(_expand, copy=False)  # type: ignore
 
+            if isinstance(mapping, TableMapping):
+                # Normalize the shared mapping once, then replace only tables in this expression.
+                normalized_keys = mapping.normalized_keys(self._dialect)
+                mapping = {
+                    key: mapping[key]
+                    for table in expression.find_all(exp.Table)
+                    if (
+                        key := normalized_keys.get(
+                            exp.normalize_table_name(table, dialect=self._dialect)
+                        )
+                    )
+                    is not None
+                }
+
             if mapping:
                 expression = exp.replace_tables(
                     expression, mapping, dialect=self._dialect, copy=False

sqlmesh/core/snapshot/evaluator.py

@@ -41,6 +41,7 @@
 from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy, DataObjectType, DataObject
 from sqlmesh.core.model.meta import GrantsTargetLayer
 from sqlmesh.core.macros import RuntimeStage
+from sqlmesh.core.renderer import TableMapping
 from sqlmesh.core.model import (
     AuditResult,
     IncrementalUnmanagedKind,
@@ -286,6 +287,8 @@ def promote(
             on_complete: A callback to call on each successfully promoted snapshot.
         """
 
+        from sqlmesh.core.snapshot import to_table_mapping
+
         tables_by_gateway: t.Dict[t.Union[str, None], t.List[exp.Table]] = defaultdict(list)
         for snapshot in target_snapshots:
             if snapshot.is_model and not snapshot.is_symbolic:
@@ -314,6 +317,16 @@ def promote(
         deployability_index = deployability_index or DeployabilityIndex.all_deployable()
         # Renderers look snapshots up by model name, not by SnapshotId.
         snapshots_by_name = {s.name: s for s in (snapshots or {}).values()}
+        # Share the physical mapping and overrides across all promotion renders and workers.
+        render_table_mapping = TableMapping(
+            {
+                **to_table_mapping(snapshots_by_name.values(), deployability_index),
+                **(table_mapping or {}),
+            }
+        )
+        for snapshot in target_snapshots:
+            if snapshot.is_model:
+                render_table_mapping.normalized_keys(snapshot.model.dialect)
         with self.concurrent_context():
             concurrent_apply_to_snapshots(
                 target_snapshots,
@@ -324,6 +337,7 @@ def promote(
                     execution_time=execution_time,
                     snapshots=snapshots_by_name,
                     table_mapping=table_mapping,
+                    render_table_mapping=render_table_mapping,
                     environment_naming_info=environment_naming_info,
                     deployability_index=deployability_index,  # type: ignore
                     on_complete=on_complete,
@@ -1264,6 +1278,7 @@ def _promote_snapshot(
         execution_time: t.Optional[TimeLike] = None,
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         table_mapping: t.Optional[t.Dict[str, str]] = None,
+        render_table_mapping: t.Optional[TableMapping] = None,
     ) -> None:
         if not snapshot.is_model:
             return
@@ -1298,11 +1313,16 @@ def _promote_snapshot(
                 environment=environment_naming_info.name,
                 snapshots=snapshots,
                 snapshot=snapshot,
+                render_table_mapping=render_table_mapping,
                 **render_kwargs,
             )
 
             adapter.execute(
-                snapshot.model.render_on_virtual_update(snapshots=snapshots, **render_kwargs)
+                snapshot.model.render_on_virtual_update(
+                    snapshots=snapshots,
+                    render_table_mapping=render_table_mapping,
+                    **render_kwargs,
+                )
             )
 
         if on_complete is not None:
@@ -2016,6 +2036,7 @@ def promote(
             snapshots=kwargs.get("snapshots"),
             deployability_index=kwargs.get("deployability_index"),
             table_mapping=kwargs.get("table_mapping"),
+            render_table_mapping=kwargs.get("render_table_mapping"),
             runtime_stage=kwargs.get("runtime_stage"),
         )
         self.adapter.create_view(

tests/core/test_model.py

@@ -15,7 +15,7 @@
 from sqlglot.schema import MappingSchema
 from sqlmesh.cli.project_init import init_example_project, ProjectTemplate
 from sqlmesh.core.environment import EnvironmentNamingInfo
-from sqlmesh.core.renderer import TableMapping
+from sqlmesh.core.renderer import ExpressionRenderer, TableMapping
 from sqlmesh.core.model.kind import TimeColumn, ModelKindName, SeedKind
 
 from sqlmesh import CustomMaterialization, CustomKind
@@ -9900,7 +9900,10 @@ def resolve_named(evaluator, name):
     )
 
 
-def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.Callable):
+@pytest.mark.parametrize("use_cached_mapping", [False, True])
+def test_resolve_tables_expand_reveals_table_after_find_check(
+    make_snapshot: t.Callable, use_cached_mapping: bool
+):
     """Embedded-model expansion (`expand=`) runs as an `expression.transform` *before* the new
     `expression.find(exp.Table)` short-circuit in `_resolve_tables`, so a table reference that
     only exists *after* inlining an embedded model's query must still be seen by `find()` and
@@ -9909,6 +9912,8 @@ def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.C
     must still resolve to its physical table name, not be silently skipped because it wasn't
     present at the time `_resolve_tables` was first called."""
 
+    from sqlmesh.core.snapshot import to_table_mapping
+
     grandparent = load_sql_based_model(d.parse("MODEL (name grandparent); SELECT 1 AS c"))
     grandparent_snapshot = make_snapshot(grandparent)
     grandparent_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
@@ -9922,7 +9927,12 @@ def test_resolve_tables_expand_reveals_table_after_find_check(make_snapshot: t.C
     child = load_sql_based_model(d.parse("MODEL (name child); SELECT c FROM mid;"))
 
     snapshots = {'"grandparent"': grandparent_snapshot, '"mid"': mid_snapshot}
-    query = child.render_query(snapshots=snapshots)
+    query = child.render_query(
+        snapshots=snapshots,
+        render_table_mapping=(
+            TableMapping(to_table_mapping(snapshots.values(), None)) if use_cached_mapping else None
+        ),
+    )
     assert query is not None
     rendered_sql = query.sql()
 
@@ -10071,6 +10081,28 @@ def items(self):
     assert ItemsCountingDict.items_call_count == 0
 
 
+@pytest.mark.parametrize(
+    "dialect",
+    [
+        "duckdb,normalization_strategy=case_insensitive",
+        "snowflake,normalization_strategy=uppercase",
+    ],
+)
+def test_cached_table_mapping_preserves_key_order(dialect):
+    # Updating an existing raw key must not move it after a later dialect-equivalent key.
+    table_mapping = {'"DB"."T"': "early_table", "DB.T": "last_table"}
+    table_mapping.update({'"DB"."T"': "override_table"})
+    renderer = ExpressionRenderer(parse_one("SELECT * FROM DB.T", dialect=dialect), dialect, [])
+
+    uncached = renderer.render(table_mapping=table_mapping)
+    cached = renderer.render(
+        table_mapping=table_mapping, render_table_mapping=TableMapping(table_mapping)
+    )
+
+    assert cached == uncached
+    assert exp.table_name(cached[0].find(exp.Table)) == "last_table"
+
+
 def test_table_mapping_normalized_keys():
     table_mapping = TableMapping({'"db"."a"': "view_a", "db.A": "view_a_upper"})
 
@@ -10414,7 +10446,8 @@ def dummy_model_entry(evaluator: MacroEvaluator) -> exp.Select:
     assert len(context.engine_adapters) == 2
 
 
-def test_model_on_virtual_update(make_snapshot: t.Callable):
+@pytest.mark.parametrize("use_cached_mapping", [False, True])
+def test_model_on_virtual_update(make_snapshot: t.Callable, use_cached_mapping: bool):
     # Macro to test resolution within virtual statement
     @macro()
     def resolve_parent_name(evaluator, name):
@@ -10495,6 +10528,7 @@ def resolve_parent_name(evaluator, name):
     rendered_on_virtual_update = model.render_on_virtual_update(
         snapshots=snapshots,
         table_mapping=table_mapping,
+        render_table_mapping=TableMapping(table_mapping) if use_cached_mapping else None,
         this_model=model_snapshot.qualified_view_name.table_for_environment(
             environment_naming_info, dialect=dialect
         ),
@@ -10526,6 +10560,7 @@ def resolve_parent_name(evaluator, name):
     rendered_parent_on_virtual_update = parent.render_on_virtual_update(
         snapshots=snapshots,
         table_mapping=table_mapping,
+        render_table_mapping=TableMapping(table_mapping) if use_cached_mapping else None,
         this_model=parent_snapshot.qualified_view_name.table_for_environment(
             environment_naming_info, dialect=dialect
         ),

tests/core/test_snapshot_evaluator.py

@@ -58,6 +58,7 @@
     SnapshotTableCleanupTask,
 )
 from sqlmesh.core.snapshot.definition import to_view_mapping
+from sqlmesh.core import snapshot as snapshot_module
 from sqlmesh.core.snapshot.evaluator import (
     CustomMaterialization,
     EngineManagedStrategy,
@@ -5766,3 +5767,95 @@ def test_promote_resolves_this_model_with_single_mapping_entry(
     assert spy.call_count == 20
     for call in spy.call_args_list:
         assert len(call.args[1]) == 1
+
+
+@pytest.mark.parametrize("ddl_concurrent_tasks", [1, 4])
+@pytest.mark.parametrize("use_view_mapping", [False, True])
+def test_promote_reuses_table_mapping(
+    mocker: MockerFixture, make_snapshot, ddl_concurrent_tasks, use_view_mapping
+):
+    @macro()
+    def resolve_upstream_query(evaluator):
+        return evaluator.resolve_tables(
+            exp.select("COUNT(*)").from_("test_schema.model_0").subquery()
+        )
+
+    adapter_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter")
+    adapter_mock.dialect = "duckdb"
+    adapter_mock.with_settings.return_value = adapter_mock
+    evaluator = SnapshotEvaluator(adapter_mock, ddl_concurrent_tasks=ddl_concurrent_tasks)
+
+    snapshots = {}
+    for i in range(10):
+        model = load_sql_based_model(
+            d.parse(
+                f"""
+                MODEL (
+                    name test_schema.model_{i},
+                    kind FULL,
+                    virtual_properties (
+                        upstream_count = (SELECT COUNT(*) FROM test_schema.model_0),
+                        macro_count = @resolve_upstream_query()
+                    ),
+                );
+                SELECT 1 AS a;
+                ON_VIRTUAL_UPDATE_BEGIN;
+                GRANT SELECT ON VIEW test_schema.model_0 TO ROLE owner;
+                ON_VIRTUAL_UPDATE_END;
+                """
+            )
+        )
+        snapshot = make_snapshot(model)
+        snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+        snapshots[snapshot.snapshot_id] = snapshot
+
+    upstream = next(iter(snapshots.values()))
+    upstream.dev_version_ = "preview"
+    # Representative snapshots use the production table even when they are not deployable.
+    deployability_index = DeployabilityIndex(
+        indexed_ids=frozenset(), representative_shared_version_ids={upstream.snapshot_id}
+    )
+    mapping_spy = mocker.spy(snapshot_module, "to_table_mapping")
+    replacement_spy = mocker.spy(exp, "replace_tables")
+
+    for environment_name in ("first", "second"):
+        environment_naming_info = EnvironmentNamingInfo(name=environment_name)
+        view_mapping = (
+            to_view_mapping(snapshots.values(), environment_naming_info)
+            if use_view_mapping
+            else None
+        )
+        expected_table = (
+            view_mapping[upstream.name]
+            if view_mapping is not None
+            else upstream.table_name(deployability_index.is_representative(upstream))
+        )
+        mapping_spy.reset_mock()
+        replacement_spy.reset_mock()
+        adapter_mock.create_view.reset_mock()
+        adapter_mock.execute.reset_mock()
+
+        evaluator.promote(
+            list(snapshots.values()),
+            environment_naming_info,
+            snapshots=snapshots,
+            table_mapping=view_mapping,
+            deployability_index=deployability_index,
+        )
+
+        full_mapping_calls = [
+            call for call in mapping_spy.call_args_list if len(call.args[0]) == len(snapshots)
+        ]
+        assert len(full_mapping_calls) == 1
+        assert all(len(call.args[1]) <= 1 for call in replacement_spy.call_args_list)
+        assert adapter_mock.create_view.call_count == len(snapshots)
+        for call in adapter_mock.create_view.call_args_list:
+            for query in call.kwargs["view_properties"].values():
+                assert exp.table_name(query.find(exp.Table)) == expected_table
+        for call in adapter_mock.execute.call_args_list:
+            for statement in call.args[0]:
+                assert exp.table_name(statement.find(exp.Table)) == expected_table
+
+        # A later promotion must see changed versions, deployability and environment overrides.
+        upstream.version = "updated"
+        deployability_index = DeployabilityIndex.none_deployable()