Pair the branches of a set operation by position

crates/flowscope-core/src/analyzer/context.rs

@@ -154,6 +154,13 @@ pub(crate) struct StatementContext {
     /// How many subqueries a predicate reads the statement has opened, each
     /// analysed against a node of its own.
     pub(crate) predicate_subqueries: usize,
+    /// While a later branch of a set operation is read, the names its columns
+    /// take. SQL pairs the branches of a UNION, an INTERSECT or an EXCEPT by
+    /// position, not by name, unless it says BY NAME.
+    pub(crate) branch_names: Option<BranchNames>,
+    /// The first branch's names of the set operation read last, for one that
+    /// holds it as its own first branch.
+    pub(crate) last_branch_names: Option<(Option<String>, Vec<String>)>,
     /// Last join/operation type for edge labeling
     pub(crate) last_operation: Option<String>,
     /// Current join information (type + condition) for edge labeling
@@ -214,6 +221,17 @@ pub(crate) struct StatementContext {
     pub(crate) instance_limit_reached: bool,
 }
 
+/// The names a later branch of a set operation gives its columns: the first
+/// branch's, by position. Only columns of `target` take them, so that a
+/// derived table or a subquery inside the branch keeps its own.
+#[derive(Debug, Clone)]
+pub(crate) struct BranchNames {
+    pub(crate) target: Option<String>,
+    pub(crate) names: Vec<String>,
+    /// The position of the next column the branch gives `target`.
+    pub(crate) next: usize,
+}
+
 /// Represents an output column in the SELECT list
 #[derive(Debug, Clone)]
 pub(crate) struct OutputColumn {
@@ -252,6 +270,8 @@ impl StatementContext {
             derived_occurrences: HashMap::new(),
             cte_occurrences: HashMap::new(),
             predicate_subqueries: 0,
+            branch_names: None,
+            last_branch_names: None,
             last_operation: None,
             current_join_info: JoinInfo::default(),
             table_node_ids: HashMap::new(),

crates/flowscope-core/src/analyzer/query.rs

@@ -1001,7 +1001,17 @@ impl<'a> Analyzer<'a> {
         ctx: &mut StatementContext,
         params: OutputColumnParams,
     ) {
-        let normalized_name = self.normalize_identifier(&params.name);
+        // A later branch of a set operation names no column: each takes the
+        // first branch's name at its position.
+        let name = match ctx.branch_names.as_mut() {
+            Some(branch) if branch.target == params.target_node => {
+                let at = branch.next;
+                branch.next += 1;
+                branch.names.get(at).cloned()
+            }
+            _ => None,
+        };
+        let normalized_name = self.normalize_identifier(name.as_deref().unwrap_or(&params.name));
         let node_id = generate_column_node_id(params.target_node.as_deref(), &normalized_name);
         let ownership_edge_id = params
             .target_node

crates/flowscope-core/src/analyzer/visitor.rs

@@ -4,20 +4,20 @@
 //! and building lineage graphs. It separates traversal logic (the `Visitor` trait)
 //! from analysis logic (the `LineageVisitor` implementation).
 
-use super::context::{OutputColumn, StatementContext};
+use super::context::{BranchNames, OutputColumn, StatementContext};
 use super::expression::ExpressionAnalyzer;
 use super::helpers::{
     alias_visibility_warning, find_cte_body_span, find_cte_definition_span,
-    find_derived_table_alias_span, generate_statement_scoped_node_id,
+    find_derived_table_alias_span, generate_column_node_id, generate_statement_scoped_node_id,
 };
 use super::query::OutputColumnParams;
 use super::select_analyzer::SelectAnalyzer;
 use super::Analyzer;
 use crate::generated::is_value_table_function;
 use crate::types::{issue_codes, Issue, Node, NodeType, Span};
 use sqlparser::ast::{
-    self, CreateView, Cte, Expr, Ident, Join, Query, Select, SetExpr, SetOperator, Statement,
-    TableAlias, TableFactor, TableWithJoins, Values,
+    self, CreateView, Cte, Expr, Ident, Join, Query, Select, SetExpr, SetOperator, SetQuantifier,
+    Statement, TableAlias, TableFactor, TableWithJoins, Values,
 };
 use std::sync::Arc;
 
@@ -723,16 +723,51 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> {
             SetExpr::Select(select) => self.visit_select(select),
             SetExpr::Query(query) => self.visit_query(query),
             SetExpr::SetOperation {
-                op, left, right, ..
+                op,
+                set_quantifier,
+                left,
+                right,
             } => {
                 let op_name = match op {
                     SetOperator::Union => "UNION",
                     SetOperator::Intersect => "INTERSECT",
                     SetOperator::Except => "EXCEPT",
                     SetOperator::Minus => "MINUS",
                 };
+                let sink = self.ctx.sink_node_id().map(|id| id.to_string());
+                let target = self.target_node.clone().or(sink);
+                let checkpoint = self.ctx.output_columns.len();
+                self.ctx.last_branch_names = None;
                 self.visit_set_expr(left);
+                // The first branch's names: a nested set operation's own first
+                // branch's, or the columns the branch just gave the target.
+                let names = match self.ctx.last_branch_names.take() {
+                    Some((of, names)) if of == target => names,
+                    _ => self.ctx.output_columns[checkpoint..]
+                        .iter()
+                        .filter(|c| {
+                            generate_column_node_id(target.as_deref(), &c.name) == c.node_id
+                        })
+                        .map(|c| c.name.clone())
+                        .collect(),
+                };
+                let by_name = matches!(
+                    set_quantifier,
+                    SetQuantifier::ByName
+                        | SetQuantifier::AllByName
+                        | SetQuantifier::DistinctByName
+                );
+                let outer = std::mem::replace(
+                    &mut self.ctx.branch_names,
+                    (!by_name).then(|| BranchNames {
+                        target: target.clone(),
+                        names: names.clone(),
+                        next: 0,
+                    }),
+                );
                 self.visit_set_expr(right);
+                self.ctx.branch_names = outer;
+                self.ctx.last_branch_names = Some((target, names));
                 if self.target_node.is_some() {
                     self.ctx.last_operation = Some(op_name.to_string());
                 }

crates/flowscope-core/tests/lineage_engine.rs

@@ -2430,46 +2430,48 @@ fn cte_self_join_alias_columns_do_not_leak_across_union_scopes() {
         .iter()
         .find(|node| node.node_type == NodeType::Output)
         .expect("Output node should exist");
-    let dept_output = stmt
+    // The branches pair by position: the output has the first branch's name
+    // only, and the second branch's `b.name` flows into it.
+    let output_columns: Vec<_> = stmt
         .nodes
         .iter()
-        .find(|node| {
+        .filter(|node| {
             node.node_type == NodeType::Column
-                && &*node.label == "dept_name"
                 && stmt.edges.iter().any(|edge| {
                     edge.edge_type == EdgeType::Ownership
                         && edge.from == output_node.id
                         && edge.to == node.id
                 })
         })
-        .expect("dept_name output column should exist");
+        .collect();
+    assert_eq!(
+        output_columns
+            .iter()
+            .map(|node| node.label.to_string())
+            .collect::<Vec<_>>(),
+        vec!["emp_name"],
+        "a UNION's output takes its first branch's names"
+    );
 
-    let dept_source_ids: HashSet<_> = stmt
+    let source_ids: HashSet<_> = stmt
         .edges
         .iter()
-        .filter(|edge| edge.edge_type == EdgeType::DataFlow && edge.to == dept_output.id)
+        .filter(|edge| edge.edge_type == EdgeType::DataFlow && edge.to == output_columns[0].id)
         .map(|edge| edge.from.clone())
         .collect();
-    assert!(
-        !dept_source_ids.is_empty(),
-        "dept_name should receive a direct data-flow edge from the dept branch source column"
-    );
-
     let source_owner_labels: HashSet<_> = stmt
         .edges
         .iter()
-        .filter(|edge| edge.edge_type == EdgeType::Ownership && dept_source_ids.contains(&edge.to))
+        .filter(|edge| edge.edge_type == EdgeType::Ownership && source_ids.contains(&edge.to))
         .filter_map(|edge| stmt.nodes.iter().find(|node| node.id == edge.from))
         .map(|node| node.label.to_string())
         .collect();
 
-    assert!(
-        source_owner_labels.contains("dept"),
-        "dept_name should be sourced from a dept CTE instance, saw owners {source_owner_labels:?}"
-    );
-    assert!(
-        !source_owner_labels.contains("emp"),
-        "dept_name should not reuse emp columns when aliases are reused across UNION branches, saw owners {source_owner_labels:?}"
+    // Each branch's `b` is its own CTE: the dept branch must not reuse emp's.
+    assert_eq!(
+        source_owner_labels,
+        HashSet::from(["emp".to_string(), "dept".to_string()]),
+        "each branch should read its own CTE instance, saw owners {source_owner_labels:?}"
     );
 }
 
@@ -13424,3 +13426,127 @@ fn an_outer_cte_named_like_an_inlined_model_s_inner_one_reads_through_it() {
         expected_sources(&["t.ts"])
     );
 }
+
+#[test]
+fn a_union_pairs_its_branches_by_position() {
+    // The second branch names its columns otherwise: SQL pairs them with the
+    // first branch's by position, and the output has the first branch's names.
+    let sql = r#"
+        SELECT k, a AS x FROM t1
+        UNION ALL
+        SELECT m.k, m.b FROM t2 AS m
+    "#;
+    let schema = schema_of(vec![
+        schema_table(None, None, "t1", &["k", "a"]),
+        schema_table(None, None, "t2", &["k", "b"]),
+    ]);
+
+    let result = run_analysis(sql, Dialect::Snowflake, Some(schema));
+    let stmt = first_statement(&result);
+    // The output has the first branch's names, none of the second's.
+    let output = stmt
+        .nodes
+        .iter()
+        .find(|node| node.node_type == NodeType::Output)
+        .expect("Output node should exist");
+    let mut outputs: Vec<String> = stmt
+        .edges
+        .iter()
+        .filter(|edge| edge.edge_type == EdgeType::Ownership && edge.from == output.id)
+        .filter_map(|edge| stmt.nodes.iter().find(|node| node.id == edge.to))
+        .map(|node| node.label.to_lowercase())
+        .collect();
+    outputs.sort();
+    assert_eq!(outputs, vec!["k", "x"]);
+    assert_eq!(
+        base_sources_of(&stmt, "x"),
+        expected_sources(&["t1.a", "t2.b"])
+    );
+    assert_eq!(
+        base_sources_of(&stmt, "k"),
+        expected_sources(&["t1.k", "t2.k"])
+    );
+}
+
+#[test]
+fn a_union_in_a_cte_pairs_by_position_through_its_stars_and_nesting() {
+    // Three branches, the last a star, inside a CTE the statement reads.
+    let sql = r#"
+        WITH u AS (
+            SELECT k, 'one' AS origin, a AS x FROM t1
+            UNION ALL
+            SELECT k, 'two', b FROM t2
+            UNION ALL
+            SELECT * FROM t3
+        )
+        SELECT u.k, u.origin, u.x FROM u
+    "#;
+    let schema = schema_of(vec![
+        schema_table(None, None, "t1", &["k", "a"]),
+        schema_table(None, None, "t2", &["k", "b"]),
+        schema_table(None, None, "t3", &["kk", "label", "c"]),
+    ]);
+
+    let result = run_analysis(sql, Dialect::Snowflake, Some(schema));
+    let stmt = first_statement(&result);
+    assert_eq!(
+        base_sources_of(&stmt, "x"),
+        expected_sources(&["t1.a", "t2.b", "t3.c"])
+    );
+    assert_eq!(
+        base_sources_of(&stmt, "k"),
+        expected_sources(&["t1.k", "t2.k", "t3.kk"])
+    );
+    assert_eq!(
+        base_sources_of(&stmt, "origin"),
+        expected_sources(&["t3.label"])
+    );
+}
+
+#[test]
+fn a_union_by_name_still_pairs_by_name() {
+    let sql = r#"
+        SELECT k, a AS x FROM t1
+        UNION ALL BY NAME
+        SELECT b AS x, k FROM t2
+    "#;
+    let schema = schema_of(vec![
+        schema_table(None, None, "t1", &["k", "a"]),
+        schema_table(None, None, "t2", &["k", "b"]),
+    ]);
+
+    let result = run_analysis(sql, Dialect::Snowflake, Some(schema));
+    let stmt = first_statement(&result);
+    assert_eq!(
+        base_sources_of(&stmt, "x"),
+        expected_sources(&["t1.a", "t2.b"])
+    );
+    assert_eq!(
+        base_sources_of(&stmt, "k"),
+        expected_sources(&["t1.k", "t2.k"])
+    );
+}
+
+#[test]
+fn a_derived_table_inside_a_later_branch_keeps_its_own_names() {
+    let sql = r#"
+        SELECT k, a AS x FROM t1
+        UNION ALL
+        SELECT d.k, d.y FROM (SELECT k, b AS y FROM t2) AS d
+    "#;
+    let schema = schema_of(vec![
+        schema_table(None, None, "t1", &["k", "a"]),
+        schema_table(None, None, "t2", &["k", "b"]),
+    ]);
+
+    let result = run_analysis(sql, Dialect::Snowflake, Some(schema));
+    let stmt = first_statement(&result);
+    assert_eq!(
+        base_sources_of(&stmt, "x"),
+        expected_sources(&["t1.a", "t2.b"])
+    );
+    assert_eq!(
+        base_sources_of(&stmt, "k"),
+        expected_sources(&["t1.k", "t2.k"])
+    );
+}

crates/flowscope-core/tests/snapshots/snapshots__run_snapshot_test@recursive_cte.snap

@@ -1,5 +1,6 @@
 ---
 source: crates/flowscope-core/tests/snapshots.rs
+assertion_line: 27
 expression: clean_result
 ---
 {
@@ -273,6 +274,15 @@ expression: clean_result
         0
       ]
     },
+    {
+      "id": "edge_8e6e75d17a8dea9d",
+      "from": "column_78a6f51170482ff9",
+      "to": "column_a5ca696854a3dcdd",
+      "type": "data_flow",
+      "statementIds": [
+        0
+      ]
+    },
     {
       "id": "edge_9837a0968ff8ff3f",
       "from": "column_f0521d6d6b3f8f6b",