Repository: TobikoData/sqlmesh
Pull request #2453 (2024-04-16 04:16:31)

fix: performance for selectors
selectors causes a lot of dag snapshot_id comparisons with upstream. pydantic is very slow with __eq__, i saw a 50 sec -> 7 sec drop.


Patch:
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index c166df20..b709c256 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -124,6 +124,16 @@ class SnapshotId(PydanticModel, frozen=True):
         """Helper method to return self."""
         return self
 
+    def __eq__(self, other: t.Any) -> bool:
+        return (
+            isinstance(other, self.__class__)
+            and self.name == other.name
+            and self.identifier == other.identifier
+        )
+
+    def __hash__(self) -> int:
+        return hash((self.__class__, self.name, self.identifier))
+
     def __lt__(self, other: SnapshotId) -> bool:
         return self.name < other.name
 
diff --git a/sqlmesh/utils/__init__.py b/sqlmesh/utils/__init__.py
index 8e6eb04e..d57a1a4c 100644
--- a/sqlmesh/utils/__init__.py
+++ b/sqlmesh/utils/__init__.py
@@ -316,22 +316,21 @@ def columns_to_types_to_struct(
     )
 
 
-def columns_to_types_all_known(columns_to_types: t.Dict[str, exp.DataType]) -> bool:
-    """
-    Checks that all column types are known and not NULL.
-    """
-
-    def check_expression_unknown(d_type: t.Union[exp.DataType, exp.ColumnDef]) -> bool:
-        if isinstance(d_type, exp.ColumnDef):
-            if not d_type.kind:
-                return False
-            d_type = d_type.kind
-        if isinstance(d_type, exp.DataTypeParam):
-            return True
-        if d_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL):
+def type_is_known(d_type: t.Union[exp.DataType, exp.ColumnDef]) -> bool:
+    """Checks that a given column type is known and not NULL."""
+    if isinstance(d_type, exp.ColumnDef):
+        if not d_type.kind:
             return False
-        if d_type.expressions:
-            return all(check_expression_unknown(expression) for expression in d_type.expressions)
+        d_type = d_type.kind
+    if isinstance(d_type, exp.DataTypeParam):
         return True
+    if d_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL):
+        return False
+    if d_type.expressions:
+        return all(type_is_known(expression) for expression in d_type.expressions)
+    return True
+
 
-    return all(check_expression_unknown(expression) for expression in columns_to_types.values())
+def columns_to_types_all_known(columns_to_types: t.Dict[str, exp.DataType]) -> bool:
+    """Checks that all column types are known and not NULL."""
+    return all(type_is_known(expression) for expression in columns_to_types.values())


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index d6443023..7583a3a0 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -18,7 +18,7 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.dialect import normalize_model_name, schema_
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model import Model, PythonModel, SqlModel
-from sqlmesh.utils import UniqueKeyDict, random_id, yaml
+from sqlmesh.utils import UniqueKeyDict, random_id, type_is_known, yaml
 from sqlmesh.utils.date import pandas_timestamp_to_pydatetime
 from sqlmesh.utils.errors import ConfigError, TestError
 
@@ -94,20 +94,39 @@ class ModelTest(unittest.TestCase):
     def setUp(self) -> None:
         """Load all input tables"""
         for name, values in self.body.get("inputs", {}).items():
-            # Types specified in the test take precedence over the corresponding inferred ones
+            all_types_are_known = False
+            known_columns_to_types: t.Dict[str, exp.DataType] = {}
+
             model = self.models.get(name)
-            columns_to_types = {
-                **(model.columns_to_types if model and model.annotated else {}),
-                **values.get("columns", {}),
-            }
+            if model:
+                inferred_columns_to_types = model.columns_to_types or {}
+                known_columns_to_types = {
+                    c: t for c, t in inferred_columns_to_types.items() if type_is_known(t)
+                }
+                all_types_are_known = bool(inferred_columns_to_types) and (
+                    len(known_columns_to_types) == len(inferred_columns_to_types)
+                )
+
+            # Types specified in the test take precedence over the corresponding inferred ones
+            known_columns_to_types.update(values.get("columns", {}))
 
             rows = values["rows"]
