From 56677a080b976be4c319f985a1aac46436ea6128 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sat, 22 Aug 2026 01:24:37 +0000 Subject: [PATCH 01/14] fix: make column mapping order deterministic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The order of `ColumnLineage.mappings` varied between runs of the same binary whenever a query had duplicate output column names: SELECT a.id, b.id FROM a JOIN b ON a.id = b.bid sometimes [ id <- a.id , id <- b.id ] sometimes [ id <- b.id , id <- a.id ] With three duplicates, four distinct orderings showed up across six runs. Callers that cache or diff results see the same input produce different output. `resolve` collected the output nodes into a `HashSet` and built the mappings by iterating it, so construction order followed hash order with a per-process random seed. The sort afterwards could not undo this: it keyed on a `HashMap` of output names, so duplicate names collided and one index won, leaving same-named mappings in whatever order the set had produced. `ordered_cols` already holds the projection order, so build the mappings straight from it. That makes the sort redundant — it can only reproduce the order the loop now has, and it is the reason duplicates were reordered in the first place — so it goes too. Co-authored-by: Claude Opus 5 --- sqllineage/src/resolve/mod.rs | 18 +++----------- sqllineage/tests/column_lineage.rs | 40 ++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index c544d0e..4ef1d2d 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1,7 +1,7 @@ mod catalog; mod topo; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; @@ -53,12 +53,11 @@ pub(crate) fn resolve( let root = ScopeTree::root(); let ordered_cols = graph.scopes.output_columns(root).to_vec(); - let final_ids: HashSet = ordered_cols.iter().map(|c| c.node_id).collect(); - let output_table = graph.tables.output.clone(); let mut mappings = Vec::new(); - for &node_id in &final_ids { + for col in &ordered_cols { + let node_id = col.node_id; match &graph.nodes[node_id] { RawNode::Output { name, .. } => { let mut visited = HashSet::new(); @@ -96,17 +95,6 @@ pub(crate) fn resolve( } } - let name_order: HashMap = ordered_cols - .iter() - .enumerate() - .filter_map(|(i, c)| match &graph.nodes[c.node_id] { - RawNode::Output { name, .. } => Some((name.clone(), i)), - RawNode::Star { .. } => Some(("*".to_string(), i)), - _ => None, - }) - .collect(); - mappings.sort_by_key(|m| name_order.get(&m.target.column).copied().unwrap_or(usize::MAX)); - if let Some(cat) = catalog { catalog::apply_catalog(&mut mappings, cat); } diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index 3418c24..a183107 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -80,6 +80,46 @@ fn select_multiple_tables_qualified() { assert_eq!(concrete_sources(m_b), vec![("t2".into(), "b".into())]); } +#[test] +fn duplicate_output_names_preserve_projection_order() { + let result = analyze_one("SELECT a.id, b.id FROM a JOIN b ON a.id = b.bid"); + let sources: Vec<_> = result + .columns + .mappings + .iter() + .map(concrete_sources) + .collect(); + + assert_eq!( + sources, + vec![ + vec![("a".into(), "id".into())], + vec![("b".into(), "id".into())] + ] + ); +} + +#[test] +fn three_duplicate_output_names_preserve_projection_order() { + let result = + analyze_one("SELECT a.id, b.id, c.id FROM a JOIN b ON a.id = b.bid JOIN c ON a.id = c.cid"); + let sources: Vec<_> = result + .columns + .mappings + .iter() + .map(concrete_sources) + .collect(); + + assert_eq!( + sources, + vec![ + vec![("a".into(), "id".into())], + vec![("b".into(), "id".into())], + vec![("c".into(), "id".into())], + ] + ); +} + #[test] fn select_case_expression() { let result = analyze_one("SELECT CASE WHEN a > 0 THEN b ELSE c END AS d FROM t"); From e8e47908536232b28f1252045394a412eb865b09 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sat, 22 Aug 2026 04:40:45 +0000 Subject: [PATCH 02/14] fix: stop fabricating column origins for unresolved columns The resolver returned ColumnOrigin::Concrete with invented table names ("?unknown?", "?cte?") when a column could not be resolved, so callers could not tell fabricated lineage from real lineage. Both sites now return ColumnOrigin::Ambiguous with an empty candidate list, which already means "catalog needed to resolve" and keeps the 0.2.x public API unchanged. Returning None instead would drop the source from the mapping entirely, and an empty source list already has a valid meaning (constant expressions). apply_catalog now skips empty candidate lists. Without that guard a CatalogProvider resolving a column by name alone would turn an unresolved origin straight back into a fabricated Concrete one. --- sqllineage-python/sqllineage.pyi | 3 ++- sqllineage/src/resolve/catalog.rs | 1 + sqllineage/src/resolve/mod.rs | 8 +++---- sqllineage/src/types.rs | 8 +++++-- sqllineage/tests/catalog.rs | 35 ++++++++++++++++++++++++++++++ sqllineage/tests/column_lineage.rs | 15 ++++++++++++- sqllineage/tests/cte.rs | 14 ++++++++++++ 7 files changed, 76 insertions(+), 8 deletions(-) diff --git a/sqllineage-python/sqllineage.pyi b/sqllineage-python/sqllineage.pyi index 04f025b..1f8bde5 100644 --- a/sqllineage-python/sqllineage.pyi +++ b/sqllineage-python/sqllineage.pyi @@ -17,7 +17,8 @@ class ColumnOrigin: Check ``kind`` to determine the variant: - ``"concrete"``: ``table`` and ``column`` are set. - - ``"ambiguous"``: ``column`` and ``candidates`` are set. + - ``"ambiguous"``: ``column`` and ``candidates`` are set. ``candidates`` + may be an empty list when the column is unresolved. - ``"wildcard"``: ``table`` is set. - ``"recursive"``: ``base_sources`` is set. """ diff --git a/sqllineage/src/resolve/catalog.rs b/sqllineage/src/resolve/catalog.rs index 97dcc71..52363ea 100644 --- a/sqllineage/src/resolve/catalog.rs +++ b/sqllineage/src/resolve/catalog.rs @@ -5,6 +5,7 @@ pub(crate) fn apply_catalog(mappings: &mut Vec, catalog: &dyn Cat for mapping in mappings.iter_mut() { for source in &mut mapping.sources { if let ColumnOrigin::Ambiguous { column, candidates } = source + && !candidates.is_empty() && let Some(owner) = catalog.resolve_column(column, candidates) { *source = ColumnOrigin::Concrete { diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 4ef1d2d..297bfe8 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -385,9 +385,9 @@ fn resolve_from_bindings( } } } else if bindings.is_empty() { - Some(ColumnOrigin::Concrete { - table: TableRef::new("?unknown?"), + Some(ColumnOrigin::Ambiguous { column: name.to_string(), + candidates: Vec::new(), }) } else { let mut table_candidates = Vec::new(); @@ -444,9 +444,9 @@ fn resolve_through_scope( origins.into_iter().next() } } else { - Some(ColumnOrigin::Concrete { - table: TableRef::new("?cte?"), + Some(ColumnOrigin::Ambiguous { column: column_name.to_string(), + candidates: Vec::new(), }) } } diff --git a/sqllineage/src/types.rs b/sqllineage/src/types.rs index fbdc051..5f392b9 100644 --- a/sqllineage/src/types.rs +++ b/sqllineage/src/types.rs @@ -122,7 +122,10 @@ pub struct ColumnMapping { pub enum ColumnOrigin { /// Fully resolved to a specific table and column. Concrete { table: TableRef, column: String }, - /// Multiple candidate tables; catalog needed to disambiguate. + /// Multiple candidate tables. A non-empty `candidates` list means genuine + /// ambiguity between known tables and can be disambiguated by a catalog. + /// An empty list means the column could not be resolved to any known table; + /// catalog refinement is not attempted. Ambiguous { column: String, candidates: Vec, @@ -225,6 +228,7 @@ impl std::error::Error for ParseError {} pub trait CatalogProvider { /// Return the column names of a table. Used to expand `SELECT *`. fn list_columns(&self, table: &TableRef) -> Option>; - /// Given a column name and candidate tables, return the owning table. + /// Given a column name and candidate tables, return the owning table. This + /// is only called with a non-empty candidate slice. fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option; } diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 73bf2a6..72f0ca9 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -23,6 +23,18 @@ impl CatalogProvider for MockCatalog { } } +struct EagerCatalog; + +impl CatalogProvider for EagerCatalog { + fn list_columns(&self, _table: &TableRef) -> Option> { + None + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + Some(TableRef::new("fabricated")) + } +} + fn opts_with_catalog() -> AnalyzeOptions { AnalyzeOptions { catalog: Some(Box::new(MockCatalog)), @@ -110,6 +122,29 @@ fn ambiguous_column_without_catalog() { } } +#[test] +fn catalog_does_not_fabricate_unresolved_column_owner() { + let result = analyze( + "SELECT missing", + AnalyzeOptions { + catalog: Some(Box::new(EagerCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn catalog_preserves_qualified_columns() { let sql = diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index a183107..5388a4a 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -1,7 +1,7 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping, table}; -use sqllineage::TransformKind; +use sqllineage::{ColumnOrigin, TransformKind}; #[test] fn select_columns() { @@ -18,6 +18,19 @@ fn select_columns() { assert_eq!(m_b.transform, TransformKind::Direct); } +#[test] +fn unresolved_column_has_empty_ambiguous_candidates() { + let result = analyze_one("SELECT missing"); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn select_expression() { let result = analyze_one("SELECT a + b AS c FROM t"); diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index 24bdfbd..cd2e6c1 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -15,6 +15,20 @@ fn single_cte() { assert_eq!(m.transform, TransformKind::Direct); } +#[test] +fn missing_column_from_cte_has_empty_ambiguous_candidates() { + let sql = "WITH cte AS (SELECT present FROM source) SELECT missing FROM cte"; + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "missing"); + match &m.sources[0] { + ColumnOrigin::Ambiguous { column, candidates } => { + assert_eq!(column, "missing"); + assert!(candidates.is_empty()); + } + other => panic!("expected Ambiguous, got {other:?}"), + } +} + #[test] fn cte_chain() { let sql = "WITH a AS (SELECT x FROM t), b AS (SELECT x FROM a) SELECT x FROM b"; From 38149bf689f0e4a3557e192c68bc94dcabc4eea8 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 03:11:48 +0000 Subject: [PATCH 03/14] fix: classify source-free aggregate calls correctly (COUNT(*)) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit determine_edge_kind correctly identifies any aggregate function call as EdgeKind::ViaAggregation regardless of its arguments, but that kind was only ever attached to the graph as an edge to an ancestor column. COUNT(*) has no column ancestor (`*` is a FunctionArgExpr::Wildcard, not an Expr, so collect_ancestors never visits it), so the correctly computed kind was silently discarded and the output's transform classification fell back to Direct — indistinguishable from a literal constant. Store the defining expression's intrinsic edge kind on the Output node itself, and use it as a fallback in derive_transform whenever no ancestor edge exists to classify from. This generalizes to any zero-ancestor aggregate/conditional/expression, not just COUNT(*). --- sqllineage/src/build/select.rs | 6 +++--- sqllineage/src/build/statement.rs | 8 ++++---- sqllineage/src/graph/mod.rs | 7 +++++-- sqllineage/src/graph/node.rs | 10 +++++++++- sqllineage/src/resolve/mod.rs | 15 ++++++++++++--- sqllineage/tests/column_lineage.rs | 9 +++++++++ 6 files changed, 42 insertions(+), 13 deletions(-) diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index 803400f..4ed3489 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -24,7 +24,7 @@ impl LineageBuilder { let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); let name = infer_column_name(expr); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -40,7 +40,7 @@ impl LineageBuilder { let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); let name = alias.value.clone(); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -57,7 +57,7 @@ impl LineageBuilder { let kind = determine_edge_kind(expr); for alias in aliases { let name = alias.value.clone(); - let output = self.graph.add_output(name.clone()); + let output = self.graph.add_output(name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index 684f052..0c2980d 100644 --- a/sqllineage/src/build/statement.rs +++ b/sqllineage/src/build/statement.rs @@ -50,7 +50,7 @@ impl LineageBuilder { let col_name = assignment_target_name(&assignment.target); let ancestors = self.collect_ancestors(&assignment.value); let kind = determine_edge_kind(&assignment.value); - let output = self.graph.add_output(col_name.clone()); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -101,7 +101,7 @@ impl LineageBuilder { let col_name = assignment_target_name(&assignment.target); let ancestors = self.collect_ancestors(&assignment.value); let kind = determine_edge_kind(&assignment.value); - let output = self.graph.add_output(col_name.clone()); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -134,7 +134,7 @@ impl LineageBuilder { .unwrap_or_else(|| format!("col{i}")); let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); - let output = self.graph.add_output(col_name.clone()); + let output = self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } @@ -280,7 +280,7 @@ impl LineageBuilder { | Statement::UNCache { .. } | Statement::UNLISTEN { .. } | Statement::Unload { .. } - | Statement::UnlockTables { .. } + | Statement::UnlockTables | Statement::Use(_) | Statement::Vacuum { .. } | Statement::WaitFor { .. } diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 6cc33cd..1e2134d 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -31,8 +31,11 @@ impl RawGraph { id } - pub fn add_output(&mut self, name: String) -> NodeId { - self.add_node(RawNode::Output { name }) + pub fn add_output(&mut self, name: String, intrinsic_kind: EdgeKind) -> NodeId { + self.add_node(RawNode::Output { + name, + intrinsic_kind, + }) } pub fn add_ref(&mut self, name: String, qualifier: Option, scope: ScopeId) -> NodeId { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index 3227495..ef7eafa 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -1,3 +1,4 @@ +use crate::graph::edge::EdgeKind; use crate::graph::scope::ScopeId; use crate::types::TableRef; @@ -6,7 +7,14 @@ pub(crate) type NodeId = usize; #[derive(Debug, Clone)] pub(crate) enum RawNode { /// Output column — produced by a projection or assignment. - Output { name: String }, + Output { + name: String, + /// The edge kind the defining expression would carry to its own + /// ancestors, kept even when it has none (e.g. `COUNT(*)` has no + /// column ancestor but is still an aggregate). Used as a fallback + /// classification when no ancestor edge exists to classify from. + intrinsic_kind: EdgeKind, + }, /// Named reference — alias, CTE reference, derived table column. Ref { name: String, diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 297bfe8..c74fa31 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -63,7 +63,7 @@ pub(crate) fn resolve( let mut visited = HashSet::new(); let (sources, edge_kinds, has_back) = collect_output_sources(node_id, &graph, &mut resolved, &incoming, &mut visited); - let transform = derive_transform(&edge_kinds); + let transform = derive_transform(&graph.nodes[node_id], &edge_kinds); if has_back { mappings.push(ColumnMapping { @@ -183,7 +183,7 @@ fn expand_scope_columns( let mut visited = HashSet::new(); let (sources, edge_kinds, _) = collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited); - let transform = derive_transform(&edge_kinds); + let transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); mappings.push(ColumnMapping { target: ColumnRef { table: output_table.cloned(), @@ -451,7 +451,16 @@ fn resolve_through_scope( } } -fn derive_transform(kinds: &[EdgeKind]) -> TransformKind { +fn derive_transform(node: &RawNode, edge_kinds: &[EdgeKind]) -> TransformKind { + let kinds = if edge_kinds.is_empty() { + match node { + RawNode::Output { intrinsic_kind, .. } => std::slice::from_ref(intrinsic_kind), + _ => edge_kinds, + } + } else { + edge_kinds + }; + if kinds.iter().any(|k| matches!(k, EdgeKind::ViaAggregation)) { TransformKind::Aggregation } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaConditional)) { diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index 5388a4a..eadc99a 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -81,6 +81,15 @@ fn select_aggregate() { assert_eq!(m.transform, TransformKind::Aggregation); } +#[test] +fn select_count_star_is_aggregation_without_sources() { + let result = analyze_one("SELECT COUNT(*) AS c FROM t"); + let m = find_mapping(&result.columns.mappings, "c"); + + assert!(m.sources.is_empty()); + assert_eq!(m.transform, TransformKind::Aggregation); +} + #[test] fn select_multiple_tables_qualified() { let result = analyze_one("SELECT t1.a, t2.b FROM t1 JOIN t2 ON t1.id = t2.id"); From c8d841ac0f8c840b27d189e37982ffa65873956f Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 04:26:37 +0000 Subject: [PATCH 04/14] fix: resolve set operations after star expansion --- sqllineage/src/build/query.rs | 50 ++--- sqllineage/src/graph/scope.rs | 24 +++ sqllineage/src/resolve/mod.rs | 395 ++++++++++++++++++++++++++++------ sqllineage/tests/catalog.rs | 153 +++++++++++++ sqllineage/tests/cte.rs | 58 +++++ 5 files changed, 579 insertions(+), 101 deletions(-) diff --git a/sqllineage/src/build/query.rs b/sqllineage/src/build/query.rs index b558d5c..59b7b96 100644 --- a/sqllineage/src/build/query.rs +++ b/sqllineage/src/build/query.rs @@ -1,9 +1,7 @@ use sqlparser::ast::{Query, SetExpr}; use crate::build::LineageBuilder; -use crate::graph::edge::RawEdge; -use crate::graph::node::NodeId; -use crate::graph::scope::{Binding, ScopeKind}; +use crate::graph::scope::{Binding, OutputPlan, ScopeKind}; impl LineageBuilder { pub(crate) fn visit_query(&mut self, query: &Query) { @@ -39,7 +37,13 @@ impl LineageBuilder { .scopes .output_columns(self.current_scope) .to_vec(); + let body_scope = self.current_scope; self.pop_scope(); + // The parent owns the query's public output. Keep the child plan + // intact and delegate through it after returning to the parent. + self.graph + .scopes + .set_output_plan(self.current_scope, OutputPlan::Delegate(body_scope)); for col in body_outputs { self.graph.scopes.add_output_column(self.current_scope, col); } @@ -54,7 +58,7 @@ impl LineageBuilder { SetExpr::SetOperation { left, right, .. } => { let left_scope = self.push_scope(ScopeKind::SetOperation); self.visit_set_expr(left); - let left_outputs: Vec<(String, NodeId)> = self + let left_outputs: Vec<(String, crate::graph::node::NodeId)> = self .graph .scopes .output_columns(left_scope) @@ -65,39 +69,17 @@ impl LineageBuilder { let right_scope = self.push_scope(ScopeKind::SetOperation); self.visit_set_expr(right); - let right_outputs: Vec<(String, NodeId)> = self - .graph - .scopes - .output_columns(right_scope) - .iter() - .map(|c| (c.name.clone(), c.node_id)) - .collect(); self.pop_scope(); let is_recursive = self.recursive_cte_name.is_some(); - - let pair_count = left_outputs.len().min(right_outputs.len()); - for i in 0..pair_count { - let left_out = left_outputs[i].1; - let right_out = right_outputs[i].1; - - let redirected: Vec<(NodeId, _)> = self - .graph - .edges - .iter() - .filter(|e| e.to == right_out) - .map(|e| (e.from, e.kind.clone())) - .collect(); - - for (from, kind) in redirected { - self.graph.edges.push(RawEdge { - from, - to: left_out, - kind, - is_recursive_back_edge: is_recursive, - }); - } - } + self.graph.scopes.set_output_plan( + self.current_scope, + OutputPlan::SetOperation { + left: left_scope, + right: right_scope, + recursive: is_recursive, + }, + ); for (name, node_id) in &left_outputs { self.graph.scopes.add_output_column( diff --git a/sqllineage/src/graph/scope.rs b/sqllineage/src/graph/scope.rs index de9f4ab..697fbce 100644 --- a/sqllineage/src/graph/scope.rs +++ b/sqllineage/src/graph/scope.rs @@ -14,6 +14,7 @@ struct Scope { bindings: HashMap, anonymous_derived: Vec, output_columns: Vec, + output_plan: OutputPlan, } #[derive(Debug, Clone)] @@ -25,6 +26,19 @@ pub(crate) enum ScopeKind { SetOperation, } +/// Describes how a scope's output columns are assembled. This is retained in +/// the raw graph until catalog expansion and positional set-operation merge. +#[derive(Debug, Clone)] +pub(crate) enum OutputPlan { + Projection, + SetOperation { + left: ScopeId, + right: ScopeId, + recursive: bool, + }, + Delegate(ScopeId), +} + #[derive(Debug, Clone)] pub(crate) enum Binding { Table(TableRef), @@ -46,6 +60,7 @@ impl ScopeTree { bindings: HashMap::new(), anonymous_derived: Vec::new(), output_columns: Vec::new(), + output_plan: OutputPlan::Projection, }], } } @@ -61,6 +76,7 @@ impl ScopeTree { bindings: HashMap::new(), anonymous_derived: Vec::new(), output_columns: Vec::new(), + output_plan: OutputPlan::Projection, }); id } @@ -91,6 +107,14 @@ impl ScopeTree { &self.scopes[scope].output_columns } + pub fn set_output_plan(&mut self, scope: ScopeId, plan: OutputPlan) { + self.scopes[scope].output_plan = plan; + } + + pub fn output_plan(&self, scope: ScopeId) -> &OutputPlan { + &self.scopes[scope].output_plan + } + pub fn add_anonymous_derived(&mut self, parent: ScopeId, child: ScopeId) { self.scopes[parent].anonymous_derived.push(child); } diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index c74fa31..f3f3a28 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -6,13 +6,14 @@ use std::collections::HashSet; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; use crate::graph::node::{NodeId, RawNode}; -use crate::graph::scope::{Binding, ScopeTree}; +use crate::graph::scope::{Binding, OutputPlan, ScopeTree}; use crate::types::{ AnalyzeResult, CatalogProvider, ColumnLineage, ColumnMapping, ColumnOrigin, ColumnRef, StatementType, TableRef, TransformKind, Warning, WarningKind, }; /// Resolve `RawGraph` into `AnalyzeResult`. +#[allow(clippy::if_not_else, clippy::useless_let_if_seq)] pub(crate) fn resolve( mut graph: RawGraph, catalog: Option<&dyn CatalogProvider>, @@ -52,52 +53,15 @@ pub(crate) fn resolve( let mut resolved: Vec> = vec![None; graph.nodes.len()]; let root = ScopeTree::root(); - let ordered_cols = graph.scopes.output_columns(root).to_vec(); let output_table = graph.tables.output.clone(); - let mut mappings = Vec::new(); - - for col in &ordered_cols { - let node_id = col.node_id; - match &graph.nodes[node_id] { - RawNode::Output { name, .. } => { - let mut visited = HashSet::new(); - let (sources, edge_kinds, has_back) = - collect_output_sources(node_id, &graph, &mut resolved, &incoming, &mut visited); - let transform = derive_transform(&graph.nodes[node_id], &edge_kinds); - - if has_back { - mappings.push(ColumnMapping { - target: ColumnRef { table: output_table.clone(), column: name.clone() }, - sources: vec![ColumnOrigin::Recursive { base_sources: sources }], - transform, - }); - } else { - mappings.push(ColumnMapping { - target: ColumnRef { table: output_table.clone(), column: name.clone() }, - sources, - transform, - }); - } - } - RawNode::Star { table, scope } => { - expand_star( - table.as_ref(), - *scope, - &graph, - &mut resolved, - &incoming, - output_table.as_ref(), - &mut mappings, - &mut HashSet::new(), - ); - } - _ => {} - } - } - - if let Some(cat) = catalog { - catalog::apply_catalog(&mut mappings, cat); - } + let mappings = resolve_scope_mappings( + root, + &graph, + &mut resolved, + &incoming, + output_table.as_ref(), + catalog, + ); AnalyzeResult { statement_type, @@ -129,6 +93,134 @@ fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> } } +/// Resolve a scope's output plan into ordered mappings. Set-operation +/// branches are resolved independently so catalog expansion happens before +/// their positional merge. +#[allow(clippy::too_many_arguments)] +fn resolve_scope_mappings( + scope: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + output_table: Option<&TableRef>, + catalog: Option<&dyn CatalogProvider>, +) -> Vec { + match graph.scopes.output_plan(scope).clone() { + OutputPlan::Projection => { + let mut mappings = Vec::new(); + for col in graph.scopes.output_columns(scope) { + match &graph.nodes[col.node_id] { + RawNode::Output { name, .. } => { + let mut visited = HashSet::new(); + let (sources, edge_kinds, has_back) = collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + ); + let transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); + mappings.push(ColumnMapping { + target: ColumnRef { + table: output_table.cloned(), + column: name.clone(), + }, + sources: if has_back { + vec![ColumnOrigin::Recursive { + base_sources: sources, + }] + } else { + sources + }, + transform, + }); + } + RawNode::Star { table, scope } => expand_star( + table.as_ref(), + *scope, + graph, + resolved, + incoming, + catalog, + output_table, + &mut mappings, + &mut HashSet::new(), + ), + _ => {} + } + } + if let Some(cat) = catalog { + catalog::apply_catalog(&mut mappings, cat); + } + mappings + } + OutputPlan::Delegate(child) => { + resolve_scope_mappings(child, graph, resolved, incoming, output_table, catalog) + } + OutputPlan::SetOperation { + left, + right, + recursive, + } => { + let left_mappings = + resolve_scope_mappings(left, graph, resolved, incoming, output_table, catalog); + let right_mappings = if recursive { + Vec::new() + } else { + resolve_scope_mappings(right, graph, resolved, incoming, output_table, catalog) + }; + // The branch resolvers expand stars independently. Preserve the + // left branch's names and order, as SQL set operations do. + let mut merged = Vec::with_capacity(left_mappings.len()); + for (idx, left_mapping) in left_mappings.into_iter().enumerate() { + let mut sources = left_mapping.sources; + let mut transform = left_mapping.transform.clone(); + if let Some(right_mapping) = right_mappings.get(idx) { + sources.extend(right_mapping.sources.clone()); + transform = merge_transform(&transform, &right_mapping.transform); + } + if recursive { + merged.push(ColumnMapping { + target: left_mapping.target, + sources: vec![ColumnOrigin::Recursive { + base_sources: sources, + }], + transform, + }); + } else { + merged.push(ColumnMapping { + target: left_mapping.target, + sources, + transform, + }); + } + } + merged + } + } +} + +fn merge_transform(left: &TransformKind, right: &TransformKind) -> TransformKind { + if matches!(left, TransformKind::Aggregation) || matches!(right, TransformKind::Aggregation) { + TransformKind::Aggregation + } else if matches!(left, TransformKind::Conditional) + || matches!(right, TransformKind::Conditional) + { + TransformKind::Conditional + } else if matches!(left, TransformKind::Expression) + || matches!(right, TransformKind::Expression) + { + TransformKind::Expression + } else if matches!(left, TransformKind::Window) || matches!(right, TransformKind::Window) { + TransformKind::Window + } else if matches!(left, TransformKind::Unknown) || matches!(right, TransformKind::Unknown) { + TransformKind::Unknown + } else { + TransformKind::Direct + } +} + /// Expand a Star node (qualified or unqualified) into `ColumnMapping`s. #[allow(clippy::too_many_arguments)] fn expand_star( @@ -137,6 +229,7 @@ fn expand_star( graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, output_table: Option<&TableRef>, mappings: &mut Vec, visited_scopes: &mut HashSet, @@ -144,7 +237,16 @@ fn expand_star( if let Some(t) = table { let binding = graph.scopes.lookup(scope, &t.table).cloned(); if let Some(Binding::Cte(s) | Binding::DerivedTable(s)) = binding { - expand_scope_columns(s, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_scope_columns( + s, + graph, + resolved, + incoming, + catalog, + output_table, + mappings, + visited_scopes, + ); } else { mappings.push(wildcard_mapping(output_table, t.clone())); } @@ -153,22 +255,42 @@ fn expand_star( match binding { Binding::Table(tref) => mappings.push(wildcard_mapping(output_table, tref)), Binding::Cte(s) | Binding::DerivedTable(s) => { - expand_scope_columns(s, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_scope_columns( + s, + graph, + resolved, + incoming, + catalog, + output_table, + mappings, + visited_scopes, + ); } } } for &child in graph.scopes.anonymous_derived(scope) { - expand_scope_columns(child, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_scope_columns( + child, + graph, + resolved, + incoming, + catalog, + output_table, + mappings, + visited_scopes, + ); } } } /// Recursively expand a scope's output columns into `ColumnMapping`s. +#[allow(clippy::too_many_arguments)] fn expand_scope_columns( scope_id: usize, graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, output_table: Option<&TableRef>, mappings: &mut Vec, visited_scopes: &mut HashSet, @@ -176,13 +298,35 @@ fn expand_scope_columns( if !visited_scopes.insert(scope_id) { return; } + if !matches!(graph.scopes.output_plan(scope_id), OutputPlan::Projection) { + let mut nested = + resolve_scope_mappings(scope_id, graph, resolved, incoming, output_table, catalog); + mappings.append(&mut nested); + return; + } for col in graph.scopes.output_columns(scope_id) { if let RawNode::Star { table, scope } = &graph.nodes[col.node_id] { - expand_star(table.as_ref(), *scope, graph, resolved, incoming, output_table, mappings, visited_scopes); + expand_star( + table.as_ref(), + *scope, + graph, + resolved, + incoming, + catalog, + output_table, + mappings, + visited_scopes, + ); } else { let mut visited = HashSet::new(); - let (sources, edge_kinds, _) = - collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited); + let (sources, edge_kinds, _) = collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + ); let transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); mappings.push(ColumnMapping { target: ColumnRef { @@ -196,12 +340,42 @@ fn expand_scope_columns( } } +/// Collect source origins for one logical output slot, retaining both sides +/// of a set operation. Unlike the public mapping path this returns origins so +/// a later CTE/derived-table reference can continue through the slot. +fn scope_column_sources( + scope: usize, + index: usize, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, +) -> (Vec, Vec, bool) { + let mappings = resolve_scope_mappings(scope, graph, resolved, incoming, None, catalog); + let Some(mapping) = mappings.get(index) else { + return (vec![], vec![], false); + }; + let mut sources = Vec::new(); + let mut has_back = false; + for source in &mapping.sources { + match source { + ColumnOrigin::Recursive { base_sources } => { + sources.extend(base_sources.clone()); + has_back = true; + } + source => sources.push(source.clone()), + } + } + (sources, vec![], has_back) +} + fn collect_output_sources( node_id: NodeId, graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> (Vec, Vec, bool) { if !visited.insert(node_id) { return (vec![], vec![], false); @@ -218,7 +392,7 @@ fn collect_output_sources( continue; } let (sub_sources, sub_back) = - collect_leaf_origins(edge.from, graph, resolved, incoming, visited); + collect_leaf_origins(edge.from, graph, resolved, incoming, visited, catalog); for _ in &sub_sources { kinds.push(edge.kind.clone()); } @@ -235,19 +409,37 @@ fn collect_leaf_origins( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> (Vec, bool) { - if let Some((target_output, _)) = find_cte_redirect(node_id, graph) { + if let Some((sources, has_back)) = + resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog) + { + return (sources, has_back); + } + + if let Some((target_output, scope)) = find_cte_redirect(node_id, graph) { + if let Some(index) = graph + .scopes + .output_columns(scope) + .iter() + .position(|c| c.node_id == target_output) + && !matches!(graph.scopes.output_plan(scope), OutputPlan::Projection) + { + let (sources, _, has_back) = + scope_column_sources(scope, index, graph, resolved, incoming, catalog); + return (sources, has_back); + } let (sources, _, has_back) = - collect_output_sources(target_output, graph, resolved, incoming, visited); + collect_output_sources(target_output, graph, resolved, incoming, visited, catalog); return (sources, has_back); } if let RawNode::Output { .. } = &graph.nodes[node_id] { let (sources, _, has_back) = - collect_output_sources(node_id, graph, resolved, incoming, visited); + collect_output_sources(node_id, graph, resolved, incoming, visited, catalog); (sources, has_back) } else { - let origin = resolve_node(node_id, graph, resolved, incoming, visited); + let origin = resolve_node(node_id, graph, resolved, incoming, visited, catalog); match origin { Some(o) => (vec![o], false), None => (vec![], false), @@ -255,12 +447,61 @@ fn collect_leaf_origins( } } +/// Resolve a named reference through a set-operation/Delegate scope using the +/// expanded mappings. Raw scope columns intentionally do not contain names +/// for individual catalog-expanded star outputs, so lookup must happen here. +fn resolve_named_scope_reference( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, +) -> Option<(Vec, bool)> { + let (name, qualifier, scope) = match &graph.nodes[node_id] { + RawNode::Ref { + name, + qualifier, + scope, + } => (name, qualifier.as_ref(), *scope), + RawNode::Unqualified { name, scope } => (name, None, *scope), + _ => return None, + }; + let binding = qualifier + .and_then(|qualifier| graph.scopes.lookup(scope, qualifier).cloned()) + .or_else(|| find_single_binding(scope, graph)); + let target_scope = match binding { + Some(Binding::Cte(target) | Binding::DerivedTable(target)) + if !matches!(graph.scopes.output_plan(target), OutputPlan::Projection) => + { + target + } + _ => return None, + }; + let mappings = resolve_scope_mappings(target_scope, graph, resolved, incoming, None, catalog); + let mapping = mappings + .iter() + .find(|mapping| mapping.target.column == *name)?; + let mut sources = Vec::new(); + let mut has_back = false; + for source in &mapping.sources { + match source { + ColumnOrigin::Recursive { base_sources } => { + sources.extend(base_sources.clone()); + has_back = true; + } + source => sources.push(source.clone()), + } + } + Some((sources, has_back)) +} + fn resolve_node( node_id: NodeId, graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> Option { if let Some(ref origin) = resolved[node_id] { return Some(origin.clone()); @@ -280,7 +521,9 @@ fn resolve_node( column: name.clone(), }), Some(Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope)) => { - resolve_through_scope(name, cte_scope, graph, resolved, incoming, visited) + resolve_through_scope( + name, cte_scope, graph, resolved, incoming, visited, catalog, + ) } None => Some(ColumnOrigin::Concrete { table: TableRef::new(qual.as_str()), @@ -288,12 +531,12 @@ fn resolve_node( }), } } else { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited) + resolve_unqualified(name, *scope, graph, resolved, incoming, visited, catalog) } } RawNode::Unqualified { name, scope } => { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited) + resolve_unqualified(name, *scope, graph, resolved, incoming, visited, catalog) } RawNode::Star { table, .. } => table @@ -361,8 +604,17 @@ fn resolve_unqualified( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> Option { - resolve_from_bindings(name, &effective_bindings(scope, graph), graph, resolved, incoming, visited) + resolve_from_bindings( + name, + &effective_bindings(scope, graph), + graph, + resolved, + incoming, + visited, + catalog, + ) } fn resolve_from_bindings( @@ -372,6 +624,7 @@ fn resolve_from_bindings( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> Option { if bindings.len() == 1 { let (_, binding) = &bindings[0]; @@ -380,9 +633,9 @@ fn resolve_from_bindings( table: table_ref.clone(), column: name.to_string(), }), - Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope) => { - resolve_through_scope(name, *cte_scope, graph, resolved, incoming, visited) - } + Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope) => resolve_through_scope( + name, *cte_scope, graph, resolved, incoming, visited, catalog, + ), } } else if bindings.is_empty() { Some(ColumnOrigin::Ambiguous { @@ -394,8 +647,15 @@ fn resolve_from_bindings( for (_, binding) in bindings { match binding { Binding::Cte(s) | Binding::DerivedTable(s) => { - if graph.scopes.output_columns(*s).iter().any(|c| c.name == name) { - return resolve_through_scope(name, *s, graph, resolved, incoming, visited); + if graph + .scopes + .output_columns(*s) + .iter() + .any(|c| c.name == name) + { + return resolve_through_scope( + name, *s, graph, resolved, incoming, visited, catalog, + ); } } Binding::Table(t) => table_candidates.push(t.clone()), @@ -422,6 +682,7 @@ fn resolve_through_scope( resolved: &mut Vec>, incoming: &[Vec], visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, ) -> Option { if let Some(col) = graph .scopes @@ -430,7 +691,7 @@ fn resolve_through_scope( .find(|c| c.name == column_name) { let (origins, _, has_back) = - collect_output_sources(col.node_id, graph, resolved, incoming, visited); + collect_output_sources(col.node_id, graph, resolved, incoming, visited, catalog); if has_back { Some(ColumnOrigin::Recursive { base_sources: origins, diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 72f0ca9..872655f 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -163,3 +163,156 @@ fn catalog_preserves_qualified_columns() { vec![("orders".into(), "amount".into())] ); } + +struct SetOperationCatalog; + +impl CatalogProvider for SetOperationCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + match table.table.as_str() { + "users" => Some(vec!["id".into(), "name".into(), "email".into()]), + "other" => Some(vec!["a".into(), "b".into(), "c".into(), "d".into()]), + "ext_a" => Some(vec!["col_x".into(), "col_y".into()]), + _ => None, + } + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + +#[test] +fn set_operation_expands_leading_star_before_positional_merge() { + let sql = "SELECT * FROM users UNION ALL SELECT a, b, c FROM other"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 3); + assert_eq!(result.columns.mappings[0].target.column, "id"); + assert_eq!(result.columns.mappings[1].target.column, "name"); + assert_eq!(result.columns.mappings[2].target.column, "email"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("other".into(), "a".into()), ("users".into(), "id".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} + +#[test] +fn set_operation_preserves_non_leading_star_contribution() { + let sql = "SELECT id, * FROM users UNION ALL SELECT a, b, c, d FROM other"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 4); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("other".into(), "b".into()), ("users".into(), "id".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[3]), + vec![ + ("other".into(), "d".into()), + ("users".into(), "email".into()) + ] + ); +} + +#[test] +fn set_operation_branches_survive_cte_and_derived_boundaries() { + let sql = "WITH combined AS (SELECT * FROM users UNION ALL SELECT a, b, c FROM other) \ + SELECT * FROM (SELECT * FROM combined) derived"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 3); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} + +#[test] +fn non_leading_star_in_right_set_branch_contributes_to_named_output() { + let sql = "WITH lit AS (SELECT 1 AS col_a), \ + u AS (SELECT col_a FROM lit UNION ALL SELECT * FROM ext_a) \ + SELECT col_a FROM u"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("ext_a".into(), "col_x".into())] + ); +} + +#[test] +fn named_lookup_through_set_operation_cte_keeps_all_branches() { + let sql = "WITH combined AS (SELECT * FROM users UNION ALL SELECT a, b, c FROM other) \ + SELECT name FROM combined"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![ + ("other".into(), "b".into()), + ("users".into(), "name".into()) + ] + ); +} diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index cd2e6c1..518730e 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -85,6 +85,12 @@ fn recursive_cte_base_case() { match &m.sources[0] { ColumnOrigin::Recursive { base_sources } => { assert!(!base_sources.is_empty()); + assert!(base_sources.iter().all(|source| { + !matches!( + source, + ColumnOrigin::Concrete { table, .. } if table.table == "cte" + ) + })); match &base_sources[0] { ColumnOrigin::Concrete { table, column } => { assert_eq!(table.table, "t"); @@ -165,6 +171,58 @@ fn union_all_columns() { ); } +#[test] +fn union_keeps_left_names_and_merges_explicit_columns_positionally() { + let sql = "SELECT a AS left_name, b AS second_name FROM t1 \ + UNION ALL SELECT c AS right_name, d AS other_name FROM t2"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 2); + assert_eq!(result.columns.mappings[0].target.column, "left_name"); + assert_eq!(result.columns.mappings[1].target.column, "second_name"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("t1".into(), "a".into()), ("t2".into(), "c".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("t1".into(), "b".into()), ("t2".into(), "d".into())] + ); +} + +#[test] +fn nested_union_preserves_outer_left_names_and_all_branch_sources() { + let sql = "SELECT a AS first_name FROM t1 \ + UNION ALL SELECT b AS second_name FROM t2 \ + UNION ALL SELECT c AS third_name FROM t3"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!(result.columns.mappings[0].target.column, "first_name"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![ + ("t1".into(), "a".into()), + ("t2".into(), "b".into()), + ("t3".into(), "c".into()) + ] + ); +} + +#[test] +fn union_transform_prefers_aggregation_across_branches() { + let sql = "SELECT SUM(a) AS value FROM t1 UNION ALL SELECT b AS other_value FROM t2"; + let result = analyze_one(sql); + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!(result.columns.mappings[0].target.column, "value"); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("t1".into(), "a".into()), ("t2".into(), "b".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + TransformKind::Aggregation + ); +} + #[test] fn union_inside_cte() { let sql = "\ From dd389415c0270f4ef38af619eb7a7f96202f7ff9 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 04:35:36 +0000 Subject: [PATCH 05/14] fix: preserve unknown set operation shapes --- sqllineage/src/lib.rs | 8 +- sqllineage/src/resolve/mod.rs | 213 ++++++++++++++++++++++++++++++++-- sqllineage/src/types.rs | 2 +- sqllineage/tests/catalog.rs | 19 +++ sqllineage/tests/cte.rs | 127 +++++++++++++++++++- 5 files changed, 357 insertions(+), 12 deletions(-) diff --git a/sqllineage/src/lib.rs b/sqllineage/src/lib.rs index 80b2b58..09ff161 100644 --- a/sqllineage/src/lib.rs +++ b/sqllineage/src/lib.rs @@ -59,7 +59,9 @@ use sqlparser::parser::Parser; /// /// # Errors /// -/// Returns [`ParseError`] if the SQL string cannot be parsed. +/// Returns [`ParseError`] if the SQL string cannot be parsed or fails semantic +/// validation during lineage analysis (for example, exact set-operation +/// branches with different column counts). #[allow(clippy::needless_pass_by_value)] pub fn analyze(sql: &str, opts: AnalyzeOptions) -> Result, ParseError> { let dialect = opts.dialect.to_sqlparser_dialect(); @@ -68,12 +70,12 @@ pub fn analyze(sql: &str, opts: AnalyzeOptions) -> Result, Pa })?; let catalog = opts.catalog; - Ok(statements + statements .iter() .map(|stmt| { let builder = build::LineageBuilder::new(opts.normalize_case); let (raw_graph, warnings, statement_type) = builder.build(stmt); resolve::resolve(raw_graph, catalog.as_deref(), warnings, statement_type) }) - .collect()) + .collect() } diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index f3f3a28..b086054 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -9,7 +9,7 @@ use crate::graph::node::{NodeId, RawNode}; use crate::graph::scope::{Binding, OutputPlan, ScopeTree}; use crate::types::{ AnalyzeResult, CatalogProvider, ColumnLineage, ColumnMapping, ColumnOrigin, ColumnRef, - StatementType, TableRef, TransformKind, Warning, WarningKind, + ParseError, StatementType, TableRef, TransformKind, Warning, WarningKind, }; /// Resolve `RawGraph` into `AnalyzeResult`. @@ -19,17 +19,17 @@ pub(crate) fn resolve( catalog: Option<&dyn CatalogProvider>, mut warnings: Vec, statement_type: StatementType, -) -> AnalyzeResult { +) -> Result { graph.tables.inputs.sort(); graph.tables.inputs.dedup(); if graph.nodes.is_empty() { - return AnalyzeResult { + return Ok(AnalyzeResult { statement_type, tables: graph.tables, columns: ColumnLineage::default(), warnings, - }; + }); } if topo::topological_sort(&graph.nodes, &graph.edges).is_err() { @@ -37,14 +37,16 @@ pub(crate) fn resolve( kind: WarningKind::UnexpectedCycle, location: None, }); - return AnalyzeResult { + return Ok(AnalyzeResult { statement_type, tables: graph.tables, columns: ColumnLineage::default(), warnings, - }; + }); } + validate_set_arities(&graph, catalog)?; + let mut incoming: Vec> = vec![vec![]; graph.nodes.len()]; for (idx, edge) in graph.edges.iter().enumerate() { incoming[edge.to].push(idx); @@ -63,12 +65,119 @@ pub(crate) fn resolve( catalog, ); - AnalyzeResult { + Ok(AnalyzeResult { statement_type, tables: graph.tables, columns: ColumnLineage { mappings }, warnings, + }) +} + +const SET_ARITY_ERROR_PREFIX: &str = "set operation arity mismatch"; + +fn validate_set_arities( + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, +) -> Result<(), ParseError> { + let mut active = HashSet::new(); + let _ = scope_arity(ScopeTree::root(), graph, catalog, &mut active)?; + Ok(()) +} + +/// Return an exact output width when every star in a scope can be expanded; +/// otherwise return `None` and leave the eventual merge conservative. +fn scope_arity( + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Result, ParseError> { + if !active.insert(scope) { + return Ok(None); + } + let result = match graph.scopes.output_plan(scope).clone() { + OutputPlan::Projection => { + let mut width = 0; + let mut exact = true; + for col in graph.scopes.output_columns(scope) { + if let RawNode::Star { + table, + scope: star_scope, + } = &graph.nodes[col.node_id] + { + match star_arity(table.as_ref(), *star_scope, graph, catalog, active)? { + Some(star_width) => width += star_width, + None => exact = false, + } + } else { + width += 1; + } + } + exact.then_some(width) + } + OutputPlan::Delegate(child) => scope_arity(child, graph, catalog, active)?, + OutputPlan::SetOperation { left, right, .. } => { + let left_width = scope_arity(left, graph, catalog, active)?; + let right_width = scope_arity(right, graph, catalog, active)?; + match (left_width, right_width) { + (Some(left), Some(right)) if left != right => { + return Err(ParseError { + message: format!( + "{SET_ARITY_ERROR_PREFIX}: left has {left} columns, right has {right} columns" + ), + }); + } + (Some(width), Some(_)) => Some(width), + _ => None, + } + } + }; + active.remove(&scope); + Ok(result) +} + +fn star_arity( + table: Option<&TableRef>, + scope: usize, + graph: &RawGraph, + catalog: Option<&dyn CatalogProvider>, + active: &mut HashSet, +) -> Result, ParseError> { + if let Some(table) = table { + let binding = graph.scopes.lookup(scope, &table.table).cloned(); + return match binding { + Some(Binding::Cte(child) | Binding::DerivedTable(child)) => { + scope_arity(child, graph, catalog, active) + } + _ => Ok(catalog + .and_then(|catalog| catalog.list_columns(table)) + .map(|columns| columns.len())), + }; + } + + let bindings = effective_bindings(scope, graph); + let mut width = 0; + for (_, binding) in bindings { + let binding_width = match binding { + Binding::Table(table) => catalog + .and_then(|catalog| catalog.list_columns(&table)) + .map(|columns| columns.len()), + Binding::Cte(child) | Binding::DerivedTable(child) => { + scope_arity(child, graph, catalog, active)? + } + }; + match binding_width { + Some(binding_width) => width += binding_width, + None => return Ok(None), + } } + for &child in graph.scopes.anonymous_derived(scope) { + match scope_arity(child, graph, catalog, active)? { + Some(child_width) => width += child_width, + None => return Ok(None), + } + } + Ok(Some(width)) } fn effective_bindings(scope: usize, graph: &RawGraph) -> Vec<(String, Binding)> { @@ -170,6 +279,17 @@ fn resolve_scope_mappings( } else { resolve_scope_mappings(right, graph, resolved, incoming, output_table, catalog) }; + // A wildcard without catalog metadata is a variable-width slot. + // Positional alignment at or after it would fabricate an ordinal + // and lose the remaining branch columns. Merge only the exact + // prefix before the first wildcard, then retain both tails and + // expose the Wildcard origin. + if !recursive + && (mappings_have_unknown_shape(&left_mappings) + || mappings_have_unknown_shape(&right_mappings)) + { + return merge_unknown_shape_mappings(left_mappings, right_mappings); + } // The branch resolvers expand stars independently. Preserve the // left branch's names and order, as SQL set operations do. let mut merged = Vec::with_capacity(left_mappings.len()); @@ -201,6 +321,85 @@ fn resolve_scope_mappings( } } +fn mappings_have_unknown_shape(mappings: &[ColumnMapping]) -> bool { + mappings.iter().any(|mapping| { + mapping.sources.iter().any(|source| match source { + ColumnOrigin::Wildcard { .. } => true, + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .any(|source| matches!(source, ColumnOrigin::Wildcard { .. })), + _ => false, + }) + }) +} + +fn merge_unknown_shape_mappings( + left: Vec, + right: Vec, +) -> Vec { + let left_barrier = first_unknown_mapping(&left).unwrap_or(left.len()); + let right_barrier = first_unknown_mapping(&right).unwrap_or(right.len()); + let prefix_len = left_barrier.min(right_barrier); + let left_unknown = wildcard_sources(&left); + let right_unknown = wildcard_sources(&right); + let mut merged = Vec::with_capacity(left.len() + right.len() - prefix_len); + + for (left_mapping, right_mapping) in left.iter().zip(right.iter()).take(prefix_len) { + let mut sources = left_mapping.sources.clone(); + sources.extend(right_mapping.sources.clone()); + merged.push(ColumnMapping { + target: left_mapping.target.clone(), + sources, + transform: merge_transform(&left_mapping.transform, &right_mapping.transform), + }); + } + merged.extend( + left.into_iter() + .skip(prefix_len) + .map(|mapping| append_unknown_sources(mapping, &right_unknown)), + ); + merged.extend( + right + .into_iter() + .skip(prefix_len) + .map(|mapping| append_unknown_sources(mapping, &left_unknown)), + ); + merged +} + +fn first_unknown_mapping(mappings: &[ColumnMapping]) -> Option { + mappings + .iter() + .position(|mapping| mappings_have_unknown_shape(std::slice::from_ref(mapping))) +} + +fn wildcard_sources(mappings: &[ColumnMapping]) -> Vec { + let mut sources = Vec::new(); + for mapping in mappings { + for source in &mapping.sources { + match source { + ColumnOrigin::Wildcard { .. } => sources.push(source.clone()), + ColumnOrigin::Recursive { base_sources } => sources.extend( + base_sources + .iter() + .filter(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + ), + _ => {} + } + } + } + sources +} + +fn append_unknown_sources( + mut mapping: ColumnMapping, + unknown_sources: &[ColumnOrigin], +) -> ColumnMapping { + mapping.sources.extend(unknown_sources.iter().cloned()); + mapping +} + fn merge_transform(left: &TransformKind, right: &TransformKind) -> TransformKind { if matches!(left, TransformKind::Aggregation) || matches!(right, TransformKind::Aggregation) { TransformKind::Aggregation diff --git a/sqllineage/src/types.rs b/sqllineage/src/types.rs index 5f392b9..dd1e992 100644 --- a/sqllineage/src/types.rs +++ b/sqllineage/src/types.rs @@ -210,7 +210,7 @@ pub enum Dialect { BigQuery, } -/// Error returned when SQL parsing fails. +/// Error returned when SQL parsing or semantic validation fails. #[derive(Debug, Clone)] pub struct ParseError { pub message: String, diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 872655f..61b6660 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -316,3 +316,22 @@ fn named_lookup_through_set_operation_cte_keeps_all_branches() { ] ); } + +#[test] +fn catalog_known_set_arity_mismatch_is_an_analysis_error() { + let result = analyze( + "SELECT * FROM users UNION ALL SELECT a, b, c, d FROM other", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ); + let error = match result { + Ok(_) => panic!("catalog-known arity mismatch should not be truncated"), + Err(error) => error, + }; + assert_eq!( + error.message, + "set operation arity mismatch: left has 3 columns, right has 4 columns" + ); +} diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index 518730e..3f7f2ab 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -1,7 +1,19 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping, table}; -use sqllineage::{ColumnOrigin, TableRef, TransformKind}; +use sqllineage::{ColumnMapping, ColumnOrigin, TableRef, TransformKind}; + +fn has_wildcard_from(mapping: &ColumnMapping, table_name: &str) -> bool { + mapping.sources.iter().any(|source| { + match source { + ColumnOrigin::Wildcard { table } => table.table == table_name, + ColumnOrigin::Recursive { base_sources } => base_sources.iter().any(|source| { + matches!(source, ColumnOrigin::Wildcard { table } if table.table == table_name) + }), + _ => false, + } + }) +} #[test] fn single_cte() { @@ -223,6 +235,119 @@ fn union_transform_prefers_aggregation_across_branches() { ); } +#[test] +fn unknown_leading_star_is_preserved_without_catalog() { + let result = analyze_one("SELECT * FROM unknown_left UNION ALL SELECT a, b FROM known"); + assert_eq!(result.columns.mappings.len(), 3); + match &result.columns.mappings[0].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "unknown_left"), + other => panic!("expected wildcard, got {other:?}"), + } + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[2], + "unknown_left" + )); +} + +#[test] +fn unknown_non_leading_star_does_not_drop_known_branch_columns() { + let result = analyze_one("SELECT id, * FROM unknown_left UNION ALL SELECT a, b, c FROM known"); + assert_eq!(result.columns.mappings.len(), 4); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Concrete { .. }, ColumnOrigin::Concrete { .. }] + )); + match &result.columns.mappings[1].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "unknown_left"), + other => panic!("expected wildcard, got {other:?}"), + } + assert!(has_wildcard_from( + &result.columns.mappings[2], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[3], + "unknown_left" + )); + assert_eq!(result.columns.mappings[2].target.column, "b"); + assert_eq!(result.columns.mappings[3].target.column, "c"); +} + +#[test] +fn unknown_right_star_marks_known_left_tail_as_unresolved() { + let result = analyze_one("SELECT a, b, c FROM known UNION ALL SELECT * FROM unknown_right"); + assert_eq!(result.columns.mappings.len(), 4); + for mapping in &result.columns.mappings[..3] { + assert!(has_wildcard_from(mapping, "unknown_right")); + } +} + +#[test] +fn unknown_stars_on_both_set_branches_are_both_retained() { + let result = analyze_one("SELECT * FROM unknown_left UNION ALL SELECT * FROM unknown_right"); + assert_eq!(result.columns.mappings.len(), 2); + for (mapping, table_name) in result + .columns + .mappings + .iter() + .zip(["unknown_left", "unknown_right"]) + { + assert!(has_wildcard_from(mapping, table_name)); + assert!(has_wildcard_from( + mapping, + if table_name == "unknown_left" { + "unknown_right" + } else { + "unknown_left" + } + )); + } +} + +#[test] +fn nested_unknown_set_keeps_every_branch_mapping() { + let result = analyze_one( + "SELECT * FROM unknown_left UNION ALL SELECT a FROM known UNION ALL SELECT * FROM unknown_right", + ); + assert_eq!(result.columns.mappings.len(), 3); + assert!(matches!( + result.columns.mappings[0].sources[0], + ColumnOrigin::Wildcard { .. } + )); + assert!(matches!( + result.columns.mappings[2].sources[0], + ColumnOrigin::Wildcard { .. } + )); + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_left" + )); + assert!(has_wildcard_from( + &result.columns.mappings[1], + "unknown_right" + )); +} + +#[test] +fn exact_set_arity_mismatch_is_an_analysis_error() { + let result = sqllineage::analyze( + "SELECT a FROM t1 UNION ALL SELECT b, c FROM t2", + sqllineage::AnalyzeOptions::default(), + ); + let error = match result { + Ok(_) => panic!("exact arity mismatch should not be truncated"), + Err(error) => error, + }; + assert_eq!( + error.message, + "set operation arity mismatch: left has 1 columns, right has 2 columns" + ); +} + #[test] fn union_inside_cte() { let sql = "\ From 2e01ae3147f628c6b49ff3a6490102483fd40771 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 05:18:21 +0000 Subject: [PATCH 06/14] fix: preserve lineage through projection scopes --- sqllineage/src/resolve/mod.rs | 147 ++++++++++++++++++++++++---------- sqllineage/tests/catalog.rs | 93 +++++++++++++++++++++ 2 files changed, 196 insertions(+), 44 deletions(-) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index b086054..f1d1eb8 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -221,15 +221,19 @@ fn resolve_scope_mappings( match &graph.nodes[col.node_id] { RawNode::Output { name, .. } => { let mut visited = HashSet::new(); - let (sources, edge_kinds, has_back) = collect_output_sources( - col.node_id, - graph, - resolved, - incoming, - &mut visited, - catalog, + let (sources, edge_kinds, has_back, inherited_transform) = + collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + &mut visited, + catalog, + ); + let transform = merge_transform( + &derive_transform(&graph.nodes[col.node_id], &edge_kinds), + &inherited_transform, ); - let transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); mappings.push(ColumnMapping { target: ColumnRef { table: output_table.cloned(), @@ -518,7 +522,7 @@ fn expand_scope_columns( ); } else { let mut visited = HashSet::new(); - let (sources, edge_kinds, _) = collect_output_sources( + let (sources, edge_kinds, _, inherited_transform) = collect_output_sources( col.node_id, graph, resolved, @@ -526,7 +530,10 @@ fn expand_scope_columns( &mut visited, catalog, ); - let transform = derive_transform(&graph.nodes[col.node_id], &edge_kinds); + let transform = merge_transform( + &derive_transform(&graph.nodes[col.node_id], &edge_kinds), + &inherited_transform, + ); mappings.push(ColumnMapping { target: ColumnRef { table: output_table.cloned(), @@ -549,10 +556,10 @@ fn scope_column_sources( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, -) -> (Vec, Vec, bool) { +) -> (Vec, Vec, bool, TransformKind) { let mappings = resolve_scope_mappings(scope, graph, resolved, incoming, None, catalog); let Some(mapping) = mappings.get(index) else { - return (vec![], vec![], false); + return (vec![], vec![], false, TransformKind::Direct); }; let mut sources = Vec::new(); let mut has_back = false; @@ -565,7 +572,7 @@ fn scope_column_sources( source => sources.push(source.clone()), } } - (sources, vec![], has_back) + (sources, vec![], has_back, mapping.transform.clone()) } fn collect_output_sources( @@ -575,14 +582,15 @@ fn collect_output_sources( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, -) -> (Vec, Vec, bool) { +) -> (Vec, Vec, bool, TransformKind) { if !visited.insert(node_id) { - return (vec![], vec![], false); + return (vec![], vec![], false, TransformKind::Direct); } let mut sources = Vec::new(); let mut kinds = Vec::new(); let mut has_back = false; + let mut inherited_transform = TransformKind::Direct; for &edge_idx in &incoming[node_id] { let edge = &graph.edges[edge_idx]; @@ -590,16 +598,17 @@ fn collect_output_sources( has_back = true; continue; } - let (sub_sources, sub_back) = + let (sub_sources, sub_back, sub_transform) = collect_leaf_origins(edge.from, graph, resolved, incoming, visited, catalog); for _ in &sub_sources { kinds.push(edge.kind.clone()); } sources.extend(sub_sources); has_back |= sub_back; + inherited_transform = merge_transform(&inherited_transform, &sub_transform); } - (sources, kinds, has_back) + (sources, kinds, has_back, inherited_transform) } fn collect_leaf_origins( @@ -609,11 +618,11 @@ fn collect_leaf_origins( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, -) -> (Vec, bool) { - if let Some((sources, has_back)) = +) -> (Vec, bool, TransformKind) { + if let Some((sources, has_back, transform)) = resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog) { - return (sources, has_back); + return (sources, has_back, transform); } if let Some((target_output, scope)) = find_cte_redirect(node_id, graph) { @@ -624,24 +633,24 @@ fn collect_leaf_origins( .position(|c| c.node_id == target_output) && !matches!(graph.scopes.output_plan(scope), OutputPlan::Projection) { - let (sources, _, has_back) = + let (sources, _, has_back, transform) = scope_column_sources(scope, index, graph, resolved, incoming, catalog); - return (sources, has_back); + return (sources, has_back, transform); } - let (sources, _, has_back) = + let (sources, _, has_back, transform) = collect_output_sources(target_output, graph, resolved, incoming, visited, catalog); - return (sources, has_back); + return (sources, has_back, transform); } if let RawNode::Output { .. } = &graph.nodes[node_id] { - let (sources, _, has_back) = + let (sources, _, has_back, transform) = collect_output_sources(node_id, graph, resolved, incoming, visited, catalog); - (sources, has_back) + (sources, has_back, transform) } else { let origin = resolve_node(node_id, graph, resolved, incoming, visited, catalog); match origin { - Some(o) => (vec![o], false), - None => (vec![], false), + Some(o) => (vec![o], false, TransformKind::Direct), + None => (vec![], false, TransformKind::Direct), } } } @@ -655,7 +664,7 @@ fn resolve_named_scope_reference( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, -) -> Option<(Vec, bool)> { +) -> Option<(Vec, bool, TransformKind)> { let (name, qualifier, scope) = match &graph.nodes[node_id] { RawNode::Ref { name, @@ -668,30 +677,48 @@ fn resolve_named_scope_reference( let binding = qualifier .and_then(|qualifier| graph.scopes.lookup(scope, qualifier).cloned()) .or_else(|| find_single_binding(scope, graph)); - let target_scope = match binding { - Some(Binding::Cte(target) | Binding::DerivedTable(target)) - if !matches!(graph.scopes.output_plan(target), OutputPlan::Projection) => - { - target - } - _ => return None, + let Some(Binding::Cte(target_scope) | Binding::DerivedTable(target_scope)) = binding else { + return None; }; let mappings = resolve_scope_mappings(target_scope, graph, resolved, incoming, None, catalog); - let mapping = mappings + if let Some(mapping) = mappings .iter() - .find(|mapping| mapping.target.column == *name)?; - let mut sources = Vec::new(); + .find(|mapping| mapping.target.column == *name) + { + let (sources, has_back) = flatten_mapping_sources(&mapping.sources); + return Some((sources, has_back, mapping.transform.clone())); + } + // A catalog-less star has no individual named mapping. Returning its + // wildcard origin is safer than falling through to a fabricated concrete + // source for a named reference through the CTE/derived scope. + let wildcard_sources = mappings + .iter() + .flat_map(|mapping| mapping.sources.iter()) + .filter_map(|source| match source { + ColumnOrigin::Wildcard { .. } => Some(source.clone()), + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .find(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + _ => None, + }) + .collect::>(); + (!wildcard_sources.is_empty()).then_some((wildcard_sources, false, TransformKind::Direct)) +} + +fn flatten_mapping_sources(sources: &[ColumnOrigin]) -> (Vec, bool) { + let mut flattened = Vec::new(); let mut has_back = false; - for source in &mapping.sources { + for source in sources { match source { ColumnOrigin::Recursive { base_sources } => { - sources.extend(base_sources.clone()); + flattened.extend(base_sources.clone()); has_back = true; } - source => sources.push(source.clone()), + source => flattened.push(source.clone()), } } - Some((sources, has_back)) + (flattened, has_back) } fn resolve_node( @@ -883,13 +910,45 @@ fn resolve_through_scope( visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, ) -> Option { + // Resolve through the same expanded output mappings used by the public + // projection path. This is important for a qualified CTE/derived + // reference whose name was introduced by a catalog-expanded star. + let mappings = resolve_scope_mappings(target_scope, graph, resolved, incoming, None, catalog); + if let Some(mapping) = mappings + .iter() + .find(|mapping| mapping.target.column == column_name) + { + let (origins, has_back) = flatten_mapping_sources(&mapping.sources); + return if has_back { + Some(ColumnOrigin::Recursive { + base_sources: origins, + }) + } else { + origins.into_iter().next() + }; + } + if let Some(source) = mappings + .iter() + .flat_map(|mapping| mapping.sources.iter()) + .find_map(|source| match source { + ColumnOrigin::Wildcard { .. } => Some(source.clone()), + ColumnOrigin::Recursive { base_sources } => base_sources + .iter() + .find(|source| matches!(source, ColumnOrigin::Wildcard { .. })) + .cloned(), + _ => None, + }) + { + return Some(source); + } + if let Some(col) = graph .scopes .output_columns(target_scope) .iter() .find(|c| c.name == column_name) { - let (origins, _, has_back) = + let (origins, _, has_back, _) = collect_output_sources(col.node_id, graph, resolved, incoming, visited, catalog); if has_back { Some(ColumnOrigin::Recursive { diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 61b6660..0d8d6e2 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -317,6 +317,99 @@ fn named_lookup_through_set_operation_cte_keeps_all_branches() { ); } +#[test] +fn named_lookup_through_projection_cte_and_derived_star_chain() { + let sql = "WITH base AS (SELECT * FROM users), wrapped AS (SELECT * FROM base) \ + SELECT name FROM (SELECT * FROM wrapped) derived"; + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("users".into(), "name".into())] + ); +} + +#[test] +fn named_lookup_through_unknown_projection_star_is_indeterminate() { + let result = analyze( + "WITH base AS (SELECT * FROM unknown) SELECT name FROM base", + AnalyzeOptions::default(), + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert!(matches!( + result.columns.mappings[0].sources.as_slice(), + [ColumnOrigin::Wildcard { table }] if table.table == "unknown" + )); +} + +#[test] +fn named_lookup_through_projection_preserves_inner_transform() { + let result = analyze( + "WITH aggregated AS (SELECT SUM(amount) AS total FROM orders) \ + SELECT total FROM aggregated", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("orders".into(), "amount".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + sqllineage::TransformKind::Aggregation + ); +} + +#[test] +fn named_lookup_through_projection_preserves_inner_expression_transform() { + let result = analyze( + "WITH transformed AS (SELECT amount + 1 AS adjusted FROM orders) \ + SELECT adjusted FROM transformed", + AnalyzeOptions { + catalog: Some(Box::new(SetOperationCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("orders".into(), "amount".into())] + ); + assert_eq!( + result.columns.mappings[0].transform, + sqllineage::TransformKind::Expression + ); +} + #[test] fn catalog_known_set_arity_mismatch_is_an_analysis_error() { let result = analyze( From 23cbb469bd328bac35c60be93e47edbe3e776e40 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 08:12:50 +0000 Subject: [PATCH 07/14] style: apply current rustfmt --- sqllineage/src/build/expr.rs | 72 +++++++++++++++++++++++++------ sqllineage/src/build/statement.rs | 3 +- 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 2b2cb3e..34e4d4b 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -24,7 +24,10 @@ impl LineageBuilder { vec![node] } - Expr::Value(_) | Expr::TypedString { .. } | Expr::Wildcard(..) | Expr::QualifiedWildcard(..) => vec![], + Expr::Value(_) + | Expr::TypedString { .. } + | Expr::Wildcard(..) + | Expr::QualifiedWildcard(..) => vec![], Expr::Cast { expr, .. } | Expr::Nested(expr) @@ -49,7 +52,12 @@ impl LineageBuilder { Expr::Extract { expr, .. } => self.collect_ancestors(expr), - Expr::Trim { expr, trim_what, trim_characters, .. } => { + Expr::Trim { + expr, + trim_what, + trim_characters, + .. + } => { let mut v = self.collect_ancestors(expr); if let Some(what) = trim_what { v.extend(self.collect_ancestors(what)); @@ -62,7 +70,12 @@ impl LineageBuilder { v } - Expr::Substring { expr, substring_from, substring_for, .. } => { + Expr::Substring { + expr, + substring_from, + substring_for, + .. + } => { let mut v = self.collect_ancestors(expr); if let Some(from) = substring_from { v.extend(self.collect_ancestors(from)); @@ -73,7 +86,13 @@ impl LineageBuilder { v } - Expr::Overlay { expr, overlay_what, overlay_from, overlay_for, .. } => { + Expr::Overlay { + expr, + overlay_what, + overlay_from, + overlay_for, + .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(overlay_what)); v.extend(self.collect_ancestors(overlay_from)); @@ -89,17 +108,36 @@ impl LineageBuilder { v } - Expr::AtTimeZone { timestamp, time_zone } => { + Expr::AtTimeZone { + timestamp, + time_zone, + } => { let mut v = self.collect_ancestors(timestamp); v.extend(self.collect_ancestors(time_zone)); v } Expr::BinaryOp { left, right, .. } - | Expr::Like { expr: left, pattern: right, .. } - | Expr::ILike { expr: left, pattern: right, .. } - | Expr::SimilarTo { expr: left, pattern: right, .. } - | Expr::RLike { expr: left, pattern: right, .. } + | Expr::Like { + expr: left, + pattern: right, + .. + } + | Expr::ILike { + expr: left, + pattern: right, + .. + } + | Expr::SimilarTo { + expr: left, + pattern: right, + .. + } + | Expr::RLike { + expr: left, + pattern: right, + .. + } | Expr::IsDistinctFrom(left, right) | Expr::IsNotDistinctFrom(left, right) => { let mut v = self.collect_ancestors(left); @@ -113,7 +151,9 @@ impl LineageBuilder { v } - Expr::InUnnest { expr, array_expr, .. } => { + Expr::InUnnest { + expr, array_expr, .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(array_expr)); v @@ -192,7 +232,12 @@ impl LineageBuilder { ancestors } - Expr::Case { operand, conditions, else_result, .. } => { + Expr::Case { + operand, + conditions, + else_result, + .. + } => { let mut v = Vec::new(); if let Some(op) = operand { v.extend(self.collect_ancestors(op)); @@ -229,7 +274,9 @@ impl LineageBuilder { vec![] } - Expr::Between { expr, low, high, .. } => { + Expr::Between { + expr, low, high, .. + } => { let mut v = self.collect_ancestors(expr); v.extend(self.collect_ancestors(low)); v.extend(self.collect_ancestors(high)); @@ -251,7 +298,6 @@ impl LineageBuilder { | Expr::Interval(_) | Expr::Lambda(_) | Expr::MatchAgainst { .. } => vec![], - } } } diff --git a/sqllineage/src/build/statement.rs b/sqllineage/src/build/statement.rs index 0c2980d..587ce93 100644 --- a/sqllineage/src/build/statement.rs +++ b/sqllineage/src/build/statement.rs @@ -134,7 +134,8 @@ impl LineageBuilder { .unwrap_or_else(|| format!("col{i}")); let ancestors = self.collect_ancestors(expr); let kind = determine_edge_kind(expr); - let output = self.graph.add_output(col_name.clone(), kind.clone()); + let output = + self.graph.add_output(col_name.clone(), kind.clone()); for &anc in &ancestors { self.graph.add_edge(anc, output, kind.clone()); } From 9a0e37145d436b343b3c0b299713f89b31214149 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 08:12:50 +0000 Subject: [PATCH 08/14] fix: preserve leading set-operation output names --- sqllineage/src/resolve/mod.rs | 12 ++++++++++++ sqllineage/tests/cte.rs | 16 ++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index f1d1eb8..952a354 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -343,6 +343,18 @@ fn merge_unknown_shape_mappings( ) -> Vec { let left_barrier = first_unknown_mapping(&left).unwrap_or(left.len()); let right_barrier = first_unknown_mapping(&right).unwrap_or(right.len()); + + // A leading unknown star determines the output names for the set + // operation. When the left side is already a merged set (and therefore + // contains named slots contributed by an earlier operand), a right-only + // tail would publish names that the leading operand never declared. Keep + // the left candidates and their existing wildcard provenance; a direct + // two-branch `SELECT * UNION SELECT a, b` still retains the right branch + // names because there are no prior named slots to preserve. + if left_barrier == 0 && left.len() > 1 && right_barrier == right.len() { + return left; + } + let prefix_len = left_barrier.min(right_barrier); let left_unknown = wildcard_sources(&left); let right_unknown = wildcard_sources(&right); diff --git a/sqllineage/tests/cte.rs b/sqllineage/tests/cte.rs index 3f7f2ab..df470e6 100644 --- a/sqllineage/tests/cte.rs +++ b/sqllineage/tests/cte.rs @@ -332,6 +332,22 @@ fn nested_unknown_set_keeps_every_branch_mapping() { )); } +#[test] +fn leading_unknown_star_hides_nonleading_only_set_names() { + let result = analyze_one( + "SELECT * FROM unknown_source \ + UNION ALL SELECT id, amt AS total FROM known_table \ + UNION ALL SELECT id, fee FROM third_table", + ); + let names = result + .columns + .mappings + .iter() + .map(|mapping| mapping.target.column.as_str()) + .collect::>(); + assert_eq!(names, vec!["*", "id", "total"]); +} + #[test] fn exact_set_arity_mismatch_is_an_analysis_error() { let result = sqllineage::analyze( From 757e532cd505000e90b521d145deb1b2caa8d091 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 11:55:24 +0000 Subject: [PATCH 09/14] perf: memoize scope mappings during resolution --- sqllineage/src/resolve/mod.rs | 380 +++++++++++++++++++++++++++------- 1 file changed, 306 insertions(+), 74 deletions(-) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 952a354..08ab35a 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1,7 +1,10 @@ mod catalog; mod topo; -use std::collections::HashSet; +#[cfg(test)] +use std::cell::Cell; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; @@ -56,14 +59,26 @@ pub(crate) fn resolve( let root = ScopeTree::root(); let output_table = graph.tables.output.clone(); - let mappings = resolve_scope_mappings( + let mut mapping_cache = ScopeMappingCache::default(); + // Scope mappings are cached in canonical form without an output table. + // Internal CTE/derived references need that form, while the root output + // table is only presentation metadata and is attached once here. Keeping + // it out of the cache key avoids materializing the same scope once per + // output column and cannot alter source resolution. + let mut mappings = resolve_scope_mappings( root, &graph, &mut resolved, &incoming, - output_table.as_ref(), catalog, - ); + &mut mapping_cache, + ) + .iter() + .cloned() + .collect::>(); + for mapping in &mut mappings { + mapping.target.table.clone_from(&output_table); + } Ok(AnalyzeResult { statement_type, @@ -202,6 +217,31 @@ fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> } } +#[derive(Default)] +struct ScopeMappingCache { + entries: HashMap, +} + +enum ScopeMappingEntry { + Computing, + Resolved(Arc<[ColumnMapping]>), +} + +#[cfg(test)] +thread_local! { + static SCOPE_MAPPING_COMPUTATIONS: Cell = const { Cell::new(0) }; +} + +#[cfg(test)] +fn reset_scope_mapping_stats() { + SCOPE_MAPPING_COMPUTATIONS.with(|computations| computations.set(0)); +} + +#[cfg(test)] +fn scope_mapping_computations() -> usize { + SCOPE_MAPPING_COMPUTATIONS.with(Cell::get) +} + /// Resolve a scope's output plan into ordered mappings. Set-operation /// branches are resolved independently so catalog expansion happens before /// their positional merge. @@ -211,10 +251,26 @@ fn resolve_scope_mappings( graph: &RawGraph, resolved: &mut Vec>, incoming: &[Vec], - output_table: Option<&TableRef>, catalog: Option<&dyn CatalogProvider>, -) -> Vec { - match graph.scopes.output_plan(scope).clone() { + mapping_cache: &mut ScopeMappingCache, +) -> Arc<[ColumnMapping]> { + if let Some(entry) = mapping_cache.entries.get(&scope) { + return match entry { + ScopeMappingEntry::Resolved(mappings) => mappings.clone(), + // A recursive scope cannot safely publish a partially materialized + // result. Returning no mappings preserves the existing fallback to + // raw output resolution and, importantly, does not cache an + // incomplete entry as resolved. + ScopeMappingEntry::Computing => Arc::from([]), + }; + } + mapping_cache + .entries + .insert(scope, ScopeMappingEntry::Computing); + #[cfg(test)] + SCOPE_MAPPING_COMPUTATIONS.with(|computations| computations.set(computations.get() + 1)); + + let mappings = match graph.scopes.output_plan(scope).clone() { OutputPlan::Projection => { let mut mappings = Vec::new(); for col in graph.scopes.output_columns(scope) { @@ -229,6 +285,7 @@ fn resolve_scope_mappings( incoming, &mut visited, catalog, + mapping_cache, ); let transform = merge_transform( &derive_transform(&graph.nodes[col.node_id], &edge_kinds), @@ -236,7 +293,7 @@ fn resolve_scope_mappings( ); mappings.push(ColumnMapping { target: ColumnRef { - table: output_table.cloned(), + table: None, column: name.clone(), }, sources: if has_back { @@ -256,7 +313,7 @@ fn resolve_scope_mappings( resolved, incoming, catalog, - output_table, + mapping_cache, &mut mappings, &mut HashSet::new(), ), @@ -269,7 +326,10 @@ fn resolve_scope_mappings( mappings } OutputPlan::Delegate(child) => { - resolve_scope_mappings(child, graph, resolved, incoming, output_table, catalog) + resolve_scope_mappings(child, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect() } OutputPlan::SetOperation { left, @@ -277,11 +337,17 @@ fn resolve_scope_mappings( recursive, } => { let left_mappings = - resolve_scope_mappings(left, graph, resolved, incoming, output_table, catalog); + resolve_scope_mappings(left, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect::>(); let right_mappings = if recursive { Vec::new() } else { - resolve_scope_mappings(right, graph, resolved, incoming, output_table, catalog) + resolve_scope_mappings(right, graph, resolved, incoming, catalog, mapping_cache) + .iter() + .cloned() + .collect() }; // A wildcard without catalog metadata is a variable-width slot. // Positional alignment at or after it would fabricate an ordinal @@ -292,37 +358,43 @@ fn resolve_scope_mappings( && (mappings_have_unknown_shape(&left_mappings) || mappings_have_unknown_shape(&right_mappings)) { - return merge_unknown_shape_mappings(left_mappings, right_mappings); - } - // The branch resolvers expand stars independently. Preserve the - // left branch's names and order, as SQL set operations do. - let mut merged = Vec::with_capacity(left_mappings.len()); - for (idx, left_mapping) in left_mappings.into_iter().enumerate() { - let mut sources = left_mapping.sources; - let mut transform = left_mapping.transform.clone(); - if let Some(right_mapping) = right_mappings.get(idx) { - sources.extend(right_mapping.sources.clone()); - transform = merge_transform(&transform, &right_mapping.transform); - } - if recursive { - merged.push(ColumnMapping { - target: left_mapping.target, - sources: vec![ColumnOrigin::Recursive { - base_sources: sources, - }], - transform, - }); - } else { - merged.push(ColumnMapping { - target: left_mapping.target, - sources, - transform, - }); + merge_unknown_shape_mappings(left_mappings, right_mappings) + } else { + // The branch resolvers expand stars independently. Preserve the + // left branch's names and order, as SQL set operations do. + let mut merged = Vec::with_capacity(left_mappings.len()); + for (idx, left_mapping) in left_mappings.into_iter().enumerate() { + let mut sources = left_mapping.sources; + let mut transform = left_mapping.transform.clone(); + if let Some(right_mapping) = right_mappings.get(idx) { + sources.extend(right_mapping.sources.clone()); + transform = merge_transform(&transform, &right_mapping.transform); + } + if recursive { + merged.push(ColumnMapping { + target: left_mapping.target, + sources: vec![ColumnOrigin::Recursive { + base_sources: sources, + }], + transform, + }); + } else { + merged.push(ColumnMapping { + target: left_mapping.target, + sources, + transform, + }); + } } + merged } - merged } - } + }; + let mappings: Arc<[ColumnMapping]> = mappings.into(); + mapping_cache + .entries + .insert(scope, ScopeMappingEntry::Resolved(mappings.clone())); + mappings } fn mappings_have_unknown_shape(mappings: &[ColumnMapping]) -> bool { @@ -445,7 +517,7 @@ fn expand_star( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, - output_table: Option<&TableRef>, + mapping_cache: &mut ScopeMappingCache, mappings: &mut Vec, visited_scopes: &mut HashSet, ) { @@ -458,17 +530,17 @@ fn expand_star( resolved, incoming, catalog, - output_table, + mapping_cache, mappings, visited_scopes, ); } else { - mappings.push(wildcard_mapping(output_table, t.clone())); + mappings.push(wildcard_mapping(None, t.clone())); } } else { for (_, binding) in effective_bindings(scope, graph) { match binding { - Binding::Table(tref) => mappings.push(wildcard_mapping(output_table, tref)), + Binding::Table(tref) => mappings.push(wildcard_mapping(None, tref)), Binding::Cte(s) | Binding::DerivedTable(s) => { expand_scope_columns( s, @@ -476,7 +548,7 @@ fn expand_star( resolved, incoming, catalog, - output_table, + mapping_cache, mappings, visited_scopes, ); @@ -490,7 +562,7 @@ fn expand_star( resolved, incoming, catalog, - output_table, + mapping_cache, mappings, visited_scopes, ); @@ -506,7 +578,7 @@ fn expand_scope_columns( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, - output_table: Option<&TableRef>, + mapping_cache: &mut ScopeMappingCache, mappings: &mut Vec, visited_scopes: &mut HashSet, ) { @@ -514,9 +586,9 @@ fn expand_scope_columns( return; } if !matches!(graph.scopes.output_plan(scope_id), OutputPlan::Projection) { - let mut nested = - resolve_scope_mappings(scope_id, graph, resolved, incoming, output_table, catalog); - mappings.append(&mut nested); + let nested = + resolve_scope_mappings(scope_id, graph, resolved, incoming, catalog, mapping_cache); + mappings.extend(nested.iter().cloned()); return; } for col in graph.scopes.output_columns(scope_id) { @@ -528,7 +600,7 @@ fn expand_scope_columns( resolved, incoming, catalog, - output_table, + mapping_cache, mappings, visited_scopes, ); @@ -541,6 +613,7 @@ fn expand_scope_columns( incoming, &mut visited, catalog, + mapping_cache, ); let transform = merge_transform( &derive_transform(&graph.nodes[col.node_id], &edge_kinds), @@ -548,7 +621,7 @@ fn expand_scope_columns( ); mappings.push(ColumnMapping { target: ColumnRef { - table: output_table.cloned(), + table: None, column: col.name.clone(), }, sources, @@ -568,8 +641,9 @@ fn scope_column_sources( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> (Vec, Vec, bool, TransformKind) { - let mappings = resolve_scope_mappings(scope, graph, resolved, incoming, None, catalog); + let mappings = resolve_scope_mappings(scope, graph, resolved, incoming, catalog, mapping_cache); let Some(mapping) = mappings.get(index) else { return (vec![], vec![], false, TransformKind::Direct); }; @@ -594,6 +668,7 @@ fn collect_output_sources( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> (Vec, Vec, bool, TransformKind) { if !visited.insert(node_id) { return (vec![], vec![], false, TransformKind::Direct); @@ -610,8 +685,15 @@ fn collect_output_sources( has_back = true; continue; } - let (sub_sources, sub_back, sub_transform) = - collect_leaf_origins(edge.from, graph, resolved, incoming, visited, catalog); + let (sub_sources, sub_back, sub_transform) = collect_leaf_origins( + edge.from, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); for _ in &sub_sources { kinds.push(edge.kind.clone()); } @@ -630,9 +712,10 @@ fn collect_leaf_origins( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> (Vec, bool, TransformKind) { if let Some((sources, has_back, transform)) = - resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog) + resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog, mapping_cache) { return (sources, has_back, transform); } @@ -645,21 +728,50 @@ fn collect_leaf_origins( .position(|c| c.node_id == target_output) && !matches!(graph.scopes.output_plan(scope), OutputPlan::Projection) { - let (sources, _, has_back, transform) = - scope_column_sources(scope, index, graph, resolved, incoming, catalog); + let (sources, _, has_back, transform) = scope_column_sources( + scope, + index, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); return (sources, has_back, transform); } - let (sources, _, has_back, transform) = - collect_output_sources(target_output, graph, resolved, incoming, visited, catalog); + let (sources, _, has_back, transform) = collect_output_sources( + target_output, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); return (sources, has_back, transform); } if let RawNode::Output { .. } = &graph.nodes[node_id] { - let (sources, _, has_back, transform) = - collect_output_sources(node_id, graph, resolved, incoming, visited, catalog); + let (sources, _, has_back, transform) = collect_output_sources( + node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); (sources, has_back, transform) } else { - let origin = resolve_node(node_id, graph, resolved, incoming, visited, catalog); + let origin = resolve_node( + node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); match origin { Some(o) => (vec![o], false, TransformKind::Direct), None => (vec![], false, TransformKind::Direct), @@ -676,6 +788,7 @@ fn resolve_named_scope_reference( resolved: &mut Vec>, incoming: &[Vec], catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option<(Vec, bool, TransformKind)> { let (name, qualifier, scope) = match &graph.nodes[node_id] { RawNode::Ref { @@ -692,7 +805,14 @@ fn resolve_named_scope_reference( let Some(Binding::Cte(target_scope) | Binding::DerivedTable(target_scope)) = binding else { return None; }; - let mappings = resolve_scope_mappings(target_scope, graph, resolved, incoming, None, catalog); + let mappings = resolve_scope_mappings( + target_scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); if let Some(mapping) = mappings .iter() .find(|mapping| mapping.target.column == *name) @@ -740,6 +860,7 @@ fn resolve_node( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { if let Some(ref origin) = resolved[node_id] { return Some(origin.clone()); @@ -760,7 +881,14 @@ fn resolve_node( }), Some(Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope)) => { resolve_through_scope( - name, cte_scope, graph, resolved, incoming, visited, catalog, + name, + cte_scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, ) } None => Some(ColumnOrigin::Concrete { @@ -769,13 +897,29 @@ fn resolve_node( }), } } else { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited, catalog) + resolve_unqualified( + name, + *scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) } } - RawNode::Unqualified { name, scope } => { - resolve_unqualified(name, *scope, graph, resolved, incoming, visited, catalog) - } + RawNode::Unqualified { name, scope } => resolve_unqualified( + name, + *scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), RawNode::Star { table, .. } => table .as_ref() @@ -835,6 +979,7 @@ fn find_single_binding(scope: usize, graph: &RawGraph) -> Option { } } +#[allow(clippy::too_many_arguments)] fn resolve_unqualified( name: &str, scope: usize, @@ -843,6 +988,7 @@ fn resolve_unqualified( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { resolve_from_bindings( name, @@ -852,9 +998,11 @@ fn resolve_unqualified( incoming, visited, catalog, + mapping_cache, ) } +#[allow(clippy::too_many_arguments)] fn resolve_from_bindings( name: &str, bindings: &[(String, Binding)], @@ -863,6 +1011,7 @@ fn resolve_from_bindings( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { if bindings.len() == 1 { let (_, binding) = &bindings[0]; @@ -872,7 +1021,14 @@ fn resolve_from_bindings( column: name.to_string(), }), Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope) => resolve_through_scope( - name, *cte_scope, graph, resolved, incoming, visited, catalog, + name, + *cte_scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, ), } } else if bindings.is_empty() { @@ -892,7 +1048,14 @@ fn resolve_from_bindings( .any(|c| c.name == name) { return resolve_through_scope( - name, *s, graph, resolved, incoming, visited, catalog, + name, + *s, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, ); } } @@ -913,6 +1076,7 @@ fn resolve_from_bindings( } } +#[allow(clippy::too_many_arguments)] fn resolve_through_scope( column_name: &str, target_scope: usize, @@ -921,11 +1085,19 @@ fn resolve_through_scope( incoming: &[Vec], visited: &mut HashSet, catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, ) -> Option { // Resolve through the same expanded output mappings used by the public // projection path. This is important for a qualified CTE/derived // reference whose name was introduced by a catalog-expanded star. - let mappings = resolve_scope_mappings(target_scope, graph, resolved, incoming, None, catalog); + let mappings = resolve_scope_mappings( + target_scope, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ); if let Some(mapping) = mappings .iter() .find(|mapping| mapping.target.column == column_name) @@ -960,8 +1132,15 @@ fn resolve_through_scope( .iter() .find(|c| c.name == column_name) { - let (origins, _, has_back, _) = - collect_output_sources(col.node_id, graph, resolved, incoming, visited, catalog); + let (origins, _, has_back, _) = collect_output_sources( + col.node_id, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ); if has_back { Some(ColumnOrigin::Recursive { base_sources: origins, @@ -1002,3 +1181,56 @@ fn derive_transform(node: &RawNode, edge_kinds: &[EdgeKind]) -> TransformKind { TransformKind::Direct } } + +#[cfg(test)] +mod tests { + use super::{reset_scope_mapping_stats, scope_mapping_computations}; + use crate::analyze; + use crate::types::{AnalyzeOptions, ColumnOrigin, Dialect}; + + #[test] + fn explicit_projection_reuses_scope_mappings() { + let columns = (0..10).map(|index| format!("c{index}")).collect::>(); + let names = columns.join(", "); + let base_projection = (0..10) + .map(|index| format!("id + 1 AS c{index}")) + .collect::>() + .join(", "); + let sql = format!( + "WITH base AS (SELECT {base_projection} FROM external_table), \ + cte0 AS (SELECT {names} FROM base), \ + cte1 AS (SELECT {names} FROM cte0), \ + cte2 AS (SELECT {names} FROM cte1), \ + cte3 AS (SELECT {names} FROM cte2), \ + cte4 AS (SELECT {names} FROM cte3) \ + SELECT {names} FROM cte4" + ); + + reset_scope_mapping_stats(); + let results = analyze( + &sql, + AnalyzeOptions { + dialect: Dialect::Generic, + ..Default::default() + }, + ) + .expect("explicit projection should resolve"); + + assert_eq!(results.len(), 1); + let mappings = &results[0].columns.mappings; + assert_eq!(mappings.len(), 10); + assert!(mappings.iter().all(|mapping| { + mapping.sources.iter().any(|source| { + matches!( + source, + ColumnOrigin::Concrete { table, column } + if table.table == "external_table" && column == "id" + ) + }) + })); + // One materialization per scope, independent of the ten requested + // columns. The exact scope count is an implementation detail, but it + // must remain bounded by the five wrappers plus the base and root. + assert!(scope_mapping_computations() <= 8); + } +} From a612e015910c1a9256a919f52b594bf3130c861e Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 12:41:10 +0000 Subject: [PATCH 10/14] fix: expand qualified alias stars from catalog --- sqllineage/src/resolve/mod.rs | 5 +++ sqllineage/tests/catalog.rs | 76 ++++++++++++++++++++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 08ab35a..7ca2185 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -164,6 +164,9 @@ fn star_arity( Some(Binding::Cte(child) | Binding::DerivedTable(child)) => { scope_arity(child, graph, catalog, active) } + Some(Binding::Table(actual_table)) => Ok(catalog + .and_then(|catalog| catalog.list_columns(&actual_table)) + .map(|columns| columns.len())), _ => Ok(catalog .and_then(|catalog| catalog.list_columns(table)) .map(|columns| columns.len())), @@ -534,6 +537,8 @@ fn expand_star( mappings, visited_scopes, ); + } else if let Some(Binding::Table(actual_table)) = binding { + mappings.push(wildcard_mapping(None, actual_table)); } else { mappings.push(wildcard_mapping(None, t.clone())); } diff --git a/sqllineage/tests/catalog.rs b/sqllineage/tests/catalog.rs index 0d8d6e2..c7b12ad 100644 --- a/sqllineage/tests/catalog.rs +++ b/sqllineage/tests/catalog.rs @@ -1,7 +1,7 @@ mod common; use common::find_mapping; -use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, TableRef, analyze}; +use sqllineage::{AnalyzeOptions, CatalogProvider, ColumnOrigin, Dialect, TableRef, analyze}; struct MockCatalog; @@ -35,6 +35,18 @@ impl CatalogProvider for EagerCatalog { } } +struct AliasCatalog; + +impl CatalogProvider for AliasCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "actual_table").then(|| vec!["id".into(), "event".into()]) + } + + fn resolve_column(&self, _column: &str, _candidates: &[TableRef]) -> Option { + None + } +} + fn opts_with_catalog() -> AnalyzeOptions { AnalyzeOptions { catalog: Some(Box::new(MockCatalog)), @@ -78,6 +90,68 @@ fn select_star_with_catalog_expands() { ); } +fn assert_qualified_alias_star_expands(sql: &str) { + let result = analyze( + sql, + AnalyzeOptions { + catalog: Some(Box::new(AliasCatalog)), + dialect: Dialect::Generic, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 2); + assert!(result.columns.mappings.iter().all(|mapping| { + mapping + .sources + .iter() + .all(|source| matches!(source, ColumnOrigin::Concrete { .. })) + })); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "id")), + vec![("actual_table".into(), "id".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "event")), + vec![("actual_table".into(), "event".into())] + ); +} + +#[test] +fn qualified_alias_star_uses_catalog_table_binding() { + for sql in [ + "SELECT a.* FROM actual_table AS a", + "WITH x AS (SELECT a.* FROM actual_table AS a) SELECT * FROM x", + ] { + assert_qualified_alias_star_expands(sql); + } +} + +#[test] +fn qualified_alias_star_without_catalog_keeps_actual_table_wildcard() { + let result = analyze( + "SELECT a.* FROM actual_table AS a", + AnalyzeOptions { + dialect: Dialect::Generic, + ..AnalyzeOptions::default() + }, + ) + .expect("parse") + .into_iter() + .next() + .unwrap(); + + assert_eq!(result.columns.mappings.len(), 1); + match &result.columns.mappings[0].sources[0] { + ColumnOrigin::Wildcard { table } => assert_eq!(table.table, "actual_table"), + other => panic!("expected Wildcard, got {other:?}"), + } +} + #[test] fn select_star_without_catalog_preserved() { let result = analyze("SELECT * FROM users", AnalyzeOptions::default()) From 968ed013853ef953bfd12198717c8d9587d5346e Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 14:18:41 +0000 Subject: [PATCH 11/14] fix: resolve compound field access ancestry --- sqllineage/src/build/expr.rs | 73 +++++++++++++++++++++++++++- sqllineage/tests/expr_coverage.rs | 81 ++++++++++++++++++++++++++++++- 2 files changed, 151 insertions(+), 3 deletions(-) diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 34e4d4b..ba6ecab 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -1,4 +1,4 @@ -use sqlparser::ast::{self, Expr, FunctionArguments, WindowType}; +use sqlparser::ast::{self, AccessExpr, Expr, FunctionArguments, Subscript, WindowType}; use crate::build::LineageBuilder; use crate::build::select::split_compound; @@ -198,7 +198,9 @@ impl LineageBuilder { v } - Expr::CompoundFieldAccess { root, .. } => self.collect_ancestors(root), + Expr::CompoundFieldAccess { root, access_chain } => { + self.collect_compound_field_ancestors(root, access_chain) + } Expr::JsonAccess { value, .. } => self.collect_ancestors(value), Expr::Function(func) => { @@ -300,6 +302,73 @@ impl LineageBuilder { | Expr::MatchAgainst { .. } => vec![], } } + + /// Collect the physical column at the root of a structured access chain. + /// + /// `base.items[0]` is ambiguous at the syntax level: `base` can be a + /// visible relation binding, in which case `items` is its physical column, + /// or it can be an unqualified top-level column (`payload.items[0]`). The + /// scope binding, rather than rendered SQL text or dialect-specific names, + /// is the structural distinction between those cases. + fn collect_compound_field_ancestors( + &mut self, + root: &Expr, + access_chain: &[AccessExpr], + ) -> Vec { + let mut ancestors = match (root, access_chain.first()) { + (Expr::Identifier(binding), Some(AccessExpr::Dot(Expr::Identifier(field)))) + if self + .graph + .scopes + .lookup(self.current_scope, &binding.value) + .is_some() => + { + let node = self.graph.add_ref( + field.value.clone(), + Some(binding.value.clone()), + self.current_scope, + ); + vec![node] + } + (Expr::Identifier(column), Some(AccessExpr::Dot(Expr::Identifier(_)))) => { + vec![ + self.graph + .add_unqualified(column.value.clone(), self.current_scope), + ] + } + _ => self.collect_ancestors(root), + }; + + for access in access_chain { + if let AccessExpr::Subscript(subscript) = access { + ancestors.extend(self.collect_subscript_ancestors(subscript)); + } + } + ancestors + } + + fn collect_subscript_ancestors(&mut self, subscript: &Subscript) -> Vec { + match subscript { + Subscript::Index { index } => self.collect_ancestors(index), + Subscript::Slice { + lower_bound, + upper_bound, + stride, + } => { + let mut ancestors = Vec::new(); + if let Some(lower) = lower_bound { + ancestors.extend(self.collect_ancestors(lower)); + } + if let Some(upper) = upper_bound { + ancestors.extend(self.collect_ancestors(upper)); + } + if let Some(step) = stride { + ancestors.extend(self.collect_ancestors(step)); + } + ancestors + } + } + } } pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind { diff --git a/sqllineage/tests/expr_coverage.rs b/sqllineage/tests/expr_coverage.rs index 181e796..610178f 100644 --- a/sqllineage/tests/expr_coverage.rs +++ b/sqllineage/tests/expr_coverage.rs @@ -1,7 +1,21 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping}; -use sqllineage::TransformKind; +use sqllineage::{AnalyzeOptions, Dialect, TransformKind, analyze}; + +fn analyze_with_dialect(sql: &str, dialect: Dialect) -> sqllineage::AnalyzeResult { + analyze( + sql, + AnalyzeOptions { + dialect, + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .into_iter() + .next() + .unwrap_or_default() +} #[test] fn extract_year() { @@ -107,3 +121,68 @@ fn json_access() { let m = find_mapping(&result.columns.mappings, "val"); assert_eq!(concrete_sources(m), vec![("t".into(), "data".into())]); } + +#[test] +fn qualified_compound_field_access_uses_binding_column() { + let result = analyze_one("SELECT base.items_array[1] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} + +#[test] +fn compound_field_access_retains_column_dependent_index() { + let result = analyze_one("SELECT base.items_array[idx] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![ + ("actual_table".into(), "idx".into()), + ("actual_table".into(), "items_array".into()), + ] + ); +} + +#[test] +fn nested_qualified_compound_field_access_keeps_top_level_column() { + let result = analyze_one("SELECT base.payload.items[1] AS item FROM actual_table AS base"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "payload".into())] + ); +} + +#[test] +fn cte_compound_field_access_uses_cte_binding_column() { + let result = analyze_one( + "WITH base AS (SELECT items_array FROM actual_table) SELECT base.items_array[1] AS item FROM base", + ); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} + +#[test] +fn unqualified_compound_field_access_uses_top_level_column() { + let result = analyze_one("SELECT payload.items[1] AS item FROM t"); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!(concrete_sources(m), vec![("t".into(), "payload".into())]); +} + +#[test] +fn bigquery_offset_compound_field_access_uses_binding_column() { + let result = analyze_with_dialect( + "SELECT base.items_array[OFFSET(0)] AS item FROM actual_table AS base", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(m), + vec![("actual_table".into(), "items_array".into())] + ); +} From 17cee6e8115ee6cd6740a38fcc632a1ff88d9bbf Mon Sep 17 00:00:00 2001 From: eitsupi Date: Sun, 23 Aug 2026 20:37:19 +0000 Subject: [PATCH 12/14] fix: resolve UNNEST value-table lineage --- sqllineage/src/build/expr.rs | 57 +++- sqllineage/src/build/select.rs | 64 ++++- sqllineage/src/graph/mod.rs | 22 +- sqllineage/src/graph/node.rs | 15 +- sqllineage/src/graph/scope.rs | 44 +++ sqllineage/src/resolve/mod.rs | 436 ++++++++++++++++++++++++++--- sqllineage/tests/column_lineage.rs | 110 ++++++++ 7 files changed, 683 insertions(+), 65 deletions(-) diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index ba6ecab..73cefcf 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -4,23 +4,40 @@ use crate::build::LineageBuilder; use crate::build::select::split_compound; use crate::graph::edge::EdgeKind; use crate::graph::node::NodeId; -use crate::graph::scope::ScopeKind; +use crate::graph::scope::{Binding, ScopeKind}; impl LineageBuilder { pub(crate) fn collect_ancestors(&mut self, expr: &Expr) -> Vec { match expr { Expr::Identifier(ident) => { - let node = self + let binding = self .graph - .add_unqualified(ident.value.clone(), self.current_scope); + .scopes + .lookup(self.current_scope, &ident.value) + .cloned(); + let binding = + binding.filter(|binding| matches!(binding, Binding::VirtualSource(_))); + let node = self.graph.add_unqualified_with_binding( + ident.value.clone(), + self.current_scope, + binding, + ); vec![node] } Expr::CompoundIdentifier(parts) => { let (qualifier, column) = split_compound(parts); - let node = self + let binding = self .graph - .add_ref(column, Some(qualifier), self.current_scope); + .scopes + .lookup(self.current_scope, &qualifier) + .cloned(); + let node = self.graph.add_ref_with_binding( + column, + Some(qualifier), + self.current_scope, + binding, + ); vec![node] } @@ -316,25 +333,39 @@ impl LineageBuilder { access_chain: &[AccessExpr], ) -> Vec { let mut ancestors = match (root, access_chain.first()) { - (Expr::Identifier(binding), Some(AccessExpr::Dot(Expr::Identifier(field)))) + (Expr::Identifier(binding_name), Some(AccessExpr::Dot(Expr::Identifier(field)))) if self .graph .scopes - .lookup(self.current_scope, &binding.value) + .lookup(self.current_scope, &binding_name.value) .is_some() => { - let node = self.graph.add_ref( + let binding = self + .graph + .scopes + .lookup(self.current_scope, &binding_name.value) + .cloned(); + let node = self.graph.add_ref_with_binding( field.value.clone(), - Some(binding.value.clone()), + Some(binding_name.value.clone()), self.current_scope, + binding, ); vec![node] } (Expr::Identifier(column), Some(AccessExpr::Dot(Expr::Identifier(_)))) => { - vec![ - self.graph - .add_unqualified(column.value.clone(), self.current_scope), - ] + let binding = self + .graph + .scopes + .lookup(self.current_scope, &column.value) + .cloned(); + let binding = + binding.filter(|binding| matches!(binding, Binding::VirtualSource(_))); + vec![self.graph.add_unqualified_with_binding( + column.value.clone(), + self.current_scope, + binding, + )] } _ => self.collect_ancestors(root), }; diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index 4ed3489..b96fffa 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -4,7 +4,7 @@ use sqlparser::ast::{ use crate::build::LineageBuilder; use crate::build::expr::determine_edge_kind; -use crate::graph::scope::{Binding, ScopeColumn, ScopeKind}; +use crate::graph::scope::{Binding, ScopeColumn, ScopeKind, VirtualColumn, VirtualColumnState}; impl LineageBuilder { /// Process a SELECT — FROM first, then projection. @@ -167,9 +167,69 @@ impl LineageBuilder { let _ = alias; } + TableFactor::UNNEST { + alias, + array_exprs, + with_offset, + with_offset_alias, + with_ordinality, + } => { + // The array expressions are evaluated in the scope visible + // before this FROM item is introduced. Capture their nodes + // first, then install the range-variable binding so it is + // visible to subsequent lateral FROM items and projection. + let dependencies = array_exprs + .iter() + .map(|expr| self.collect_ancestors(expr)) + .collect::>(); + + let Some(alias) = alias else { + return; + }; + + let mut columns = Vec::with_capacity( + array_exprs.len() + usize::from(*with_offset) + usize::from(*with_ordinality), + ); + for (index, deps) in dependencies.into_iter().enumerate() { + let name = alias.columns.get(index).map_or_else( + || alias.name.value.clone(), + |column| column.name.value.clone(), + ); + columns.push(VirtualColumn { + name, + state: if deps.is_empty() { + VirtualColumnState::KnownEmpty + } else { + VirtualColumnState::Unknown + }, + dependencies: deps, + }); + } + if *with_offset { + columns.push(VirtualColumn { + name: with_offset_alias + .as_ref() + .map_or_else(|| "offset".to_string(), |ident| ident.value.clone()), + dependencies: Vec::new(), + state: VirtualColumnState::KnownEmpty, + }); + } else if *with_ordinality { + columns.push(VirtualColumn { + name: "ordinality".to_string(), + dependencies: Vec::new(), + state: VirtualColumnState::KnownEmpty, + }); + } + + let virtual_id = self + .graph + .scopes + .add_virtual_source(self.current_scope, columns); + self.add_binding(alias.name.value.clone(), Binding::VirtualSource(virtual_id)); + } + TableFactor::TableFunction { .. } | TableFactor::Function { .. } - | TableFactor::UNNEST { .. } | TableFactor::JsonTable { .. } | TableFactor::OpenJsonTable { .. } | TableFactor::Pivot { .. } diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 1e2134d..3cc9d3b 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -38,16 +38,32 @@ impl RawGraph { }) } - pub fn add_ref(&mut self, name: String, qualifier: Option, scope: ScopeId) -> NodeId { + pub fn add_ref_with_binding( + &mut self, + name: String, + qualifier: Option, + scope: ScopeId, + binding: Option, + ) -> NodeId { self.add_node(RawNode::Ref { name, qualifier, scope, + binding, }) } - pub fn add_unqualified(&mut self, name: String, scope: ScopeId) -> NodeId { - self.add_node(RawNode::Unqualified { name, scope }) + pub fn add_unqualified_with_binding( + &mut self, + name: String, + scope: ScopeId, + binding: Option, + ) -> NodeId { + self.add_node(RawNode::Unqualified { + name, + scope, + binding, + }) } pub fn add_star(&mut self, table: Option, scope: ScopeId) -> NodeId { diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index ef7eafa..f6b1899 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -1,5 +1,5 @@ use crate::graph::edge::EdgeKind; -use crate::graph::scope::ScopeId; +use crate::graph::scope::{Binding, ScopeId}; use crate::types::TableRef; pub(crate) type NodeId = usize; @@ -20,6 +20,10 @@ pub(crate) enum RawNode { name: String, qualifier: Option, scope: ScopeId, + /// Binding captured while building a FROM expression. This prevents + /// a later table alias from changing the meaning of a lateral + /// dependency (for example `base, UNNEST(base.items) AS base`). + binding: Option, }, /// SELECT * or table.* — expandable with catalog. Star { @@ -27,5 +31,12 @@ pub(crate) enum RawNode { scope: ScopeId, }, /// Unqualified column in multi-table scope. - Unqualified { name: String, scope: ScopeId }, + Unqualified { + name: String, + scope: ScopeId, + /// Binding visible while the expression was built. This keeps a + /// lateral FROM dependency attached to the preceding relation even + /// when a later range variable shadows its name. + binding: Option, + }, } diff --git a/sqllineage/src/graph/scope.rs b/sqllineage/src/graph/scope.rs index 697fbce..2073e60 100644 --- a/sqllineage/src/graph/scope.rs +++ b/sqllineage/src/graph/scope.rs @@ -5,6 +5,30 @@ use crate::types::TableRef; pub(crate) type ScopeId = usize; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct VirtualSourceId { + scope: ScopeId, + index: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VirtualColumnState { + KnownEmpty, + Unknown, +} + +#[derive(Debug, Clone)] +pub(crate) struct VirtualColumn { + pub name: String, + pub dependencies: Vec, + pub state: VirtualColumnState, +} + +#[derive(Debug, Clone)] +pub(crate) struct VirtualSource { + pub columns: Vec, +} + pub(crate) struct ScopeTree { scopes: Vec, } @@ -15,6 +39,7 @@ struct Scope { anonymous_derived: Vec, output_columns: Vec, output_plan: OutputPlan, + virtual_sources: Vec, } #[derive(Debug, Clone)] @@ -44,6 +69,7 @@ pub(crate) enum Binding { Table(TableRef), Cte(ScopeId), DerivedTable(ScopeId), + VirtualSource(VirtualSourceId), } #[derive(Debug, Clone)] @@ -61,6 +87,7 @@ impl ScopeTree { anonymous_derived: Vec::new(), output_columns: Vec::new(), output_plan: OutputPlan::Projection, + virtual_sources: Vec::new(), }], } } @@ -77,6 +104,7 @@ impl ScopeTree { anonymous_derived: Vec::new(), output_columns: Vec::new(), output_plan: OutputPlan::Projection, + virtual_sources: Vec::new(), }); id } @@ -89,6 +117,22 @@ impl ScopeTree { self.scopes[scope].bindings.insert(name, binding); } + pub fn add_virtual_source( + &mut self, + scope: ScopeId, + columns: Vec, + ) -> VirtualSourceId { + let index = self.scopes[scope].virtual_sources.len(); + self.scopes[scope] + .virtual_sources + .push(VirtualSource { columns }); + VirtualSourceId { scope, index } + } + + pub fn virtual_source(&self, id: VirtualSourceId) -> &VirtualSource { + &self.scopes[id.scope].virtual_sources[id.index] + } + pub fn add_output_column(&mut self, scope: ScopeId, col: ScopeColumn) { self.scopes[scope].output_columns.push(col); } diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index 7ca2185..c9c2b09 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use crate::graph::RawGraph; use crate::graph::edge::EdgeKind; use crate::graph::node::{NodeId, RawNode}; -use crate::graph::scope::{Binding, OutputPlan, ScopeTree}; +use crate::graph::scope::{Binding, OutputPlan, ScopeTree, VirtualColumnState, VirtualSourceId}; use crate::types::{ AnalyzeResult, CatalogProvider, ColumnLineage, ColumnMapping, ColumnOrigin, ColumnRef, ParseError, StatementType, TableRef, TransformKind, Warning, WarningKind, @@ -183,6 +183,9 @@ fn star_arity( Binding::Cte(child) | Binding::DerivedTable(child) => { scope_arity(child, graph, catalog, active)? } + Binding::VirtualSource(source) => { + Some(graph.scopes.virtual_source(source).columns.len()) + } }; match binding_width { Some(binding_width) => width += binding_width, @@ -220,6 +223,80 @@ fn wildcard_mapping(output_table: Option<&TableRef>, source_table: TableRef) -> } } +#[allow(clippy::too_many_arguments)] +fn expand_virtual_source( + source: VirtualSourceId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, + mappings: &mut Vec, +) { + for column in &graph.scopes.virtual_source(source).columns { + let mut origins = Vec::new(); + let mut visited = HashSet::new(); + for &dependency in &column.dependencies { + let (dependency_origins, _, _) = collect_leaf_origins( + dependency, + graph, + resolved, + incoming, + &mut visited, + catalog, + mapping_cache, + ); + origins.extend(dependency_origins); + } + if origins.is_empty() + && matches!(column.state, VirtualColumnState::Unknown) + && !column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) + { + origins.push(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }); + } + mappings.push(ColumnMapping { + target: ColumnRef { + table: None, + column: column.name.clone(), + }, + sources: origins, + transform: TransformKind::Direct, + }); + } +} + +fn known_empty_dependency(node_id: NodeId, graph: &RawGraph) -> bool { + let (name, binding) = match &graph.nodes[node_id] { + RawNode::Ref { name, binding, .. } | RawNode::Unqualified { name, binding, .. } => { + (name, binding.as_ref()) + } + _ => return false, + }; + let Some(Binding::VirtualSource(source)) = binding else { + return false; + }; + let Some(column) = graph + .scopes + .virtual_source(*source) + .columns + .iter() + .find(|column| column.name == *name) + else { + return false; + }; + matches!(column.state, VirtualColumnState::KnownEmpty) + && column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) +} + #[derive(Default)] struct ScopeMappingCache { entries: HashMap, @@ -539,6 +616,16 @@ fn expand_star( ); } else if let Some(Binding::Table(actual_table)) = binding { mappings.push(wildcard_mapping(None, actual_table)); + } else if let Some(Binding::VirtualSource(source)) = binding { + expand_virtual_source( + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + ); } else { mappings.push(wildcard_mapping(None, t.clone())); } @@ -558,6 +645,15 @@ fn expand_star( visited_scopes, ); } + Binding::VirtualSource(source) => expand_virtual_source( + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + mappings, + ), } } for &child in graph.scopes.anonymous_derived(scope) { @@ -719,6 +815,12 @@ fn collect_leaf_origins( catalog: Option<&dyn CatalogProvider>, mapping_cache: &mut ScopeMappingCache, ) -> (Vec, bool, TransformKind) { + if let Some(result) = + resolve_virtual_reference(node_id, graph, resolved, incoming, catalog, mapping_cache) + { + return result; + } + if let Some((sources, has_back, transform)) = resolve_named_scope_reference(node_id, graph, resolved, incoming, catalog, mapping_cache) { @@ -784,6 +886,149 @@ fn collect_leaf_origins( } } +#[allow(clippy::too_many_arguments)] +fn resolve_virtual_reference( + node_id: NodeId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option<(Vec, bool, TransformKind)> { + let (name, binding, scope) = match &graph.nodes[node_id] { + RawNode::Ref { + name, + qualifier, + scope, + binding, + } => ( + name, + binding + .clone() + .or_else(|| { + qualifier + .as_deref() + .and_then(|qualifier| graph.scopes.lookup(*scope, qualifier).cloned()) + }) + .or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }), + *scope, + ), + RawNode::Unqualified { + name, + scope, + binding, + } => ( + name, + binding.clone().or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }), + *scope, + ), + _ => return None, + }; + let source = match binding { + Some(Binding::VirtualSource(source)) => source, + Some(_) => return None, + None => match find_virtual_sources_for_column(scope, name, graph).as_slice() { + [source] => *source, + [] => return None, + _ => { + return Some(( + vec![ColumnOrigin::Ambiguous { + column: name.clone(), + candidates: Vec::new(), + }], + false, + TransformKind::Direct, + )); + } + }, + }; + resolve_virtual_column_sources( + name, + source, + graph, + resolved, + incoming, + catalog, + mapping_cache, + ) +} + +#[allow(clippy::too_many_arguments)] +fn resolve_virtual_column_sources( + name: &str, + source: VirtualSourceId, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option<(Vec, bool, TransformKind)> { + let column = graph + .scopes + .virtual_source(source) + .columns + .iter() + .find(|column| column.name == name)?; + let mut origins = Vec::new(); + for &dependency in &column.dependencies { + let (dependency_origins, _, _) = collect_leaf_origins( + dependency, + graph, + resolved, + incoming, + &mut HashSet::new(), + catalog, + mapping_cache, + ); + origins.extend(dependency_origins); + } + if origins.is_empty() + && matches!(column.state, VirtualColumnState::Unknown) + && !column + .dependencies + .iter() + .all(|&dependency| known_empty_dependency(dependency, graph)) + { + origins.push(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }); + } + Some((origins, false, TransformKind::Direct)) +} + +fn virtual_column_origin( + name: &str, + source: VirtualSourceId, + graph: &RawGraph, +) -> Option { + let column = graph + .scopes + .virtual_source(source) + .columns + .iter() + .find(|column| column.name == name)?; + match column.state { + VirtualColumnState::KnownEmpty => None, + VirtualColumnState::Unknown => Some(ColumnOrigin::Ambiguous { + column: column.name.clone(), + candidates: Vec::new(), + }), + } +} + /// Resolve a named reference through a set-operation/Delegate scope using the /// expanded mappings. Raw scope columns intentionally do not contain names /// for individual catalog-expanded star outputs, so lookup must happen here. @@ -800,13 +1045,24 @@ fn resolve_named_scope_reference( name, qualifier, scope, + .. } => (name, qualifier.as_ref(), *scope), - RawNode::Unqualified { name, scope } => (name, None, *scope), + RawNode::Unqualified { name, scope, .. } => (name, None, *scope), _ => return None, }; - let binding = qualifier - .and_then(|qualifier| graph.scopes.lookup(scope, qualifier).cloned()) - .or_else(|| find_single_binding(scope, graph)); + let binding = match &graph.nodes[node_id] { + RawNode::Ref { + binding: Some(binding), + .. + } + | RawNode::Unqualified { + binding: Some(binding), + .. + } => Some(binding.clone()), + _ => None, + } + .or_else(|| qualifier.and_then(|qualifier| graph.scopes.lookup(scope, qualifier).cloned())) + .or_else(|| find_single_binding(scope, graph)); let Some(Binding::Cte(target_scope) | Binding::DerivedTable(target_scope)) = binding else { return None; }; @@ -876,31 +1132,29 @@ fn resolve_node( name, qualifier, scope, + binding, } => { - if let Some(qual) = qualifier { - let binding = graph.scopes.lookup(*scope, qual).cloned(); - match binding { - Some(Binding::Table(table_ref)) => Some(ColumnOrigin::Concrete { - table: table_ref, - column: name.clone(), - }), - Some(Binding::Cte(cte_scope) | Binding::DerivedTable(cte_scope)) => { - resolve_through_scope( - name, - cte_scope, - graph, - resolved, - incoming, - visited, - catalog, - mapping_cache, - ) - } - None => Some(ColumnOrigin::Concrete { - table: TableRef::new(qual.as_str()), - column: name.clone(), - }), - } + let binding = binding.clone().or_else(|| { + qualifier + .as_deref() + .and_then(|qual| graph.scopes.lookup(*scope, qual).cloned()) + }); + if let Some(binding) = binding { + resolve_captured_binding( + name, + binding, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else if let Some(qual) = qualifier { + Some(ColumnOrigin::Concrete { + table: TableRef::new(qual.as_str()), + column: name.clone(), + }) } else { resolve_unqualified( name, @@ -915,16 +1169,41 @@ fn resolve_node( } } - RawNode::Unqualified { name, scope } => resolve_unqualified( + RawNode::Unqualified { name, - *scope, - graph, - resolved, - incoming, - visited, - catalog, - mapping_cache, - ), + scope, + binding, + } => { + if let Some(binding) = binding.clone().or_else(|| { + graph + .scopes + .lookup(*scope, name) + .filter(|binding| matches!(binding, Binding::VirtualSource(_))) + .cloned() + }) { + resolve_captured_binding( + name, + binding, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else { + resolve_unqualified( + name, + *scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } + } RawNode::Star { table, .. } => table .as_ref() @@ -937,18 +1216,52 @@ fn resolve_node( origin } +#[allow(clippy::too_many_arguments)] +fn resolve_captured_binding( + name: &str, + binding: Binding, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option { + match binding { + Binding::Table(table) => Some(ColumnOrigin::Concrete { + table, + column: name.to_string(), + }), + Binding::Cte(scope) | Binding::DerivedTable(scope) => resolve_through_scope( + name, + scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + Binding::VirtualSource(source) => virtual_column_origin(name, source, graph), + } +} + fn find_cte_redirect(node_id: NodeId, graph: &RawGraph) -> Option<(NodeId, usize)> { match &graph.nodes[node_id] { RawNode::Ref { name, qualifier, scope, + binding, } => { - let binding = if let Some(qual) = qualifier { - graph.scopes.lookup(*scope, qual).cloned() - } else { - find_single_binding(*scope, graph) - }; + let binding = binding + .clone() + .or_else(|| { + qualifier + .as_deref() + .and_then(|qual| graph.scopes.lookup(*scope, qual).cloned()) + }) + .or_else(|| find_single_binding(*scope, graph)); match binding { Some(Binding::Cte(s) | Binding::DerivedTable(s)) => graph .scopes @@ -959,8 +1272,14 @@ fn find_cte_redirect(node_id: NodeId, graph: &RawGraph) -> Option<(NodeId, usize _ => None, } } - RawNode::Unqualified { name, scope } => { - let binding = find_single_binding(*scope, graph); + RawNode::Unqualified { + name, + scope, + binding, + } => { + let binding = binding + .clone() + .or_else(|| find_single_binding(*scope, graph)); match binding { Some(Binding::Cte(s) | Binding::DerivedTable(s)) => graph .scopes @@ -1035,6 +1354,7 @@ fn resolve_from_bindings( catalog, mapping_cache, ), + Binding::VirtualSource(_) => None, } } else if bindings.is_empty() { Some(ColumnOrigin::Ambiguous { @@ -1065,6 +1385,7 @@ fn resolve_from_bindings( } } Binding::Table(t) => table_candidates.push(t.clone()), + Binding::VirtualSource(_) => {} } } if table_candidates.len() == 1 { @@ -1081,6 +1402,31 @@ fn resolve_from_bindings( } } +fn virtual_has_column(name: &str, source: VirtualSourceId, graph: &RawGraph) -> bool { + graph + .scopes + .virtual_source(source) + .columns + .iter() + .any(|column| column.name == name) +} + +fn find_virtual_sources_for_column( + scope: usize, + name: &str, + graph: &RawGraph, +) -> Vec { + effective_bindings(scope, graph) + .into_iter() + .filter_map(|(_, binding)| match binding { + Binding::VirtualSource(source) if virtual_has_column(name, source, graph) => { + Some(source) + } + _ => None, + }) + .collect() +} + #[allow(clippy::too_many_arguments)] fn resolve_through_scope( column_name: &str, diff --git a/sqllineage/tests/column_lineage.rs b/sqllineage/tests/column_lineage.rs index eadc99a..3bc41b2 100644 --- a/sqllineage/tests/column_lineage.rs +++ b/sqllineage/tests/column_lineage.rs @@ -164,3 +164,113 @@ fn select_cast_passthrough() { assert_eq!(concrete_sources(m), vec![("t".into(), "a".into())]); assert_eq!(m.transform, TransformKind::Direct); } + +#[test] +fn unnest_source_free_alias_has_no_physical_sources() { + let result = analyze_one( + "SELECT item FROM UNNEST(GENERATE_DATE_ARRAY(DATE('2020-01-01'), DATE('2020-01-03'))) AS item", + ); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert!(mapping.sources.is_empty()); + assert_eq!(result.tables.inputs, Vec::::new()); +} + +#[test] +fn unnest_alias_depends_on_array_column() { + let result = analyze_one("SELECT item FROM base, UNNEST(base.items_array) AS item"); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert_eq!( + concrete_sources(mapping), + vec![("base".into(), "items_array".into())] + ); +} + +#[test] +fn unnest_unresolved_array_is_ambiguous_not_alias_column() { + let result = analyze_one("SELECT item FROM UNNEST(missing_array) AS item"); + let mapping = find_mapping(&result.columns.mappings, "item"); + assert!(matches!( + mapping.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "item" && candidates.is_empty() + )); +} + +#[test] +fn unnest_unqualified_array_column_keeps_prior_table_binding() { + let result = analyze_one("SELECT item FROM base, UNNEST(items_array) AS item"); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "item")), + vec![("base".into(), "items_array".into())] + ); +} + +#[test] +fn unnest_known_empty_virtual_dependency_stays_source_free() { + let result = analyze_one("SELECT item FROM UNNEST([1, 2]) AS source, UNNEST(source) AS item"); + assert!( + find_mapping(&result.columns.mappings, "item") + .sources + .is_empty() + ); +} + +#[test] +fn unqualified_identifier_does_not_capture_relation_alias() { + let result = analyze_one("SELECT a FROM table1 AS a, table2 AS b"); + let mapping = find_mapping(&result.columns.mappings, "a"); + assert!(matches!( + mapping.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "a" && candidates.len() == 2 + )); +} + +#[test] +fn duplicate_virtual_slots_are_ambiguous_but_qualified_slots_resolve() { + let result = + analyze_one("SELECT x FROM base, UNNEST(base.first) AS u(x), UNNEST(base.second) AS v(x)"); + assert!(matches!( + find_mapping(&result.columns.mappings, "x").sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "x" && candidates.is_empty() + )); + + let result = analyze_one( + "SELECT u.x, v.x FROM base, UNNEST(base.first) AS u(x), UNNEST(base.second) AS v(x)", + ); + assert_eq!( + concrete_sources(&result.columns.mappings[0]), + vec![("base".into(), "first".into())] + ); + assert_eq!( + concrete_sources(&result.columns.mappings[1]), + vec![("base".into(), "second".into())] + ); +} + +#[test] +fn unnest_alias_columns_keep_array_expression_ordinals() { + let result = analyze_one("SELECT x, y FROM base, UNNEST(base.first, base.second) AS u(x, y)"); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "x")), + vec![("base".into(), "first".into())] + ); + assert_eq!( + concrete_sources(find_mapping(&result.columns.mappings, "y")), + vec![("base".into(), "second".into())] + ); +} + +#[test] +fn unnest_offset_is_a_source_free_generated_slot() { + let result = analyze_one("SELECT item, off FROM UNNEST([1, 2]) AS item WITH OFFSET AS off"); + assert!( + find_mapping(&result.columns.mappings, "item") + .sources + .is_empty() + ); + assert!( + find_mapping(&result.columns.mappings, "off") + .sources + .is_empty() + ); +} From 360e1e7c2c30220e4e72a869e82225dcd52403ce Mon Sep 17 00:00:00 2001 From: eitsupi Date: Mon, 24 Aug 2026 00:40:45 +0000 Subject: [PATCH 13/14] fix: resolve structured field references --- sqllineage/src/build/expr.rs | 127 ++++++++++++--- sqllineage/src/build/mod.rs | 6 +- sqllineage/src/build/select.rs | 14 +- sqllineage/src/graph/mod.rs | 13 ++ sqllineage/src/graph/node.rs | 9 ++ sqllineage/src/lib.rs | 2 +- sqllineage/src/resolve/mod.rs | 83 ++++++++++ sqllineage/src/types.rs | 13 ++ sqllineage/tests/expr_coverage.rs | 249 +++++++++++++++++++++++++++++- 9 files changed, 478 insertions(+), 38 deletions(-) diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index 73cefcf..d6eb32d 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -1,10 +1,10 @@ use sqlparser::ast::{self, AccessExpr, Expr, FunctionArguments, Subscript, WindowType}; use crate::build::LineageBuilder; -use crate::build::select::split_compound; use crate::graph::edge::EdgeKind; use crate::graph::node::NodeId; use crate::graph::scope::{Binding, ScopeKind}; +use crate::types::TableRef; impl LineageBuilder { pub(crate) fn collect_ancestors(&mut self, expr: &Expr) -> Vec { @@ -15,6 +15,19 @@ impl LineageBuilder { .scopes .lookup(self.current_scope, &ident.value) .cloned(); + if binding.as_ref().is_some_and(|binding| { + self.dialect.supports_relation_alias_row_value() + && matches!( + binding, + Binding::Table(_) | Binding::Cte(_) | Binding::DerivedTable(_) + ) + }) { + return vec![self.graph.add_row_value_candidate( + ident.value.clone(), + self.current_scope, + binding, + )]; + } let binding = binding.filter(|binding| matches!(binding, Binding::VirtualSource(_))); let node = self.graph.add_unqualified_with_binding( @@ -25,21 +38,7 @@ impl LineageBuilder { vec![node] } - Expr::CompoundIdentifier(parts) => { - let (qualifier, column) = split_compound(parts); - let binding = self - .graph - .scopes - .lookup(self.current_scope, &qualifier) - .cloned(); - let node = self.graph.add_ref_with_binding( - column, - Some(qualifier), - self.current_scope, - binding, - ); - vec![node] - } + Expr::CompoundIdentifier(parts) => self.collect_compound_identifier_ancestors(parts), Expr::Value(_) | Expr::TypedString { .. } @@ -345,13 +344,11 @@ impl LineageBuilder { .scopes .lookup(self.current_scope, &binding_name.value) .cloned(); - let node = self.graph.add_ref_with_binding( + vec![self.add_bound_field_ancestor( + binding_name.value.clone(), field.value.clone(), - Some(binding_name.value.clone()), - self.current_scope, binding, - ); - vec![node] + )] } (Expr::Identifier(column), Some(AccessExpr::Dot(Expr::Identifier(_)))) => { let binding = self @@ -378,6 +375,82 @@ impl LineageBuilder { ancestors } + /// Resolve a plain dotted identifier by separating its relation binding + /// from the top-level physical column. The parser represents both + /// `alias.column` and `alias.struct.field` as a flat compound identifier, + /// so rendering all but the final component as one qualifier loses the + /// distinction between a relation name and a nested field path. + fn collect_compound_identifier_ancestors( + &mut self, + parts: &[sqlparser::ast::Ident], + ) -> Vec { + if let Some((prefix_len, binding)) = self.find_compound_binding(parts) { + let qualifier = parts[..prefix_len] + .iter() + .map(|part| part.value.as_str()) + .collect::>() + .join("."); + let column = parts[prefix_len].value.clone(); + return vec![self.add_bound_field_ancestor(qualifier, column, Some(binding))]; + } + + // No relation prefix was found: `struct.field` is an unqualified + // top-level column followed by a nested field path. Only the + // top-level column can be represented by the public ColumnOrigin API. + let column = parts[0].value.clone(); + let binding = self + .graph + .scopes + .lookup(self.current_scope, &column) + .cloned() + .filter(|binding| matches!(binding, Binding::VirtualSource(_))); + vec![ + self.graph + .add_unqualified_with_binding(column, self.current_scope, binding), + ] + } + + /// Find the longest visible relation prefix in a compound identifier. + /// + /// A one-component prefix is a SQL alias. Longer prefixes are matched + /// against the physical parts of a table binding, allowing references such + /// as `catalog.schema.table.column` without turning the relation into a + /// single quoted string containing dots. + fn find_compound_binding(&self, parts: &[sqlparser::ast::Ident]) -> Option<(usize, Binding)> { + let visible = self.graph.scopes.visible_bindings(self.current_scope); + (1..parts.len()).rev().find_map(|prefix_len| { + let prefix = parts[..prefix_len] + .iter() + .map(|part| part.value.as_str()) + .collect::>(); + + // Aliases are single identifiers and therefore only match the + // first component of a compound identifier. + if prefix_len == 1 + && let Some((_, binding)) = visible.iter().find(|(name, _)| name == prefix[0]) + { + return Some((prefix_len, binding.clone())); + } + + visible.iter().find_map(|(_, binding)| { + let Binding::Table(table) = binding else { + return None; + }; + (table_parts(table) == prefix).then(|| (prefix_len, binding.clone())) + }) + }) + } + + fn add_bound_field_ancestor( + &mut self, + qualifier: String, + column: String, + binding: Option, + ) -> NodeId { + self.graph + .add_ref_with_binding(column, Some(qualifier), self.current_scope, binding) + } + fn collect_subscript_ancestors(&mut self, subscript: &Subscript) -> Vec { match subscript { Subscript::Index { index } => self.collect_ancestors(index), @@ -402,6 +475,18 @@ impl LineageBuilder { } } +fn table_parts(table: &TableRef) -> Vec<&str> { + let mut parts = Vec::with_capacity(3); + if let Some(catalog) = &table.catalog { + parts.push(catalog.as_str()); + } + if let Some(schema) = &table.schema { + parts.push(schema.as_str()); + } + parts.push(table.table.as_str()); + parts +} + pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind { match expr { Expr::Identifier(_) | Expr::CompoundIdentifier(_) | Expr::Value(_) => EdgeKind::Direct, diff --git a/sqllineage/src/build/mod.rs b/sqllineage/src/build/mod.rs index 605b29c..b0fd16b 100644 --- a/sqllineage/src/build/mod.rs +++ b/sqllineage/src/build/mod.rs @@ -5,7 +5,7 @@ pub(crate) mod statement; use crate::graph::RawGraph; use crate::graph::scope::{Binding, ScopeId, ScopeKind, ScopeTree}; -use crate::types::{StatementType, Warning}; +use crate::types::{Dialect, StatementType, Warning}; use sqlparser::ast::Statement; pub(crate) struct LineageBuilder { @@ -15,10 +15,11 @@ pub(crate) struct LineageBuilder { pub(crate) warnings: Vec, pub(crate) normalize_case: bool, pub(crate) inner_statement_type: Option, + pub(crate) dialect: Dialect, } impl LineageBuilder { - pub fn new(normalize_case: bool) -> Self { + pub fn new(normalize_case: bool, dialect: Dialect) -> Self { let graph = RawGraph::new(); let root = ScopeTree::root(); Self { @@ -28,6 +29,7 @@ impl LineageBuilder { warnings: Vec::new(), normalize_case, inner_statement_type: None, + dialect, } } diff --git a/sqllineage/src/build/select.rs b/sqllineage/src/build/select.rs index b96fffa..69e0f0e 100644 --- a/sqllineage/src/build/select.rs +++ b/sqllineage/src/build/select.rs @@ -1,5 +1,5 @@ use sqlparser::ast::{ - Expr, Ident, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, TableWithJoins, + Expr, Select, SelectItem, SelectItemQualifiedWildcardKind, TableFactor, TableWithJoins, }; use crate::build::LineageBuilder; @@ -253,15 +253,3 @@ fn infer_column_name(expr: &Expr) -> String { _ => "?column?".to_string(), } } - -/// Split a compound identifier into (qualifier, `column_name`). -pub(crate) fn split_compound(parts: &[Ident]) -> (String, String) { - let len = parts.len(); - let column = parts[len - 1].value.clone(); - let qualifier = parts[..len - 1] - .iter() - .map(|p| p.value.as_str()) - .collect::>() - .join("."); - (qualifier, column) -} diff --git a/sqllineage/src/graph/mod.rs b/sqllineage/src/graph/mod.rs index 3cc9d3b..bedf2ee 100644 --- a/sqllineage/src/graph/mod.rs +++ b/sqllineage/src/graph/mod.rs @@ -66,6 +66,19 @@ impl RawGraph { }) } + pub fn add_row_value_candidate( + &mut self, + name: String, + scope: ScopeId, + binding: Option, + ) -> NodeId { + self.add_node(RawNode::RowValueCandidate { + name, + scope, + binding, + }) + } + pub fn add_star(&mut self, table: Option, scope: ScopeId) -> NodeId { self.add_node(RawNode::Star { table, scope }) } diff --git a/sqllineage/src/graph/node.rs b/sqllineage/src/graph/node.rs index f6b1899..6061105 100644 --- a/sqllineage/src/graph/node.rs +++ b/sqllineage/src/graph/node.rs @@ -39,4 +39,13 @@ pub(crate) enum RawNode { /// when a later range variable shadows its name. binding: Option, }, + /// A relation alias used where the dialect permits a whole-row value. + /// Resolution must distinguish this from a source-free expression: a + /// catalog or derived scope can still prove that the alias is an ordinary + /// physical/output column with the same name. + RowValueCandidate { + name: String, + scope: ScopeId, + binding: Option, + }, } diff --git a/sqllineage/src/lib.rs b/sqllineage/src/lib.rs index 09ff161..e98860b 100644 --- a/sqllineage/src/lib.rs +++ b/sqllineage/src/lib.rs @@ -73,7 +73,7 @@ pub fn analyze(sql: &str, opts: AnalyzeOptions) -> Result, Pa statements .iter() .map(|stmt| { - let builder = build::LineageBuilder::new(opts.normalize_case); + let builder = build::LineageBuilder::new(opts.normalize_case, opts.dialect); let (raw_graph, warnings, statement_type) = builder.build(stmt); resolve::resolve(raw_graph, catalog.as_deref(), warnings, statement_type) }) diff --git a/sqllineage/src/resolve/mod.rs b/sqllineage/src/resolve/mod.rs index c9c2b09..c617df8 100644 --- a/sqllineage/src/resolve/mod.rs +++ b/sqllineage/src/resolve/mod.rs @@ -1205,6 +1205,22 @@ fn resolve_node( } } + RawNode::RowValueCandidate { + name, + scope, + binding, + } => resolve_row_value_candidate( + name, + *scope, + binding.clone(), + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ), + RawNode::Star { table, .. } => table .as_ref() .map(|t| ColumnOrigin::Wildcard { table: t.clone() }), @@ -1216,6 +1232,73 @@ fn resolve_node( origin } +#[allow(clippy::too_many_arguments)] +fn resolve_row_value_candidate( + name: &str, + scope: usize, + binding: Option, + graph: &RawGraph, + resolved: &mut Vec>, + incoming: &[Vec], + visited: &mut HashSet, + catalog: Option<&dyn CatalogProvider>, + mapping_cache: &mut ScopeMappingCache, +) -> Option { + let binding = binding.or_else(|| graph.scopes.lookup(scope, name).cloned()); + let Some(binding) = binding else { + return Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }); + }; + + match binding { + Binding::Table(table) => { + if let Some(owner) = catalog + .and_then(|catalog| catalog.resolve_column(name, std::slice::from_ref(&table))) + { + Some(ColumnOrigin::Concrete { + table: owner, + column: name.to_string(), + }) + } else { + Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }) + } + } + Binding::Cte(target_scope) | Binding::DerivedTable(target_scope) => { + let is_named_column = graph + .scopes + .output_columns(target_scope) + .iter() + .any(|column| column.name == name); + if is_named_column { + resolve_through_scope( + name, + target_scope, + graph, + resolved, + incoming, + visited, + catalog, + mapping_cache, + ) + } else { + Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }) + } + } + Binding::VirtualSource(_) => Some(ColumnOrigin::Ambiguous { + column: name.to_string(), + candidates: Vec::new(), + }), + } +} + #[allow(clippy::too_many_arguments)] fn resolve_captured_binding( name: &str, diff --git a/sqllineage/src/types.rs b/sqllineage/src/types.rs index dd1e992..2504ee4 100644 --- a/sqllineage/src/types.rs +++ b/sqllineage/src/types.rs @@ -210,6 +210,19 @@ pub enum Dialect { BigQuery, } +impl Dialect { + /// Whether an unqualified relation alias can denote the complete row. + /// + /// `BigQuery` and PostgreSQL permit expressions such as `ARRAY_AGG(t)` when + /// `t` is a range-variable alias. Such an expression is a row/record + /// value, not a physical column named after the alias. The lineage API + /// has no whole-row origin, so callers must retain honest uncertainty + /// instead of fabricating a concrete `table.alias` source. + pub(crate) const fn supports_relation_alias_row_value(self) -> bool { + matches!(self, Self::BigQuery | Self::PostgreSql) + } +} + /// Error returned when SQL parsing or semantic validation fails. #[derive(Debug, Clone)] pub struct ParseError { diff --git a/sqllineage/tests/expr_coverage.rs b/sqllineage/tests/expr_coverage.rs index 610178f..6f6276c 100644 --- a/sqllineage/tests/expr_coverage.rs +++ b/sqllineage/tests/expr_coverage.rs @@ -1,7 +1,9 @@ mod common; use common::{analyze_one, concrete_sources, find_mapping}; -use sqllineage::{AnalyzeOptions, Dialect, TransformKind, analyze}; +use sqllineage::{ + AnalyzeOptions, CatalogProvider, ColumnOrigin, Dialect, TableRef, TransformKind, analyze, +}; fn analyze_with_dialect(sql: &str, dialect: Dialect) -> sqllineage::AnalyzeResult { analyze( @@ -186,3 +188,248 @@ fn bigquery_offset_compound_field_access_uses_binding_column() { vec![("actual_table".into(), "items_array".into())] ); } + +#[test] +fn qualified_struct_field_access_uses_binding_column() { + let result = analyze_with_dialect( + "SELECT agg.event.qualified_field AS field FROM upstream_model AS agg", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "field"); + assert_eq!( + concrete_sources(m), + vec![("upstream_model".into(), "event".into())] + ); +} + +#[test] +fn unqualified_struct_field_access_uses_top_level_column() { + let result = analyze_with_dialect( + "SELECT event.bare_field AS field FROM upstream_model", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "field"); + assert_eq!( + concrete_sources(m), + vec![("upstream_model".into(), "event".into())] + ); +} + +#[test] +fn cte_struct_field_access_uses_cte_binding_column() { + let result = analyze_one( + "WITH upstream AS (SELECT event FROM source) SELECT upstream.event.field AS value FROM upstream", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); +} + +#[test] +fn derived_struct_field_access_uses_derived_binding_column() { + let result = analyze_one( + "SELECT derived.event.field AS value FROM (SELECT event FROM source) AS derived", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); +} + +#[test] +fn physical_relation_prefix_struct_field_access_uses_table_parts() { + let result = + analyze_one("SELECT catalog.schema.source.event.field AS value FROM catalog.schema.source"); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!(concrete_sources(m), vec![("source".into(), "event".into())]); + assert_eq!(result.tables.inputs[0].catalog.as_deref(), Some("catalog")); + assert_eq!(result.tables.inputs[0].schema.as_deref(), Some("schema")); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog.as_deref(), Some("catalog")); + assert_eq!(table.schema.as_deref(), Some("schema")); + assert_eq!(table.table, "source"); + assert_eq!(column, "event"); + } + other => panic!("expected concrete physical relation source, got {other:?}"), + } +} + +#[test] +fn quoted_single_component_relation_name_keeps_embedded_dot() { + let result = analyze_with_dialect( + "SELECT \"orders.v2\".payload.field AS value FROM \"orders.v2\"", + Dialect::PostgreSql, + ); + let m = find_mapping(&result.columns.mappings, "value"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog, None); + assert_eq!(table.schema, None); + assert_eq!(table.table, "orders.v2"); + assert_eq!(column, "payload"); + } + other => panic!("expected quoted relation source, got {other:?}"), + } +} + +#[test] +fn qualified_struct_field_access_keeps_normal_alias_column_resolution() { + let result = analyze_one("SELECT source.user_id AS value FROM source"); + let m = find_mapping(&result.columns.mappings, "value"); + assert_eq!( + concrete_sources(m), + vec![("source".into(), "user_id".into())] + ); +} + +#[test] +fn unqualified_struct_field_access_keeps_binding_ambiguity() { + let result = analyze_one( + "SELECT event.field AS value FROM first_source JOIN second_source ON first_source.id = second_source.id", + ); + let m = find_mapping(&result.columns.mappings, "value"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] if column == "event" && candidates.len() == 2 + )); +} + +struct UnqualifiedStructCatalog; + +impl CatalogProvider for UnqualifiedStructCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + match table.table.as_str() { + "first_source" => Some(vec!["event".into()]), + "second_source" => Some(vec!["other".into()]), + _ => None, + } + } + + fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option { + (column == "event") + .then(|| { + candidates + .iter() + .find(|table| table.table == "first_source") + .cloned() + }) + .flatten() + } +} + +#[test] +fn unqualified_struct_field_access_uses_catalog_owner_for_ambiguous_root() { + let result = analyze( + "SELECT event.field AS value FROM first_source JOIN second_source ON first_source.id = second_source.id", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(UnqualifiedStructCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .remove(0); + let m = find_mapping(&result.columns.mappings, "value"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.table, "first_source"); + assert_eq!(column, "event"); + } + other => panic!("expected catalog-resolved source, got {other:?}"), + } +} + +struct RowValueCatalog; + +impl CatalogProvider for RowValueCatalog { + fn list_columns(&self, table: &TableRef) -> Option> { + (table.table == "source_table").then(|| vec!["source".into()]) + } + + fn resolve_column(&self, column: &str, candidates: &[TableRef]) -> Option { + (column == "source") + .then(|| candidates.first().cloned()) + .flatten() + } +} + +#[test] +fn bigquery_row_value_alias_prefers_catalog_column_with_same_name() { + let result = analyze( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + AnalyzeOptions { + dialect: Dialect::BigQuery, + catalog: Some(Box::new(RowValueCatalog)), + ..AnalyzeOptions::default() + }, + ) + .expect("SQL should parse") + .remove(0); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("source_table".into(), "source".into())] + ); +} + +#[test] +fn bigquery_row_value_alias_prefers_cte_output_column_with_same_name() { + let result = analyze_with_dialect( + "WITH source AS (SELECT source_table AS source FROM base) SELECT ARRAY_AGG(source) AS event FROM source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("base".into(), "source_table".into())] + ); +} + +#[test] +fn bigquery_row_value_alias_prefers_derived_output_column_with_same_name() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM (SELECT source_table AS source FROM base) AS source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("base".into(), "source_table".into())] + ); +} + +#[test] +fn bigquery_row_value_relation_alias_is_not_a_column() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + Dialect::BigQuery, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "source" && candidates.is_empty() + )); +} + +#[test] +fn postgresql_row_value_relation_alias_is_not_a_column() { + let result = analyze_with_dialect( + "SELECT ARRAY_AGG(source) AS event FROM source_table AS source", + Dialect::PostgreSql, + ); + let m = find_mapping(&result.columns.mappings, "event"); + assert!(matches!( + m.sources.as_slice(), + [ColumnOrigin::Ambiguous { column, candidates }] + if column == "source" && candidates.is_empty() + )); +} + +#[test] +fn generic_row_value_relation_alias_preserves_existing_behavior() { + let result = analyze_one("SELECT ARRAY_AGG(source) AS event FROM source_table AS source"); + let m = find_mapping(&result.columns.mappings, "event"); + assert_eq!( + concrete_sources(m), + vec![("source_table".into(), "source".into())] + ); +} From deb8ab51434e6715144f93ed9b1dbeecc9dad4a9 Mon Sep 17 00:00:00 2001 From: eitsupi Date: Mon, 24 Aug 2026 00:57:43 +0000 Subject: [PATCH 14/14] fix: preserve unbound qualified references --- sqllineage/src/build/expr.rs | 44 ++++++++++++++++++++++++++++++ sqllineage/tests/expr_coverage.rs | 45 +++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/sqllineage/src/build/expr.rs b/sqllineage/src/build/expr.rs index d6eb32d..6dcbabf 100644 --- a/sqllineage/src/build/expr.rs +++ b/sqllineage/src/build/expr.rs @@ -394,6 +394,50 @@ impl LineageBuilder { return vec![self.add_bound_field_ancestor(qualifier, column, Some(binding))]; } + // With no visible relation bindings, preserve the traditional + // qualified-reference fallback. This is used by callers that feed + // already-qualified expressions without a FROM clause (for example + // `orders.id`): the qualifier is still a physical relation name, + // rather than an unqualified struct root. + if self + .graph + .scopes + .visible_bindings(self.current_scope) + .is_empty() + && parts.len() >= 2 + { + let relation_parts = &parts[..parts.len() - 1]; + let qualifier = parts[..parts.len() - 1] + .iter() + .map(|part| part.value.as_str()) + .collect::>() + .join("."); + let binding = match relation_parts { + [table] => Some(Binding::Table(TableRef::new(table.value.clone()))), + [schema, table] => Some(Binding::Table(TableRef::with_schema( + schema.value.clone(), + table.value.clone(), + ))), + [catalog, schema, table] => Some(Binding::Table(TableRef { + catalog: Some(catalog.value.clone()), + schema: Some(schema.value.clone()), + table: table.value.clone(), + })), + // Keep the legacy display-only fallback for an unsupported + // number of relation components. Standard SQL relation + // names are at most catalog.schema.table, and this branch + // avoids inventing a lossy structured interpretation beyond + // that shape. + _ => None, + }; + return vec![self.graph.add_ref_with_binding( + parts[parts.len() - 1].value.clone(), + Some(qualifier), + self.current_scope, + binding, + )]; + } + // No relation prefix was found: `struct.field` is an unqualified // top-level column followed by a nested field path. Only the // top-level column can be represented by the public ColumnOrigin API. diff --git a/sqllineage/tests/expr_coverage.rs b/sqllineage/tests/expr_coverage.rs index 6f6276c..3a97f47 100644 --- a/sqllineage/tests/expr_coverage.rs +++ b/sqllineage/tests/expr_coverage.rs @@ -176,6 +176,51 @@ fn unqualified_compound_field_access_uses_top_level_column() { assert_eq!(concrete_sources(m), vec![("t".into(), "payload".into())]); } +#[test] +fn compound_identifier_without_visible_binding_keeps_qualified_relation_fallback() { + let result = analyze_one("SELECT orders.id AS id"); + let m = find_mapping(&result.columns.mappings, "id"); + match m.sources.as_slice() { + [ColumnOrigin::Concrete { table, column }] => { + assert_eq!(table.catalog, None); + assert_eq!(table.schema, None); + assert_eq!(table.table, "orders"); + assert_eq!(column, "id"); + } + other => panic!("expected structured orders.id source, got {other:?}"), + } +} + +#[test] +fn compound_identifier_without_visible_binding_preserves_relation_parts() { + for (sql, catalog, schema, table) in [ + ("SELECT raw.orders.id AS id", None, Some("raw"), "orders"), + ( + "SELECT warehouse.raw.orders.id AS id", + Some("warehouse"), + Some("raw"), + "orders", + ), + ] { + let result = analyze_one(sql); + let m = find_mapping(&result.columns.mappings, "id"); + match m.sources.as_slice() { + [ + ColumnOrigin::Concrete { + table: source_table, + column, + }, + ] => { + assert_eq!(source_table.catalog.as_deref(), catalog); + assert_eq!(source_table.schema.as_deref(), schema); + assert_eq!(source_table.table, table); + assert_eq!(column, "id"); + } + other => panic!("expected structured relation source, got {other:?}"), + } + } +} + #[test] fn bigquery_offset_compound_field_access_uses_binding_column() { let result = analyze_with_dialect(