fix: include every branch's columns in a BY NAME set operation's columns_to_types

sqlmesh/core/model/definition.py

@@ -1490,7 +1490,7 @@ def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
             unknown = exp.DataType.build("unknown")
 
             columns_to_types = {}
-            for select in query.selects:
+            for select in _output_projections(query):
                 output_name = select.output_name
 
                 # If model validation is disabled, we cannot assume that projections
@@ -1520,7 +1520,7 @@ def column_descriptions(self) -> t.Dict[str, str]:
 
         return {
             select.alias_or_name: select.comments[-1].strip()
-            for select in query.selects
+            for select in _output_projections(query)
             if select.comments
         }
 
@@ -2895,6 +2895,74 @@ def _list_of_calls_to_exp(value: t.List[t.Tuple[str, t.Dict[str, t.Any]]]) -> ex
     )
 
 
+def _output_projections(query: exp.Query) -> t.List[exp.Expr]:
+    """Return the projections that define a query's output columns.
+
+    A set operation takes its columns from its leftmost branch, unless it matches its branches'
+    columns BY NAME. Then a column that only a later branch has is also an output column (engines
+    such as DuckDB fill it with NULL for the other branches), so `query.selects`, which only covers
+    the leftmost branch, misses it. Each column is represented by the projection of the first
+    branch that has it, which carries its type and comments.
+
+    Like `query.selects`, this walks down a chain of set operations' left branches in a loop, so
+    long chains don't recurse; only a parenthesized right branch is visited recursively.
+    """
+    set_operations: t.List[exp.SetOperation] = []
+    expression = query.unnest()
+    while isinstance(expression, exp.SetOperation):
+        set_operations.append(expression)
+        expression = expression.left.unnest()
+
+    projections = expression.selects if isinstance(expression, exp.Query) else []
+    for set_operation in reversed(set_operations):
+        if set_operation.args.get("by_name"):
+            projections = _by_name_projections(
+                set_operation, projections, _output_projections(set_operation.right)
+            )
+    return projections
+
+
+def _by_name_projections(
+    set_operation: exp.SetOperation, left: t.List[exp.Expr], right: t.List[exp.Expr]
+) -> t.List[exp.Expr]:
+    """Combine the projections of a BY NAME set operation's branches into its output's.
+
+    For BigQuery's modifiers, this follows sqlglot's `Resolver.get_source_columns_from_set_op`: an
+    ON list names the columns, LEFT keeps the left branch's and INNER the ones both branches have.
+    With FULL or no modifier, as in DuckDB's `UNION BY NAME`, the output has every branch's columns
+    in order of first appearance. sqlglot 30.8's resolver takes the left branch's columns instead
+    in that case, which is the bug this works around.
+    """
+    on = set_operation.args.get("on")
+    if not on and set_operation.side == "LEFT":
+        return left
+
+    # An unexpanded star hides some of a branch's columns, so the output's columns can't be known.
+    # Keeping it among the projections makes `columns_to_types` return None.
+    stars = [projection for projection in (*left, *right) if projection.is_star]
+    if stars:
+        return [*left, *stars]
+
+    if on:
+        projections: t.Dict[str, exp.Expr] = {}
+        for projection in (*left, *right):
+            projections.setdefault(projection.output_name, projection)
+        return [projections[column.name] for column in on if column.name in projections]
+
+    if set_operation.kind == "INNER":
+        right_names = {projection.output_name for projection in right}
+        return [projection for projection in left if projection.output_name in right_names]
+
+    # Keep the left branch's projections as they are, so that a duplicate or missing name there
+    # still makes the columns unknown, and add the columns that only the right branch has
+    left_names = {projection.output_name for projection in left}
+    added: t.Dict[str, exp.Expr] = {}
+    for projection in right:
+        if projection.output_name not in left_names:
+            added.setdefault(projection.output_name, projection)
+    return [*left, *added.values()]
+
+
 def _has_ordinal_references(query: exp.Query) -> bool:
     order = query.args.get("order")
     if order and any(
@@ -3030,6 +3098,12 @@ def _is_only_projection_additions(
                     and isinstance(this_expression, exp.Select)
                     and arg_key == "expressions"
                 ):
+                    # In a branch of a set operation BY NAME, an added projection can fill in a
+                    # column that another branch already outputs, changing its values
+                    if len(previous_value) != len(this_value) and _is_by_name_branch(
+                        this_expression
+                    ):
+                        return False
                     if not _projections_only_safely_added(previous_expression, this_expression):
                         return False
                 elif len(previous_value) != len(this_value):
@@ -3049,6 +3123,16 @@ def _is_only_projection_additions(
     return True
 
 
+def _is_by_name_branch(query: exp.Select) -> bool:
+    """Return whether the SELECT is a branch, at any depth, of a set operation BY NAME."""
+    parent = query.parent
+    while isinstance(parent, (exp.SetOperation, exp.Subquery)):
+        if isinstance(parent, exp.SetOperation) and parent.args.get("by_name"):
+            return True
+        parent = parent.parent
+    return False
+
+
 def _single_expr_or_tuple(values: t.Sequence[exp.Expr]) -> exp.Expr | exp.Tuple:
     return values[0] if len(values) == 1 else exp.Tuple(expressions=values)
 

tests/core/test_context.py

@@ -3965,6 +3965,27 @@ def test_table_diff_ignores_extra_args(sushi_context: Context):
     )
 
 
+def test_plan_union_by_name_keeps_later_branch_columns(tmp_path: Path):
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+    (models_dir / "x.sql").write_text("MODEL (name t.x, kind FULL); SELECT 1::INT AS a")
+    (models_dir / "y.sql").write_text(
+        "MODEL (name t.y, kind FULL); SELECT 2::INT AS a, 'two'::TEXT AS b"
+    )
+    (models_dir / "u.sql").write_text(
+        "MODEL (name t.u, kind FULL); SELECT a FROM t.x UNION ALL BY NAME SELECT a, b FROM t.y"
+    )
+
+    context = Context(
+        paths=tmp_path, config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
+    )
+    context.plan(auto_apply=True, no_prompts=True)
+
+    df = context.fetchdf("SELECT * FROM t.u ORDER BY a")
+    assert list(df.columns) == ["a", "b"]
+    assert df["b"].tolist() == [None, "two"]
+
+
 def test_plan_min_intervals(tmp_path: Path):
     init_example_project(tmp_path, engine_type="duckdb", dialect="duckdb")
 

tests/core/test_model.py

@@ -309,6 +309,157 @@ def test_model_union_query(sushi_context, assert_exp_eq):
     )
 
 
