diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py index e8e122dece..85afc8d44c 100644 --- a/sqlmesh/core/model/definition.py +++ b/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) diff --git a/tests/core/test_context.py b/tests/core/test_context.py index 71276b7396..61cc7a6f00 100644 --- a/tests/core/test_context.py +++ b/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") diff --git a/tests/core/test_model.py b/tests/core/test_model.py index c5d1a71bee..00a04f9f21 100644 --- a/tests/core/test_model.py +++ b/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", diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py index 64bee7f472..77f1d5b225 100644 --- a/tests/core/test_snapshot.py +++ b/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(