perf: speed up SQL auto-categorization

Signed-off-by: dirk.nilius <dirk.nilius@nc-group.net>

Author: dirk.nilius

sqlmesh/core/model/definition.py

@@ -9,8 +9,7 @@
 from pathlib import Path
 
 from pydantic import Field
-from sqlglot import diff, exp
-from sqlglot.diff import Insert
+from sqlglot import exp
 from sqlglot.helper import seq_get
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.optimizer.simplify import gen
@@ -1580,37 +1579,12 @@ def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
             # Can't determine if there's a breaking change if we can't render the query.
             return None
 
-        if previous_query is this_query:
-            edits = []
-        else:
-            edits = diff(
-                previous_query,
-                this_query,
-                matchings=[(previous_query, this_query)],
-                delta_only=True,
-                dialect=self.dialect if self.dialect == previous.dialect else None,
-            )
-        inserted_expressions = {e.expression for e in edits if isinstance(e, Insert)}
-
-        for edit in edits:
-            if not isinstance(edit, Insert):
-                return _additive_projection_change(previous_query, this_query, self.dialect)
-
-            expr = edit.expression
-            if isinstance(expr, exp.UDTF):
-                # projection subqueries do not change cardinality, engines don't allow these to return
-                # more than one row of data
-                parent = expr.find_ancestor(exp.Subquery)
-
-                if not parent:
-                    return None
-
-                expr = parent
-
-            if not _is_projection(expr) and expr.parent not in inserted_expressions:
-                return _additive_projection_change(previous_query, this_query, self.dialect)
+        if previous_query is this_query or _is_only_projection_additions(
+            previous_query, this_query
+        ):
+            return False
 
-        return False
+        return None
 
     def is_metadata_only_change(self, previous: _Node) -> bool:
         if self._is_metadata_only_change_cache.get(id(previous), None) is not None:
@@ -2907,11 +2881,6 @@ def _list_of_calls_to_exp(value: t.List[t.Tuple[str, t.Dict[str, t.Any]]]) -> ex
     )
 
 
-def _is_projection(expr: exp.Expr) -> bool:
-    parent = expr.parent
-    return isinstance(parent, exp.Select) and expr.arg_key == "expressions"
-
-
 def _has_ordinal_references(query: exp.Select) -> bool:
     order = query.args.get("order")
     if order and any(
@@ -2924,84 +2893,134 @@ def _has_ordinal_references(query: exp.Select) -> bool:
     )
 
 