-            if not columns_to_types and rows:
-                for i, v in rows[0].items():
-                    v_type = annotate_types(exp.convert(v)).type or type(v).__name__
-                    columns_to_types[i] = exp.maybe_parse(
-                        v_type, into=exp.DataType, dialect=self.dialect
-                    )
+            if not all_types_are_known and rows:
+                for col, value in rows[0].items():
+                    if col not in known_columns_to_types:
+                        v_type = annotate_types(exp.convert(value)).type or type(value).__name__
+                        v_type = exp.maybe_parse(v_type, into=exp.DataType, dialect=self.dialect)
+
+                        if not type_is_known(v_type):
+                            _raise_error(
+                                f"Failed to infer the data type of column '{col}' for '{name}'. This issue can be "
+                                "mitigated by casting the column in the model definition, setting its type in "
+                                "schema.yaml if it's an external model, setting the model's 'columns' property, "
+                                "or setting its 'columns' mapping in the test itself",
+                                self.path,
+                            )
+
+                        known_columns_to_types[col] = v_type
 
             test_fixture_table = self._test_fixture_table(name)
             if test_fixture_table.db:
@@ -115,8 +134,8 @@ class ModelTest(unittest.TestCase):
                     schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
                 )
 
-            df = _create_df(rows, columns=columns_to_types)
-            self.engine_adapter.create_view(test_fixture_table, df, columns_to_types)
+            df = _create_df(rows, columns=known_columns_to_types)
+            self.engine_adapter.create_view(test_fixture_table, df, known_columns_to_types)
 
     def tearDown(self) -> None:
         """Drop all fixture tables."""
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 216a27f3..1efb1a07 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -548,6 +548,84 @@ test_foo:
     )
 
 
+def test_partially_inferred_schemas(sushi_context: Context, mocker: MockerFixture) -> None:
+    parent = _create_model(
+        "SELECT a, b, s::ROW(d DATE) AS s FROM sushi.unknown",
+        meta="MODEL (name sushi.parent, kind FULL, dialect trino)",
+        default_catalog="memory",
+    )
+    parent = t.cast(SqlModel, sushi_context.upsert_model(parent))
+
+    child = _create_model(
+        "SELECT a, b, s.d AS d FROM sushi.parent",
+        meta="MODEL (name sushi.child, kind FULL)",
+        default_catalog="memory",
+    )
+    child = t.cast(SqlModel, sushi_context.upsert_model(child))
+
+    body = load_yaml(
+        """
+test_child:
+  model: sushi.child
+  inputs:
+    sushi.parent:
+      - a: 1
+        b: bla
+        s:
+          d: 2020-01-01
+  outputs:
+    query:
+      - a: 1
+        b: bla
+        d: 2020-01-01
+        """
+    )
+    test = _create_test(body, "test_child", child, sushi_context)
+
+    random_id = "jzngz56a"
+    test._test_id = random_id
+
+    spy_execute = mocker.spy(test.engine_adapter, "_execute")
+    _check_successful_or_raise(test.run())
+
+    spy_execute.assert_any_call(
+        f'CREATE OR REPLACE VIEW "memory"."sushi"."parent__fixture__{random_id}" ("s", "a", "b") AS '
+        "SELECT "
+        'CAST("s" AS STRUCT("d" DATE)) AS "s", '
+        'CAST("a" AS INT) AS "a", '
+        'CAST("b" AS TEXT) AS "b" '
+        """FROM (VALUES ({'d': CAST('2020-01-01' AS DATE)}, 1, 'bla')) AS "t"("s", "a", "b")"""
+    )
+
+
+def test_uninferrable_schema() -> None:
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      - value: null
+  outputs:
+    query:
+      - value: null
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("SELECT value FROM raw"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            """Failed to infer the data type of column 'value' for '"raw"'. This issue can """
+            "be mitigated by casting the column in the model definition, setting its type in "
+            "schema.yaml if it's an external model, setting the model's 'columns' property, "
+            "or setting its 'columns' mapping in the test itself\n"
+        ),
+    )
+
+
 def test_missing_column_failure(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
     _check_successful_or_raise(
         _create_test(
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index 571d5aba..4bccc9b8 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -167,7 +167,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
         "models_to_backfill": ['"test_model"'],
         "end_bounded": False,
         "ensure_finalized_snapshots": False,
-        "directly_modified_snapshots": [snapshot.snapshot_id],
+        "directly_modified_snapshots": [{"identifier": "1736759018", "name": '"test_model"'}],
         "indirectly_modified_snapshots": {},
         "removed_snapshots": [],
     }