+def test_columns_to_types_union_by_name():
+    # BY NAME matches the branches' columns by name, so a column that only a later branch
+    # has is still one of the query's columns, after the columns of the branches before it
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect duckdb,
+        );
+
+        SELECT 1::INT AS a
+        UNION ALL BY NAME
+        SELECT
+          'x'::TEXT AS b, -- only in the second branch
+          2::INT AS a
+        UNION ALL BY NAME
+        SELECT 3::BIGINT AS c
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert list(model.columns_to_types or {}) == ["a", "b", "c"]
+    assert model.columns_to_types == {
+        "a": exp.DataType.build("int"),
+        "b": exp.DataType.build("text"),
+        "c": exp.DataType.build("bigint"),
+    }
+    assert model.column_descriptions == {"b": "only in the second branch"}
+
+
+def test_columns_to_types_union_by_position_of_union_by_name():
+    # Without BY NAME, the leftmost branch names the columns, even when that branch is
+    # itself a union BY NAME
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect duckdb,
+        );
+
+        (SELECT 1::INT AS a UNION ALL BY NAME SELECT 2::INT AS a, 'x'::TEXT AS b)
+        UNION ALL
+        SELECT 3::INT AS c, 'y'::TEXT AS d
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert model.columns_to_types == {
+        "a": exp.DataType.build("int"),
+        "b": exp.DataType.build("text"),
+    }
+
+
+@pytest.mark.parametrize(
+    "set_operation, expected_columns",
+    [
+        ("FULL UNION ALL BY NAME", ["a", "z", "b"]),
+        ("LEFT UNION ALL BY NAME", ["a", "z"]),
+        ("INNER UNION ALL BY NAME", ["a"]),
+        ("FULL UNION ALL BY NAME ON (z, b)", ["z", "b"]),
+    ],
+)
+def test_columns_to_types_union_by_name_modifiers(
+    set_operation: str, expected_columns: t.List[str]
+):
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect bigquery,
+        );
+
+        SELECT 1 AS a, 0 AS z
+        {set_operation}
+        SELECT 2 AS a, 3 AS b
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert list(model.columns_to_types or {}) == expected_columns
+
+
+def test_columns_to_types_union_by_name_parenthesized_right_branch():
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect duckdb,
+        );
+
+        SELECT 1::INT AS a
+        UNION ALL BY NAME
+        (SELECT 2::INT AS a, 'x'::TEXT AS b UNION ALL BY NAME SELECT 'y'::TEXT AS c)
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert list(model.columns_to_types or {}) == ["a", "b", "c"]
+
+
+@pytest.mark.parametrize(
+    "set_operation",
+    [
+        "UNION ALL BY NAME",
+        "INNER UNION ALL BY NAME",
+        "FULL UNION ALL BY NAME ON (a, b)",
+    ],
+)
+def test_columns_to_types_union_by_name_unexpanded_star(set_operation: str):
+    # A star that can't be expanded hides some of a branch's columns, so the columns are unknown
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect bigquery,
+        );
+
+        SELECT * FROM db.unknown_source
+        {set_operation}
+        SELECT 1 AS a
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert model.columns_to_types is None
+
+
+@pytest.mark.parametrize("set_operation", ["UNION ALL", "UNION ALL BY NAME"])
+def test_columns_to_types_long_union_chain(set_operation: str):
+    branches = f" {set_operation} ".join(f"SELECT {i}::INT AS a" for i in range(1000))
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            dialect duckdb,
+        );
+
+        {branches}
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert model.columns_to_types == {"a": exp.DataType.build("int")}
+
+
 @time_machine.travel("1996-02-10 00:00:00 UTC")
 @pytest.mark.parametrize(
     "test_id, condition, union_type, table_count, expected_result",

tests/core/test_snapshot.py

@@ -1882,6 +1882,18 @@ def test_categorize_change_sql_additive_projection_edge_cases(make_snapshot):
             SnapshotChangeCategory.NON_BREAKING,
             id="union-order-by-ordinal-append",
         ),
+        pytest.param(
+            "SELECT a FROM x UNION ALL BY NAME SELECT a, b FROM y",
+            "SELECT a, b FROM x UNION ALL BY NAME SELECT a, b FROM y",
+            None,
+            id="union-by-name-fills-existing-column",
+        ),
+        pytest.param(
+            "SELECT a FROM x UNION ALL BY NAME SELECT a FROM y",
+            "SELECT a, b FROM x UNION ALL BY NAME SELECT a FROM y",
+            None,
+            id="union-by-name-addition-not-auto-categorized",
+        ),
     ],
 )
 def test_categorize_change_sql_nested_projection_additions(