fix: classify source-free aggregate calls correctly (COUNT(*))

sqllineage/src/build/expr.rs

@@ -303,7 +303,11 @@ impl LineageBuilder {
     }
 }
 
-pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind {
+/// Classify what an expression does to the values it reads.
+///
+/// The answer is stamped on the output column it defines and on every edge it
+/// draws to an ancestor, so it survives an expression that reads no column.
+pub(crate) fn classify_expr(expr: &Expr) -> EdgeKind {
     match expr {
         Expr::Identifier(_) | Expr::CompoundIdentifier(_) | Expr::Value(_) => EdgeKind::Direct,
         Expr::Function(f) => {
@@ -317,7 +321,7 @@ pub(crate) fn determine_edge_kind(expr: &Expr) -> EdgeKind {
         }
         Expr::Case { .. } => EdgeKind::ViaConditional,
         Expr::Cast { expr, .. } | Expr::Nested(expr) | Expr::Collate { expr, .. } => {
-            determine_edge_kind(expr)
+            classify_expr(expr)
         }
         _ => EdgeKind::ViaExpression,
     }

sqllineage/src/build/select.rs

@@ -3,7 +3,7 @@ use sqlparser::ast::{
 };
 
 use crate::build::LineageBuilder;
-use crate::build::expr::determine_edge_kind;
+use crate::build::expr::classify_expr;
 use crate::graph::scope::{Binding, ScopeColumn, ScopeKind};
 
 impl LineageBuilder {
@@ -22,9 +22,9 @@ impl LineageBuilder {
             match item {
                 SelectItem::UnnamedExpr(expr) => {
                     let ancestors = self.collect_ancestors(expr);
-                    let kind = determine_edge_kind(expr);
+                    let kind = classify_expr(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());
                     }
@@ -38,9 +38,9 @@ impl LineageBuilder {
                 }
                 SelectItem::ExprWithAlias { expr, alias } => {
                     let ancestors = self.collect_ancestors(expr);
-                    let kind = determine_edge_kind(expr);
+                    let kind = classify_expr(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());
                     }
@@ -54,10 +54,10 @@ impl LineageBuilder {
                 }
                 SelectItem::ExprWithAliases { expr, aliases } => {
                     let ancestors = self.collect_ancestors(expr);
-                    let kind = determine_edge_kind(expr);
+                    let kind = classify_expr(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());
                         }

sqllineage/src/build/statement.rs

@@ -1,7 +1,7 @@
 use sqlparser::ast::{self, AssignmentTarget, FunctionArguments, Ident, MergeAction, Statement};
 
 use crate::build::LineageBuilder;
-use crate::build::expr::determine_edge_kind;
+use crate::build::expr::classify_expr;
 use crate::graph::scope::ScopeColumn;
 use crate::types::{StatementType, TableRef};
 
@@ -49,8 +49,8 @@ impl LineageBuilder {
                 for assignment in &update.assignments {
                     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 kind = classify_expr(&assignment.value);
+                    let output = self.graph.add_output(col_name.clone(), kind.clone());
                     for &anc in &ancestors {
                         self.graph.add_edge(anc, output, kind.clone());
                     }
@@ -101,8 +101,9 @@ impl LineageBuilder {
                                 for assignment in assignments {
                                     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 kind = classify_expr(&assignment.value);
+                                    let output =
+                                        self.graph.add_output(col_name.clone(), kind.clone());
                                     for &anc in &ancestors {
                                         self.graph.add_edge(anc, output, kind.clone());
                                     }
@@ -140,8 +141,9 @@ impl LineageBuilder {
                                             .cloned()
                                             .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 kind = classify_expr(expr);
+                                        let output =
+                                            self.graph.add_output(col_name.clone(), kind.clone());
                                         for &anc in &ancestors {
                                             self.graph.add_edge(anc, output, kind.clone());
                                         }

sqllineage/src/graph/mod.rs

@@ -31,8 +31,8 @@ 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, kind: EdgeKind) -> NodeId {
+        self.add_node(RawNode::Output { name, kind })
     }
 
     pub fn add_ref(&mut self, name: String, qualifier: Option<String>, scope: ScopeId) -> NodeId {

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,
+        /// What the defining expression does, independent of what it draws
+        /// from. An expression that reaches no column still has one —
+        /// `COUNT(*)` has no column ancestor but is still an aggregate — so
+        /// this is where the classification survives when no edge carries it.
+        kind: EdgeKind,
+    },
     /// Named reference — alias, CTE reference, derived table column.
     Ref {
         name: String,

sqllineage/src/resolve/mod.rs

@@ -59,11 +59,11 @@ pub(crate) fn resolve(
     for col in &ordered_cols {
         let node_id = col.node_id;
         match &graph.nodes[node_id] {
-            RawNode::Output { name, .. } => {
+            RawNode::Output { name, kind } => {
                 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(kind, &edge_kinds);
 
                 if has_back {
                     mappings.push(ColumnMapping {
@@ -209,8 +209,8 @@ fn expand_scope_columns(
         return;
     }
     for col in graph.scopes.output_columns(scope_id) {
-        if let RawNode::Star { table, scope } = &graph.nodes[col.node_id] {
-            expand_star(
+        match &graph.nodes[col.node_id] {
+            RawNode::Star { table, scope } => expand_star(
                 table.as_ref(),
                 *scope,
                 graph,
@@ -219,20 +219,22 @@ fn expand_scope_columns(
                 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 transform = derive_transform(&edge_kinds);
-            mappings.push(ColumnMapping {
-                target: ColumnRef {
-                    table: output_table.cloned(),
-                    column: col.name.clone(),
-                },
-                sources,
-                transform,
-            });
+            ),
+            RawNode::Output { kind, .. } => {
+                let mut visited = HashSet::new();
+                let (sources, edge_kinds, _) =
+                    collect_output_sources(col.node_id, graph, resolved, incoming, &mut visited);
+                let transform = derive_transform(kind, &edge_kinds);
+                mappings.push(ColumnMapping {
+                    target: ColumnRef {
+                        table: output_table.cloned(),
+                        column: col.name.clone(),
+                    },
+                    sources,
+                    transform,
+                });
+            }
+            _ => {}
         }
     }
 }
@@ -511,12 +513,22 @@ fn resolve_through_scope(
     }
 }
 
-fn derive_transform(kinds: &[EdgeKind]) -> TransformKind {
-    if kinds.iter().any(|k| matches!(k, EdgeKind::ViaAggregation)) {
+/// Classify a column from its own kind together with the kinds of the edges
+/// that reached a source.
+///
+/// The two normally say the same thing — a projection stamps one kind on
+/// itself and on every edge it draws. They differ at a set operation, where
+/// the edges redirected from the other branch carry that branch's kind, and
+/// wherever a branch reached no source at all and so left no edge behind.
+/// Both have to count, so this is a union and not a preference.
+fn derive_transform(own_kind: &EdgeKind, edge_kinds: &[EdgeKind]) -> TransformKind {
+    let has = |probe: fn(&EdgeKind) -> bool| probe(own_kind) || edge_kinds.iter().any(probe);
+
+    if has(|k| matches!(k, EdgeKind::ViaAggregation)) {
         TransformKind::Aggregation
-    } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaConditional)) {
+    } else if has(|k| matches!(k, EdgeKind::ViaConditional)) {
         TransformKind::Conditional
-    } else if kinds.iter().any(|k| matches!(k, EdgeKind::ViaExpression)) {
+    } else if has(|k| matches!(k, EdgeKind::ViaExpression)) {
         TransformKind::Expression
     } else {
         TransformKind::Direct

sqllineage/tests/column_lineage.rs

@@ -68,6 +68,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");
@@ -153,3 +162,56 @@ fn proven_and_unresolved_sources_coexist_in_one_mapping() {
         m.sources
     );
 }
+
+/// A source-free projection is classified by what the expression is, not by
+/// the fact that it has no ancestors. A literal really is a direct value; a
+/// function call or an operator is not.
+#[test]
+fn source_free_projections_are_classified_by_their_own_kind() {
+    for (sql, expected) in [
+        ("SELECT 1 AS c FROM t", TransformKind::Direct),
+        ("SELECT NULL AS c FROM t", TransformKind::Direct),
+        ("SELECT CAST(1 AS INT) AS c FROM t", TransformKind::Direct),
+        ("SELECT 1 + 2 AS c FROM t", TransformKind::Expression),
+        ("SELECT NOW() AS c FROM t", TransformKind::Expression),
+        ("SELECT COUNT(1) AS c FROM t", TransformKind::Aggregation),
+        (
+            "SELECT CASE WHEN 1 = 1 THEN 2 ELSE 3 END AS c FROM t",
+            TransformKind::Conditional,
+        ),
+    ] {
+        let result = analyze_one(sql);
+        let m = find_mapping(&result.columns.mappings, "c");
+        assert!(m.sources.is_empty(), "{sql}: expected no sources");
+        assert_eq!(m.transform, expected, "{sql}");
+    }
+}
+
+/// A set operation classifies the column from every branch, including one that
+/// reached no source and so left no edge behind — the branch's own kind still
+/// counts.
+#[test]
+fn a_set_operation_is_classified_by_every_branch() {
+    for (sql, expected) in [
+        (
+            "SELECT COUNT(*) AS c FROM t UNION ALL SELECT a FROM u",
+            TransformKind::Aggregation,
+        ),
+        (
+            "SELECT a AS c FROM t UNION ALL SELECT SUM(b) FROM u",
+            TransformKind::Aggregation,
+        ),
+        (
+            "SELECT a AS c FROM t UNION ALL SELECT b + 1 FROM u",
+            TransformKind::Expression,
+        ),
+        (
+            "SELECT a AS c FROM t UNION ALL SELECT b FROM u",
+            TransformKind::Direct,
+        ),
+    ] {
+        let result = analyze_one(sql);
+        let m = find_mapping(&result.columns.mappings, "c");
+        assert_eq!(m.transform, expected, "{sql}");
+    }
+}

sqllineage/tests/cte.rs

@@ -314,3 +314,34 @@ fn column_behind_an_unexpanded_star_is_unresolved() {
         );
     }
 }
+
+/// Known limitation: only the last hop classifies the column.
+///
+/// `collect_output_sources` records the kind of the *immediate* incoming edge
+/// and drops every kind met deeper in the walk, so an aggregate below a CTE or
+/// a derived table reads as `Direct`. The column's own kind does not rescue
+/// this — it describes the outer projection, which really is a plain
+/// reference; the aggregation is a hop further down and never travels.
+///
+/// Pinned so the change is visible when that lands.
+#[test]
+fn only_the_last_hop_classifies_the_column() {
+    for sql in [
+        "WITH x AS (SELECT COUNT(*) AS c FROM t) SELECT c FROM x",
+        "SELECT c FROM (SELECT COUNT(*) AS c FROM t) d",
+        // Not a question of missing sources: `t.x` survives the hop, the
+        // aggregation does not.
+        "SELECT c FROM (SELECT SUM(x) AS c FROM t) d",
+    ] {
+        let result = analyze_one(sql);
+        let m = find_mapping(&result.columns.mappings, "c");
+        assert_eq!(m.transform, TransformKind::Direct, "{sql}");
+    }
+
+    // Expanding a star takes no second hop — it classifies each inner column
+    // from that column's own node — so the same query disagrees with itself
+    // depending on how the column is selected.
+    let result = analyze_one("SELECT * FROM (SELECT COUNT(*) AS c FROM t) d");
+    let m = find_mapping(&result.columns.mappings, "c");
+    assert_eq!(m.transform, TransformKind::Aggregation);
+}