Add scoped UNION and UNION ALL queries

crates/lance-graph/README.md

@@ -134,6 +134,16 @@ namespace utilities.
 
 Most APIs return `Result<T, GraphError>`. Errors include parsing failures, missing mappings, and execution issues surfaced from DataFusion.
 
+## Constructing query ASTs directly
+
+Prefer `CypherQuery::new` or the query builder when you do not need direct AST access. Pagination fields
+now use `Option<ast::RowCount>`: wrap literal values with `RowCount::Literal(n)`, or use
+`RowCount::Parameter(name)` before parameter substitution. AST struct literals also set
+`union_branches: Vec::new()` for a single query. `UnionBranch` holds each following query and its `all`
+flag; each branch has its own variable scope. Exhaustive matches on `LogicalOperator` must handle the
+new `SingleRow` and `QueryUnion` variants. These are Rust source compatibility changes; Python method
+signatures are unchanged by UNION.
+
 ## Testing
 
 ```bash

crates/lance-graph/src/ast.rs

@@ -56,9 +56,27 @@ pub struct CypherQuery {
     pub order_by: Option<OrderByClause>,
     /// SKIP/OFFSET clause (optional)
     pub skip: Option<RowCount>,
+    /// Following UNION branches, evaluated in order with independent variable scopes.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub union_branches: Vec<UnionBranch>,
+}
+
+/// One right-hand query of UNION (distinct) or UNION ALL (bag concatenation).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct UnionBranch {
+    pub all: bool,
+    pub query: CypherQuery,
 }
 
 impl CypherQuery {
+    /// Visit each independently scoped query without cloning its AST.
+    pub(crate) fn for_each_query(&self, visit: &mut dyn FnMut(&Self)) {
+        visit(self);
+        for branch in &self.union_branches {
+            branch.query.for_each_query(visit);
+        }
+    }
+
     /// Extract all node labels referenced in the query
     pub fn get_node_labels(&self) -> Vec<String> {
         let mut labels = Vec::new();
@@ -92,6 +110,13 @@ impl CypherQuery {
                 }
             }
         }
+        for branch in &self.union_branches {
+            for label in branch.query.get_node_labels() {
+                if !labels.contains(&label) {
+                    labels.push(label);
+                }
+            }
+        }
         labels
     }
 
@@ -105,6 +130,13 @@ impl CypherQuery {
                 }
             }
         }
+        for branch in &self.union_branches {
+            for rel_type in branch.query.get_relationship_types() {
+                if !types.contains(&rel_type) {
+                    types.push(rel_type);
+                }
+            }
+        }
         types
     }
 

crates/lance-graph/src/datafusion_planner/analysis.rs

@@ -205,7 +205,8 @@ fn analyze_operator(
         | LogicalOperator::Distinct { input } => {
             analyze_operator(input, analysis, rel_counter)?;
         }
-        LogicalOperator::Join { left, right, .. } => {
+        LogicalOperator::Join { left, right, .. }
+        | LogicalOperator::QueryUnion { left, right, .. } => {
             analyze_operator(left, analysis, rel_counter)?;
             analyze_operator(right, analysis, rel_counter)?;
         }

crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs

@@ -9,6 +9,34 @@ use crate::error::Result;
 use crate::logical_plan::*;
 use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, SortExpr};
 
+fn union_compatible(left: &arrow_schema::DataType, right: &arrow_schema::DataType) -> bool {
+    use arrow_schema::DataType;
+    if left == right
+        || left.is_null()
+        || right.is_null()
+        || (left.is_numeric() && right.is_numeric())
+    {
+        return true;
+    }
+    let text =
+        |t: &DataType| matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View);
+    if text(left) && text(right) {
+        return true;
+    }
+    match (left, right) {
+        (DataType::List(left), DataType::List(right)) => {
+            union_compatible(left.data_type(), right.data_type())
+        }
+        (DataType::Struct(left), DataType::Struct(right)) => {
+            left.len() == right.len()
+                && left.iter().zip(right).all(|(l, r)| {
+                    l.name() == r.name() && union_compatible(l.data_type(), r.data_type())
+                })
+        }
+        _ => false,
+    }
+}
+
 impl DataFusionPlanner {
     /// UNION ALL of alternation branches. Each branch is planned with its own variable
     /// labels (the analysis keeps only one per variable), then all branches are projected
@@ -79,6 +107,76 @@ impl DataFusionPlanner {
             .map_err(|e| self.plan_error("Failed to build UNION", e))
     }
 
+    /// Language-level UNION validates the positional output contract before concatenation.
+    pub(crate) fn build_query_union(
+        &self,
+        ctx: &mut PlanningContext,
+        left: &LogicalOperator,
+        right: &LogicalOperator,
+        all: bool,
+    ) -> Result<LogicalPlan> {
+        let mut plans = Vec::with_capacity(2);
+        for input in [left, right] {
+            let mut analysis = ctx.analysis.clone();
+            Self::collect_labels(input, &mut analysis.var_to_label);
+            let mut branch = PlanningContext::new(&analysis);
+            branch.relationship_instance_idx = ctx.relationship_instance_idx.clone();
+            plans.push(self.build_operator(&mut branch, input)?);
+            ctx.relationship_instance_idx = branch.relationship_instance_idx;
+        }
+        let right = plans.pop().unwrap();
+        let left = plans.pop().unwrap();
+        let l = left.schema();
+        let r = right.schema();
+        if l.fields().len() != r.fields().len()
+            || l.fields()
+                .iter()
+                .zip(r.fields())
+                .any(|(l, r)| !l.name().eq_ignore_ascii_case(r.name()))
+        {
+            return Err(crate::error::GraphError::PlanError {
+                message: "UNION branches must return the same column names in the same order"
+                    .to_string(),
+                location: snafu::Location::new(file!(), line!(), column!()),
+            });
+        }
+        for (left, right) in l.fields().iter().zip(r.fields()) {
+            if !union_compatible(left.data_type(), right.data_type()) {
+                return Err(crate::error::GraphError::PlanError {
+                    message: format!(
+                        "UNION column '{}' has incompatible types {} and {}",
+                        left.name(),
+                        left.data_type(),
+                        right.data_type()
+                    ),
+                    location: snafu::Location::new(file!(), line!(), column!()),
+                });
+            }
+        }
+        // Consumers need the final output schema before collecting UNION results.
+        // DataFusion's UNION builder deliberately leaves coercion until analysis;
+        // perform its standard UNION coercion now, without optimizing parameter values.
+        let union = datafusion::logical_expr::Union::try_new_with_loose_types(vec![
+            std::sync::Arc::new(left),
+            std::sync::Arc::new(right),
+        ])
+        .and_then(
+            datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter::coerce_union,
+        )
+        .map_err(|e| self.plan_error("Failed to coerce query UNION", e))?;
+        let union = LogicalPlanBuilder::from(union);
+        let union = if all {
+            union
+        } else {
+            union
+                .distinct()
+                .map_err(|e| self.plan_error("Failed to deduplicate UNION", e))?
+        };
+        union
+            .build()
+            .map_err(|e| self.plan_error("Failed to build query UNION", e))
+    }
+
     pub(crate) fn build_filter(
         &self,
         ctx: &mut PlanningContext,

crates/lance-graph/src/datafusion_planner/builder/helpers.rs

@@ -81,7 +81,8 @@ impl DataFusionPlanner {
                 }
             }
             // Binary operator: recurse into both left and right
-            LogicalOperator::Join { left, right, .. } => {
+            LogicalOperator::Join { left, right, .. }
+            | LogicalOperator::QueryUnion { left, right, .. } => {
                 Self::collect_variables(left, vars);
                 Self::collect_variables(right, vars);
             }
@@ -153,7 +154,8 @@ impl DataFusionPlanner {
             | LogicalOperator::Limit { input, .. }
             | LogicalOperator::Offset { input, .. }
             | LogicalOperator::Distinct { input } => Self::collect_labels(input, labels),
-            LogicalOperator::Join { left, right, .. } => {
+            LogicalOperator::Join { left, right, .. }
+            | LogicalOperator::QueryUnion { left, right, .. } => {
                 Self::collect_labels(left, labels);
                 Self::collect_labels(right, labels);
             }

crates/lance-graph/src/datafusion_planner/builder/mod.rs

@@ -108,6 +108,9 @@ impl DataFusionPlanner {
                 label,
             } => self.build_unwind(ctx, input, expression, alias, label.as_deref()),
             LogicalOperator::Union { inputs } => self.build_union(ctx, inputs),
+            LogicalOperator::QueryUnion { left, right, all } => {
+                self.build_query_union(ctx, left, right, *all)
+            }
             LogicalOperator::ListComprehension {
                 input,
                 list,

crates/lance-graph/src/logical_plan.rs

@@ -20,6 +20,12 @@ use std::collections::HashMap;
 pub enum LogicalOperator {
     /// One row with no bound variables, for queries without reading clauses.
     SingleRow,
+    /// Language-level UNION; unlike pattern alternation, branches must share output schema.
+    QueryUnion {
+        left: Box<LogicalOperator>,
+        right: Box<LogicalOperator>,
+        all: bool,
+    },
     /// Scan all nodes with a specific label
     ScanByLabel {
         variable: String,
@@ -375,6 +381,15 @@ impl<'a> LogicalPlanner<'a> {
             };
         }
 
+        for branch in &query.union_branches {
+            let mut planner = LogicalPlanner::new(self.config).with_fan_out(self.fan_out);
+            let right = planner.plan(&branch.query)?;
+            plan = LogicalOperator::QueryUnion {
+                left: Box::new(plan),
+                right: Box::new(right),
+                all: branch.all,
+            };
+        }
         Ok(plan)
     }
 
@@ -1283,7 +1298,8 @@ impl<'a> LogicalPlanner<'a> {
             LogicalOperator::Sort { input, .. } => self.extract_variable_from_plan(input),
             LogicalOperator::Offset { input, .. } => self.extract_variable_from_plan(input),
             LogicalOperator::Limit { input, .. } => self.extract_variable_from_plan(input),
-            LogicalOperator::Join { left, right, .. } => {
+            LogicalOperator::Join { left, right, .. }
+            | LogicalOperator::QueryUnion { left, right, .. } => {
                 // Prefer the right branch's tail variable, else fall back to left
                 self.extract_variable_from_plan(right)
                     .or_else(|_| self.extract_variable_from_plan(left))

crates/lance-graph/src/parameter_substitution.rs

@@ -7,6 +7,9 @@ pub fn substitute_parameters(
     query: &mut CypherQuery,
     parameters: &HashMap<String, serde_json::Value>,
 ) -> Result<()> {
+    for branch in &mut query.union_branches {
+        substitute_parameters(&mut branch.query, parameters)?;
+    }
     substitute_row_count(&mut query.skip, parameters)?;
     substitute_row_count(&mut query.limit, parameters)?;
 

crates/lance-graph/src/parser.rs

@@ -36,6 +36,11 @@ pub fn parse_cypher_query(input: &str) -> Result<CypherQuery> {
     }
 
     let mut query = query;
+    normalize_inline_properties(&mut query);
+    Ok(query)
+}
+
+fn normalize_inline_properties(query: &mut CypherQuery) {
     let stages = query
         .further_with_stages
         .iter_mut()
@@ -51,7 +56,9 @@ pub fn parse_cypher_query(input: &str) -> Result<CypherQuery> {
             move_variable_properties_to_where(m, &mut anonymous);
         }
     }
-    Ok(query)
+    for branch in &mut query.union_branches {
+        normalize_inline_properties(&mut branch.query);
+    }
 }
 
 /// Inline node properties whose value is a variable or property from an earlier clause
@@ -106,6 +113,24 @@ fn move_variable_properties_to_where(clause: &mut MatchClause, anonymous: &mut u
 
 // Top-level parser for a complete Cypher query
 fn cypher_query(input: &str) -> IResult<&str, CypherQuery> {
+    let (input, mut query) = single_query(input)?;
+    let (input, branches) = many0(pair(union_header, single_query))(input)?;
+    query.union_branches = branches
+        .into_iter()
+        .map(|(all, query)| UnionBranch { all, query })
+        .collect();
+    Ok((input, query))
+}
+
+fn union_header(input: &str) -> IResult<&str, bool> {
+    let (input, _) = multispace0(input)?;
+    let (input, _) = tag_no_case("UNION")(input)?;
+    let (input, _) = multispace1(input)?;
+    let (input, all) = opt(terminated(tag_no_case("ALL"), multispace1))(input)?;
+    Ok((input, all.is_some()))
+}
+
+fn single_query(input: &str) -> IResult<&str, CypherQuery> {
     let (input, _) = multispace0(input)?;
     let (input, reading_clauses) = many0(reading_clause)(input)?;
     let (input, pre_with_where) = opt(where_clause)(input)?;
@@ -140,6 +165,7 @@ fn cypher_query(input: &str) -> IResult<&str, CypherQuery> {
             post_with_reading_clauses,
             post_with_where_clause: post_with_where,
             further_with_stages,
+            union_branches: Vec::new(),
             return_clause,
             limit,
             order_by,
@@ -2295,6 +2321,18 @@ mod tests {
         }
     }
 
+    #[test]
+    fn unions_preserve_modes_and_quoted_keywords() {
+        let query = parse_cypher_query(
+            "RETURN 'UNION' AS n UNION ALL RETURN 'ALL' AS n UNION RETURN 'UNION ALL' AS n",
+        )
+        .unwrap();
+        assert_eq!(query.union_branches.len(), 2);
+        assert!(query.union_branches[0].all);
+        assert!(!query.union_branches[1].all);
+        assert!(parse_cypher_query("RETURN 1 AS n UNION ALL").is_err());
+    }
+
     #[test]
     fn pagination_accepts_parameters_without_losing_names() {
         let query = parse_cypher_query("RETURN 1 AS n SKIP $offset LIMIT $count").unwrap();