Feat: Cache the optimized model query (#975)

* Feat: Cache the optimized model query

* fix failing tests

* disable gc when loading

* add optimized query test

---------

Co-authored-by: tobymao <toby.mao@gmail.com>

Author: izeigerman

setup.py

@@ -44,7 +44,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=16.0.0",
+        "sqlglot~=16.1.0",
         "fsspec",
     ],
     extras_require={

sqlmesh/core/context.py

@@ -34,6 +34,7 @@
 
 import abc
 import contextlib
+import gc
 import typing as t
 import unittest.result
 from io import StringIO
@@ -291,7 +292,7 @@ def upsert_model(self, model: t.Union[str, Model], **kwargs: t.Any) -> Model:
 
         self._models.update({model.name: model})
         self.dag.add(model.name, model.depends_on)
-        update_model_schemas(self.dag, self._models)
+        update_model_schemas(self.dag, self._models, self.path)
 
         return model
 
@@ -356,11 +357,13 @@ def refresh(self) -> None:
     def load(self, update_schemas: bool = True) -> Context:
         """Load all files in the context's path."""
         with sys_path(*self.configs):
+            gc.disable()
             project = self._loader.load(self, update_schemas)
             self._macros = project.macros
             self._models = project.models
             self._audits = project.audits
             self.dag = project.dag
+            gc.enable()
 
         return self
 

sqlmesh/core/loader.py

@@ -22,6 +22,7 @@
 from sqlmesh.core.model import (
     Model,
     ModelCache,
+    OptimizedQueryCache,
     SeedModel,
     create_external_model,
     load_model,
@@ -38,8 +39,11 @@
     from sqlmesh.core.context import Context
 
 
-def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> None:
+def update_model_schemas(
+    dag: DAG[str], models: UniqueKeyDict[str, Model], context_path: Path
+) -> None:
     schema = MappingSchema(normalize=False)
+    optimized_query_cache: OptimizedQueryCache = OptimizedQueryCache(context_path / c.CACHE)
 
     for name in dag.sorted:
         model = models.get(name)
@@ -49,18 +53,18 @@ def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> No
             continue
 
         model.update_schema(schema)
+        cache_hit = optimized_query_cache.with_optimized_query(model)
         schema.add_table(name, model.columns_to_types, dialect=model.dialect)
 
-        external = any(dep not in models for dep in model.depends_on)
-        if external:
+        if any(dep not in models for dep in model.depends_on):
             if "*" in model.columns_to_types:
                 raise ConfigError(
                     f"Can't expand SELECT * expression for model '{name}' at '{model._path}'."
                     " Either specify external source projections expliticly or"
                     ' add source tables as "external models" using the command'
                     " 'sqlmesh create_external_models'."
                 )
-        elif model.mapping_schema:
+        elif model.mapping_schema and not cache_hit:
             try:
                 validate_qualify_columns(model.render_query())
             except SqlglotError as e:
@@ -120,7 +124,7 @@ def load(self, context: Context, update_schemas: bool = True) -> LoadedProject:
             self._add_model_to_dag(model)
 
         if update_schemas:
-            update_model_schemas(self._dag, models)
+            update_model_schemas(self._dag, models, self._context.path)
 
         audits = self._load_audits()
 

sqlmesh/core/model/__init__.py

@@ -1,4 +1,4 @@
-from sqlmesh.core.model.cache import ModelCache
+from sqlmesh.core.model.cache import ModelCache, OptimizedQueryCache
 from sqlmesh.core.model.common import parse_model_name
 from sqlmesh.core.model.decorator import model
 from sqlmesh.core.model.definition import (

sqlmesh/core/model/cache.py

@@ -7,6 +7,7 @@
 
 from sqlmesh.core.model.definition import Model, SqlModel
 from sqlmesh.utils.cache import FileCache
+from sqlmesh.utils.hashing import crc32
 from sqlmesh.utils.pydantic import PydanticModel
 
 
@@ -44,14 +45,73 @@ def get_or_load(self, name: str, entry_id: str, loader: t.Callable[[], Model]) -
         cache_entry = self._file_cache.get(name, entry_id)
         if cache_entry:
             model = cache_entry.model
-            model._query_renderer.update_cache(cache_entry.rendered_query)
+            model._query_renderer.update_cache(cache_entry.rendered_query, optimized=False)
             return model
 
         loaded_model = loader()
         if isinstance(loaded_model, SqlModel):
             new_entry = SqlModelCacheEntry(
-                model=loaded_model, rendered_query=loaded_model.render_query()
+                model=loaded_model, rendered_query=loaded_model.render_query(optimize=False)
             )
             self._file_cache.put(name, entry_id, new_entry)
 
         return loaded_model
+
+
+class OptimizedQueryCacheEntry(PydanticModel):
+    optimized_rendered_query: exp.Expression
+
+
+class OptimizedQueryCache:
+    """File-based cache implementation for optimized model queries.
+
+    Args:
+        path: The path to the cache folder.
+    """
+
+    def __init__(self, path: Path):
+        self.path = path
+        self._file_cache: FileCache[OptimizedQueryCacheEntry] = FileCache(
+            path, OptimizedQueryCacheEntry, prefix="optimized_query"
+        )
+
+    def with_optimized_query(self, model: Model) -> bool:
+        """Adds an optimized query to the model's in-memory cache.
+
+        Args:
+            model: The model to add the optimized query to.
+        """
+        if not isinstance(model, SqlModel):
+            return True
+
+        entry_id = self._entry_id(model)
+        cache_entry = self._file_cache.get(model.name, entry_id)
+        if cache_entry:
+            model._query_renderer.update_cache(cache_entry.optimized_rendered_query, optimized=True)
+            return True
+
+        new_entry = OptimizedQueryCacheEntry(
+            optimized_rendered_query=model.render_query(optimize=True)
+        )
+        self._file_cache.put(model.name, entry_id, new_entry)
+        return False
+
+    @staticmethod
+    def _entry_id(model: SqlModel) -> str:
+        data = OptimizedQueryCache._mapping_schema_hash_data(model.mapping_schema)
+        data.append(model.render_query(optimize=False).sql())
+        return crc32(data)
+
+    @staticmethod
+    def _mapping_schema_hash_data(schema: t.Dict[str, t.Any]) -> t.List[str]:
+        keys = sorted(schema) if all(isinstance(v, dict) for v in schema.values()) else schema
+
+        data = []
+        for k in keys:
+            data.append(k)
+            if isinstance(schema[k], dict):
+                data.extend(OptimizedQueryCache._mapping_schema_hash_data(schema[k]))
+            else:
+                data.append(str(schema[k]))
+
+        return data

sqlmesh/core/model/definition.py

@@ -433,7 +433,7 @@ def depends_on(self) -> t.Set[str]:
             return self.depends_on_ - {self.name}
 
         if self._depends_on is None:
-            self._depends_on = _find_tables(self._render_all_sql()) - {self.name}
+            self._depends_on = _find_tables(self.render_query(optimize=False)) - {self.name}
         return self._depends_on
 
     @property
@@ -481,7 +481,7 @@ def depends_on_past(self) -> bool:
         if self._depends_on_past is None:
             self._depends_on_past = (
                 self.kind.is_incremental_by_unique_key
-                or self.name in _find_tables([self.render_query()])
+                or self.name in _find_tables(self.render_query(optimize=False))
             )
         return self._depends_on_past
 
@@ -529,14 +529,6 @@ def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
         """
         raise NotImplementedError
 
-    def _render_all_sql(self) -> t.List[exp.Expression]:
-        """Renders all the SQL expressions of this model."""
-        return [
-            *self.render_pre_statements(),
-            self.render_query(),
-            *self.render_post_statements(),
-        ]
-
 
 class _SqlBasedModel(_Model):
     pre_statements_: t.Optional[t.List[exp.Expression]] = Field(
@@ -713,17 +705,18 @@ def column_descriptions(self) -> t.Dict[str, str]:
         if self._column_descriptions is None:
             self._column_descriptions = {
                 select.alias: "\n".join(comment.strip() for comment in select.comments)
-                for select in self.render_query().selects
+                for select in self.render_query(optimize=False).selects
                 if select.comments
             }
         return self._column_descriptions
 
     def update_schema(self, schema: MappingSchema) -> None:
         super().update_schema(schema)
         self._columns_to_types = None
+        self._query_renderer._optimized_cache = {}
 
     def validate_definition(self) -> None:
-        query = self._query_renderer.render()
+        query = self._query_renderer.render(optimize=False)
 
         if not isinstance(query, exp.Subqueryable):
             raise_config_error("Missing SELECT query in the model definition", self._path)
@@ -1444,7 +1437,7 @@ def _validate_model_fields(klass: t.Type[_Model], provided_fields: t.Set[str], p
         raise_config_error(f"Invalid extra fields {extra_fields} in the model definition", path)
 
 
-def _find_tables(expressions: t.List[exp.Expression]) -> t.Set[str]:
+def _find_tables(expression: exp.Expression) -> t.Set[str]:
     """Find all tables referenced in a query.
 
     Args:
@@ -1455,8 +1448,7 @@ def _find_tables(expressions: t.List[exp.Expression]) -> t.Set[str]:
     """
     return {
         exp.table_name(table)
-        for e in expressions
-        for scope in traverse_scope(e)
+        for scope in traverse_scope(expression)
         for table in scope.tables
         if isinstance(table.this, exp.Identifier) and exp.table_name(table) not in scope.cte_sources
     }