-def _additive_projection_change(
-    previous_query: exp.Query,
-    this_query: exp.Query,
-    dialect: DialectType,
-) -> t.Optional[bool]:
-    """Fallback for when SQLGlot's tree diff can't express an additive projection change.
-
-    SQLGlot's diff matches nodes by structural similarity, so interchangeable leaves (e.g. two
-    identical ``CAST(... AS T)`` target types) can be cross-matched. Inserting a same-type cast
-    above an existing one therefore yields spurious ``Move`` / ``Update`` edits even though a
-    column was simply added to the SELECT list. In that case the edit-based check above is
-    inconclusive, so we verify additivity directly against the output projections.
-
-    Returns ``False`` (non-breaking) only when the change is provably additive:
-      * both queries are simple ``SELECT`` statements,
-      * everything other than the projection list is structurally identical,
-      * no added projection is a (potentially cardinality-changing) ``UDTF``,
-      * every previous projection is preserved, in order, within the new projection list, and
-      * no mid-list insert shifts ordinal ``ORDER BY`` / ``GROUP BY`` references.
-
-    Otherwise returns ``None`` (undetermined), preserving the conservative default.
+def _is_valid_added_projection(projection: exp.Expr) -> bool:
+    """Return whether an added projection preserves the query's row cardinality.
+
+    A directly projected UDTF can emit multiple rows. SQLMesh treats it as safe when its nearest
+    subquery ancestor is contained by the added projection because engines require that projection
+    subquery to return at most one row.
     """
-    # UNIONs or other query expressions, are left to the caller's conservative diff result.
-    if not isinstance(previous_query, exp.Select) or not isinstance(this_query, exp.Select):
-        return None
+    udtfs = list(projection.find_all(exp.UDTF))
+    if not udtfs:
+        return True
 
+    projection_node_ids = {id(node) for node in projection.walk()}
+    return all(
+        (subquery := udtf.find_ancestor(exp.Subquery)) is not None
+        and id(subquery) in projection_node_ids
+        for udtf in udtfs
+    )
+
+
+def _projection_lists_match(previous_query: exp.Select, this_query: exp.Select) -> bool:
+    """Return whether a SELECT's projections differ only through safe additions.
+
+    Every previous projection must occur unchanged and in the same order in the current list.
+    Unmatched current projections are additions, subject to the UDTF cardinality check. Additions
+    before an existing projection are unsafe when the query uses ordinal GROUP BY or ORDER BY
+    references because they can change which output those ordinals address.
+    """
     previous_projections = previous_query.expressions
     this_projections = this_query.expressions
-    # If the new query has not gained any projections, this cannot be an additive projection-only
-    # change, so there is nothing for this fallback to prove.
-    if len(this_projections) <= len(previous_projections):
-        return None
+    this_index = 0
+    added_at: list[int] = []
+    matched_at: list[int] = []
 
-    # Adding a UDTF projection (e.g. EXPLODE / UNNEST) can change row cardinality, so such a
-    # change is not safely non-breaking even when it appears as an extra SELECT item.
-    for projection in this_projections:
-        bare = projection.this if isinstance(projection, exp.Alias) else projection
-        if isinstance(bare, exp.UDTF):
-            return None
+    # Match each previous projection to the earliest identical current projection. Any current
+    # projections skipped along the way are additions.
+    for previous_projection in previous_projections:
+        while (
+            this_index < len(this_projections)
+            and previous_projection != this_projections[this_index]
+        ):
+            if not _is_valid_added_projection(this_projections[this_index]):
+                return False
 
-    # Everything other than the projection list must be structurally identical. Replacing each
-    # SELECT list with the same dummy literal lets the expression equality check focus on the
-    # FROM / WHERE / GROUP BY / ORDER BY / etc. parts of the query.
-    previous_skeleton = previous_query.copy()
-    this_skeleton = this_query.copy()
-    previous_skeleton.set("expressions", [exp.Literal.number(1)])
-    this_skeleton.set("expressions", [exp.Literal.number(1)])
-    if previous_skeleton != this_skeleton:
-        return None
+            added_at.append(this_index)
+            this_index += 1
 
-    # Every previous projection must appear, in order, within the new projection list. Comparing
-    # dialect-normalized SQL makes semantically equivalent projection nodes match even when the
-    # parser built distinct object identities.
-    this_projection_sql = [p.sql(dialect=dialect, comments=False) for p in this_projections]
-    search_start = 0
-    matched_at: list[int] = []
-    for projection in previous_projections:
-        target_sql = projection.sql(dialect=dialect, comments=False)
-        # Continue after the previous match so added columns can appear before, between, or after
-        # the original projections, but existing projections cannot be reordered or rewritten.
-        for index in range(search_start, len(this_projection_sql)):
-            if this_projection_sql[index] == target_sql:
-                matched_at.append(index)
-                search_start = index + 1
-                break
-        else:
-            return None
+        if this_index == len(this_projections):
+            return False
+
+        matched_at.append(this_index)
+        this_index += 1
+
+    # Once all previous projections are matched, every remaining projection was appended.
+    for index in range(this_index, len(this_projections)):
+        if not _is_valid_added_projection(this_projections[index]):
+            return False
+        added_at.append(index)
+
+    # Be conservative about every mid-list addition when ordinals are present. Determining whether
+    # a particular ordinal was shifted would couple this comparison to dialect-specific semantics.
+    if (
+        added_at
+        and matched_at
+        and _has_ordinal_references(this_query)
+        and any(index < matched_at[-1] for index in added_at)
+    ):
+        return False
+
+    return True
 
-    # Mid-list inserts shift ordinal references in ORDER BY / GROUP BY clauses.
-    if _has_ordinal_references(this_query):
-        matched_set = set(matched_at)
-        last_matched = matched_at[-1]
-        if any(i < last_matched for i in range(len(this_projections)) if i not in matched_set):
-            return None
 
-    # At this point the query shape is unchanged and all prior outputs are preserved, so the only
-    # remaining difference is one or more additional, non-UDTF projections.
-    return False
+def _is_only_projection_additions(
+    previous_query: exp.Query,
+    this_query: exp.Query,
+) -> bool:
+    """Return whether a query changed exclusively through safe projection additions.
+
+    The two ASTs are walked in lockstep. Node types, scalar arguments, and non-projection child
+    lists must match exactly. SELECT projection lists may contain additional expressions as long
+    as all previous projections remain unchanged and ordered and the additions pass the
+    cardinality and ordinal-reference safeguards.
+
+    This specialized comparison avoids the candidate matching performed by SQLGlot's general tree
+    diff while remaining conservative for every change other than an added projection.
+    """
+    expression_pairs: t.List[t.Tuple[exp.Expr, exp.Expr]] = [(previous_query, this_query)]
+
+    while expression_pairs:
+        previous_expression, this_expression = expression_pairs.pop()
+
+        if type(previous_expression) is not type(this_expression):
+            return False
+
+        for arg_key in previous_expression.args.keys() | this_expression.args.keys():
+            previous_value = previous_expression.args.get(arg_key)
+            this_value = this_expression.args.get(arg_key)
+
+            if isinstance(previous_value, exp.Expr):
+                if not isinstance(this_value, exp.Expr):
+                    return False
+
+                expression_pairs.append((previous_value, this_value))
+            elif isinstance(previous_value, list):
+                if not isinstance(this_value, list):
+                    return False
+
+                if (
+                    isinstance(previous_expression, exp.Select)
+                    and isinstance(this_expression, exp.Select)
+                    and arg_key == "expressions"
+                ):
+                    if not _projection_lists_match(previous_expression, this_expression):
+                        return False
+                elif len(previous_value) != len(this_value):
+                    return False
+                else:
+                    for previous_item, this_item in zip(previous_value, this_value):
+                        if isinstance(previous_item, exp.Expr):
+                            if not isinstance(this_item, exp.Expr):
+                                return False
+
+                            expression_pairs.append((previous_item, this_item))
+                        elif previous_item != this_item:
+                            return False
+            elif previous_value != this_value:
+                return False
+
+    return True
 
 
 def _single_expr_or_tuple(values: t.Sequence[exp.Expr]) -> exp.Expr | exp.Tuple:

tests/core/test_snapshot.py

@@ -1633,10 +1633,7 @@ def test_categorize_change_sql(make_snapshot):
     )
 
 
-def test_categorize_change_sql_redundant_cast(make_snapshot):
-    # Adding a column with a redundant CAST above an existing same-type cast makes SQLGlot's tree
-    # diff emit spurious Move/Update edits (interchangeable DataType leaves get cross-matched),
-    # even though the change is purely an added projection. These must remain NON_BREAKING.
+def test_categorize_change_sql_additive_projection_edge_cases(make_snapshot):
     config = CategorizerConfig(sql=AutoCategorizationMode.SEMI)
 
     old_snapshot = make_snapshot(
@@ -1655,6 +1652,18 @@ def test_categorize_change_sql_redundant_cast(make_snapshot):
         == SnapshotChangeCategory.NON_BREAKING
     )
 
+    # An added projection identical to an existing projection remains non-breaking.
+    assert (
+        categorize_change(
+            new=make_snapshot(
+                SqlModel(name="a", query=parse_one("SELECT s::TEXT, a::DATE, s::TEXT FROM t"))
+            ),
+            old=old_snapshot,
+            config=config,
+        )
+        == SnapshotChangeCategory.NON_BREAKING
+    )
+
     # Multiple same-type cast columns have been inserted mid-list.
     assert (
         categorize_change(
@@ -1771,7 +1780,7 @@ def test_categorize_change_sql_redundant_cast(make_snapshot):
         is None
     )
 
-    # Aliased UDTF projection via fallback path: undetermined.
+    # Aliased UDTF projection: undetermined.
     assert (
         categorize_change(
             new=make_snapshot(
@@ -1810,6 +1819,66 @@ def test_categorize_change_sql_redundant_cast(make_snapshot):
     )
 
 
+@pytest.mark.parametrize(
+    ("previous_query", "this_query", "expected"),
+    [
+        pytest.param(
+            "WITH x AS (SELECT a FROM t) SELECT a FROM x",
+            "WITH x AS (SELECT a, b FROM t) SELECT a FROM x",
+            SnapshotChangeCategory.NON_BREAKING,
+            id="cte",
+        ),
+        pytest.param(
+            "SELECT a FROM t WHERE EXISTS (SELECT x FROM u)",
+            "SELECT a FROM t WHERE EXISTS (SELECT x, y FROM u)",
+            SnapshotChangeCategory.NON_BREAKING,
+            id="exists",
+        ),
+        pytest.param(
+            "SELECT (SELECT a FROM t) AS x",
+            "SELECT (SELECT a, b FROM t) AS x",
+            None,
+            id="scalar-subquery-projection",
+        ),
+        pytest.param(
+            "SELECT a FROM x UNION ALL SELECT a FROM y",
+            "SELECT a, b FROM x UNION ALL SELECT a, b FROM y",
+            SnapshotChangeCategory.NON_BREAKING,
+            id="union",
+        ),
+        pytest.param(
+            "WITH x AS (SELECT a::DATE, s::TEXT FROM t ORDER BY 2) SELECT * FROM x",
+            "WITH x AS (SELECT a::DATE, x::TEXT, s::TEXT FROM t ORDER BY 2) SELECT * FROM x",
+            None,
+            id="nested-order-by-ordinal",
+        ),
+        pytest.param(
+            "SELECT a FROM (SELECT x FROM t) AS nested",
+            "SELECT a FROM (SELECT x, EXPLODE(y) FROM t) AS nested",
+            None,
+            id="derived-table-udtf",
+        ),
+    ],
+)
+def test_categorize_change_sql_nested_projection_additions(
+    make_snapshot,
+    previous_query,
+    this_query,
+    expected,
+):
+    config = CategorizerConfig(sql=AutoCategorizationMode.SEMI)
+    old_snapshot = make_snapshot(SqlModel(name="a", query=parse_one(previous_query)))
+
+    assert (
+        categorize_change(
+            new=make_snapshot(SqlModel(name="a", query=parse_one(this_query))),
+            old=old_snapshot,
+            config=config,
+        )
+        is expected
+    )
+
+
 def test_categorize_change_seed(make_snapshot, tmp_path):
     config = CategorizerConfig(seed=AutoCategorizationMode.SEMI)
     model_name = "test_db.test_seed_model"