From eefd2ebc395e629c2ae29ce36380a72cd8a4b9bf Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:08:11 -0700 Subject: [PATCH 1/2] feat(cypher): execute union queries with isolated scopes Parse UNION and UNION ALL as independent branches, validate matching output columns and compatible Arrow types, and preserve distinct versus bag semantics. Include all branches in query metadata and document direct Rust AST migration. Part of #20. Submitted by David Hughes (@stugorf), with AI Collaboration. Full combined validation and SF1 measurements remain pending. --- crates/lance-graph/README.md | 10 ++ crates/lance-graph/src/ast.rs | 32 +++++++ .../src/datafusion_planner/analysis.rs | 3 +- .../datafusion_planner/builder/basic_ops.rs | 89 +++++++++++++++++ .../src/datafusion_planner/builder/helpers.rs | 6 +- .../src/datafusion_planner/builder/mod.rs | 3 + crates/lance-graph/src/logical_plan.rs | 18 +++- .../lance-graph/src/parameter_substitution.rs | 3 + crates/lance-graph/src/parser.rs | 40 +++++++- crates/lance-graph/src/query.rs | 95 +++++++++++++----- crates/lance-graph/src/semantic.rs | 36 +++++++ docs/api-reference.md | 2 +- docs/cypher.md | 12 ++- python/python/tests/test_union.py | 96 +++++++++++++++++++ 14 files changed, 416 insertions(+), 29 deletions(-) create mode 100644 python/python/tests/test_union.py diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index e5172fb..66fe955 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -134,6 +134,16 @@ namespace utilities. Most APIs return `Result`. 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`: 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 diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 7d4e183..6804d83 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -56,9 +56,27 @@ pub struct CypherQuery { pub order_by: Option, /// SKIP/OFFSET clause (optional) pub skip: Option, + /// Following UNION branches, evaluated in order with independent variable scopes. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub union_branches: Vec, +} + +/// 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 { 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 } diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 3329517..d77edb6 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/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)?; } diff --git a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs index c10b116..3d38d04 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/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,67 @@ 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 { + 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!()), + }); + } + } + let union = LogicalPlanBuilder::from(left) + .union(right) + .map_err(|e| self.plan_error("Failed to build query UNION", e))?; + 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, diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 34647d5..7c7501e 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/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); } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index c2ca6a4..3fc8ccb 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/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, diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index a85a642..0b8612c 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/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, + right: Box, + 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)) diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index d0a0df4..61ba9e7 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -7,6 +7,9 @@ pub fn substitute_parameters( query: &mut CypherQuery, parameters: &HashMap, ) -> 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)?; diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 260240b..ac42ecc 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -36,6 +36,11 @@ pub fn parse_cypher_query(input: &str) -> Result { } 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 { 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(); diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index b99f200..817d82a 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1099,14 +1099,15 @@ impl CypherQuery { pub fn referenced_node_labels(&self) -> Vec { let mut labels = Vec::new(); - for clause in &self.ast.reading_clauses { - if let ReadingClause::Match(match_clause) = clause { - for pattern in &match_clause.patterns { - self.collect_node_labels_from_pattern(pattern, &mut labels); + self.ast.for_each_query(&mut |ast| { + for clause in &ast.reading_clauses { + if let ReadingClause::Match(match_clause) = clause { + for pattern in &match_clause.patterns { + self.collect_node_labels_from_pattern(pattern, &mut labels); + } } } - } - + }); labels.sort(); labels.dedup(); labels @@ -1116,14 +1117,15 @@ impl CypherQuery { pub fn referenced_relationship_types(&self) -> Vec { let mut types = Vec::new(); - for clause in &self.ast.reading_clauses { - if let ReadingClause::Match(match_clause) = clause { - for pattern in &match_clause.patterns { - self.collect_relationship_types_from_pattern(pattern, &mut types); + self.ast.for_each_query(&mut |ast| { + for clause in &ast.reading_clauses { + if let ReadingClause::Match(match_clause) = clause { + for pattern in &match_clause.patterns { + self.collect_relationship_types_from_pattern(pattern, &mut types); + } } } - } - + }); types.sort(); types.dedup(); types @@ -1133,19 +1135,20 @@ impl CypherQuery { pub fn variables(&self) -> Vec { let mut variables = Vec::new(); - for clause in &self.ast.reading_clauses { - match clause { - ReadingClause::Match(match_clause) => { - for pattern in &match_clause.patterns { - self.collect_variables_from_pattern(pattern, &mut variables); + self.ast.for_each_query(&mut |ast| { + for clause in &ast.reading_clauses { + match clause { + ReadingClause::Match(match_clause) => { + for pattern in &match_clause.patterns { + self.collect_variables_from_pattern(pattern, &mut variables); + } + } + ReadingClause::Unwind(unwind_clause) => { + variables.push(unwind_clause.alias.clone()); } - } - ReadingClause::Unwind(unwind_clause) => { - variables.push(unwind_clause.alias.clone()); } } - } - + }); variables.sort(); variables.dedup(); variables @@ -1356,6 +1359,7 @@ impl CypherQueryBuilder { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: crate::ast::ReturnClause { distinct: self.distinct, items: self.return_items, @@ -1392,6 +1396,53 @@ mod tests { use super::*; use crate::config::GraphConfig; + #[tokio::test] + async fn union_distinct_all_and_schema_validation_execute() { + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])), + vec![Arc::new(Int64Array::from(vec![1, 2, 2]))], + ) + .unwrap(); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for (query, rows) in [ + ("MATCH (p:Person) RETURN p.id AS n UNION RETURN 2 AS n", 2), + ( + "MATCH (p:Person) RETURN p.id AS n UNION ALL RETURN 2 AS n", + 4, + ), + ( + "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", + 2, + ), + ("RETURN NULL AS n UNION RETURN NULL AS n", 1), + ] { + let result = CypherQuery::new(query) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .unwrap(); + assert_eq!(result.num_rows(), rows, "{query}"); + } + for query in [ + "RETURN 1 AS x UNION RETURN 2 AS y", + "RETURN 1 AS x UNION RETURN 'a' AS x", + ] { + assert!(CypherQuery::new(query) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .is_err()); + } + } + #[test] fn test_parse_simple_cypher_query() { let query = CypherQuery::new("MATCH (n:Person) RETURN n.name").unwrap(); diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index ac30dff..bfef949 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -192,6 +192,19 @@ impl SemanticAnalyzer { // Phase 8: Type checking self.validate_types(&mut errors); + for branch in &mut analyzed_query.union_branches { + let mut analyzer = SemanticAnalyzer::new(self.config.clone()); + let analyzed = analyzer.analyze(&branch.query, parameters)?; + branch.query = analyzed.ast; + errors.extend( + analyzed + .errors + .into_iter() + .map(|e| format!("UNION branch: {e}")), + ); + warnings.extend(analyzed.warnings); + } + Ok(SemanticResult { ast: analyzed_query, variables: self.variables.clone(), @@ -993,6 +1006,18 @@ mod tests { .unwrap() } + #[test] + fn union_branches_do_not_inherit_variables() { + let query = crate::parser::parse_cypher_query( + "MATCH (p:Person) RETURN p.id AS id UNION RETURN p.id AS id", + ) + .unwrap(); + let result = SemanticAnalyzer::new(test_config()) + .analyze(&query, &HashMap::new()) + .unwrap(); + assert!(result.errors.iter().any(|e| e.contains("UNION branch"))); + } + // Helper: analyze a query that only has a single RETURN expression fn analyze_return_expr(expr: ValueExpression) -> Result { let query = CypherQuery { @@ -1002,6 +1027,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { @@ -1035,6 +1061,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { @@ -1071,6 +1098,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1124,6 +1152,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1159,6 +1188,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1205,6 +1235,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1235,6 +1266,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1278,6 +1310,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1327,6 +1360,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1398,6 +1432,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1442,6 +1477,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { diff --git a/docs/api-reference.md b/docs/api-reference.md index 4972240..b000671 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -197,7 +197,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde : The query text and the bound parameters. `#!python variables() -> list[str]`, `#!python node_labels() -> list[str]`, `#!python relationship_types() -> list[str]` -: Names referenced by the query. +: Names referenced by the query, including independently scoped UNION branches. List ordering is not guaranteed. --- diff --git a/docs/cypher.md b/docs/cypher.md index c980728..f8c7655 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -190,12 +190,13 @@ read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). | `WHERE` | After `MATCH`, `OPTIONAL MATCH` (filters inside the optional part) or `WITH`. | | `WITH [DISTINCT]` | With `ORDER BY`, `LIMIT` and `WHERE`; chain as many as you need. No `SKIP` inside `WITH`. | | `UNWIND list AS x` | | +| `UNION`, `UNION ALL` | Combine independent queries with matching output names/order and compatible Arrow types. | | `RETURN [DISTINCT]` | With `AS` aliases. Without reading clauses, expressions run on one implicit row. | | `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take nonnegative integer literals or parameters. | !!! failure "Not supported" - Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, `CALL`, and `FOREACH`. + Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `CALL`, and `FOREACH`. Integer list literals retain signed 64-bit values: `RETURN [1, 2, 3] AS ids` returns integers, and `UNWIND [1131, 19791209300775] AS id RETURN id` preserves both IDs exactly. Numeric literals with a decimal @@ -203,6 +204,15 @@ point or exponent, such as `[0.1, 1e-3]`, form float32 vectors. Whole-number lis vectors; conversion to float32 happens inside the vector operation. A configured engine is still required for standalone `RETURN` queries. +`UNION` removes duplicate rows; `UNION ALL` preserves them. Branches have independent variable scopes and +keep their own `ORDER BY`, `SKIP` and `LIMIT`. Combined output has no guaranteed order. Mixed UNION modes +are evaluated from left to right. Numeric columns can promote to a common numeric type; incompatible +column counts, names or types raise `ValueError`. This follows the existing homogeneous Arrow value model. + +```python +engine.execute("RETURN 1 AS n UNION RETURN 1 AS n").to_pylist() # [{'n': 1}] +``` + ### Patterns ```cypher diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py new file mode 100644 index 0000000..73ad036 --- /dev/null +++ b/python/python/tests/test_union.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Language UNION has independent scopes and a strict output-column contract.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +CONFIG = ( + GraphConfig.builder().node_table("Person", "id").node_table("City", "id").build() +) +TABLES = {"Person": pa.table({"id": [1, 2, 2]}), "City": pa.table({"id": [2, 3]})} + + +@pytest.fixture(params=[False, True]) +def engine(request): + return CypherEngine(CONFIG, TABLES, use_csr=request.param) + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN 1 AS n UNION RETURN 1 AS n", [1]), + ("RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), + ("RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", [1, 2]), + ("RETURN 1 AS n UNION RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), + ( + "MATCH (p:Person) RETURN p.id AS n UNION MATCH (p:City) RETURN p.id AS n", + [1, 2, 3], + ), + ( + "MATCH (p:Person) RETURN p.id AS n UNION ALL " + "MATCH (p:City) RETURN p.id AS n", + [1, 2, 2, 2, 3], + ), + ("RETURN NULL AS n UNION RETURN NULL AS n", [None]), + ("RETURN NULL AS n UNION ALL RETURN 1 AS n", [None, 1]), + ("RETURN 1 AS n LIMIT 0 UNION RETURN 2 AS n", [2]), + ( + "UNWIND [1, 2] AS n RETURN n ORDER BY n DESC LIMIT 1 " + "UNION ALL RETURN 3 AS n", + [2, 3], + ), + ("RETURN 'UNION ALL' AS n UNION RETURN 'UNION' AS n", ["UNION", "UNION ALL"]), + ], +) +def test_union_results(engine, query, expected): + actual = engine.execute(query).column(0).to_pylist() + assert sorted(actual, key=lambda v: (v is not None, v)) == sorted( + expected, key=lambda v: (v is not None, v) + ) + assert engine.execute_batches(query).read_all().num_rows == len(expected) + + +def test_numeric_promotion_and_parameters(engine): + query = "RETURN $a AS n UNION ALL RETURN $b AS n" + assert sorted(engine.execute(query, {"a": 1, "b": 2.5}).column(0).to_pylist()) == [ + 1.0, + 2.5, + ] + one_shot = ( + CypherQuery(query) + .with_config(CONFIG) + .with_parameter("a", 1) + .with_parameter("b", 2) + ) + assert sorted(one_shot.execute(TABLES).column(0).to_pylist()) == [1, 2] + + +@pytest.mark.parametrize( + "query", + [ + "RETURN 1 AS x UNION RETURN 2 AS y", + "RETURN 1 AS x UNION RETURN 2 AS x, 3 AS y", + "RETURN 1 AS x UNION RETURN 'a' AS x", + "WITH 1 AS x RETURN x UNION RETURN x", + "RETURN 1 AS x UNION ALL", + ], +) +def test_invalid_union(engine, query): + with pytest.raises(ValueError): + engine.execute(query) + + +def test_union_query_metadata(): + query = CypherQuery( + "MATCH (p:Person) RETURN p.id AS n UNION MATCH (c:City) RETURN c.id AS n" + ) + assert set(query.node_labels()) == {"City", "Person"} + assert query.variables() == ["c", "p"] + relationships = CypherQuery( + "MATCH (p:Person)-[:X]->(c:City) RETURN p.id AS n " + "UNION MATCH (c:City)-[:Y]->(p:Person) RETURN p.id AS n" + ) + assert set(relationships.relationship_types()) == {"X", "Y"} From 6e62a429da270649aa6877e81fc4a67df941bdc8 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:21:34 -0700 Subject: [PATCH 2/2] fix(query): collect UNION with its final output schema Use DataFusion UNION coercion before exposing child schemas and capture the physical result schema for collection. Preserve promoted numeric types and nullable outputs regardless of branch batch arrival order. Add repeated literal and parameter-null checks for collected and streamed results. The equivalent combined CALL/UNION implementation passed 10 Rust and 88 Python focused cases. Coordinator will validate this isolated commit on the standalone UNION branch. Submitted by David Hughes (@stugorf), with AI Collaboration. --- .../datafusion_planner/builder/basic_ops.rs | 15 +++++++-- crates/lance-graph/src/query.rs | 32 ++++++++++++------- python/python/tests/test_union.py | 21 ++++++++++++ 3 files changed, 54 insertions(+), 14 deletions(-) diff --git a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs index 3d38d04..a25daa0 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -153,9 +153,18 @@ impl DataFusionPlanner { }); } } - let union = LogicalPlanBuilder::from(left) - .union(right) - .map_err(|e| self.plan_error("Failed to build query UNION", e))?; + // 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 { diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 817d82a..b7ce175 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -624,14 +624,24 @@ impl CypherQuery { .dataframe_with_catalog_and_context(catalog, ctx) .await?; - // Get schema before collecting (in case result is empty) - let result_schema = df.schema().inner().clone(); - - // Collect results - let batches = df.collect().await.map_err(|e| GraphError::ExecutionError { - message: format!("Failed to collect query results: {}", e), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + // Capture the final schema after type coercion, exactly as DataFrame::collect + // creates its physical plan. UNION batches can differ in nullability, while + // the unoptimized logical schema can still have a narrower data type. + let task_context = Arc::new(df.task_ctx()); + let physical = df + .create_physical_plan() + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to plan query results: {e}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let result_schema = physical.schema(); + let batches = datafusion::physical_plan::collect(physical, task_context) + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to collect query results: {e}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; if batches.is_empty() { // Return empty batch with the schema from the DataFrame @@ -639,9 +649,9 @@ impl CypherQuery { return Ok(arrow::record_batch::RecordBatch::new_empty(result_schema)); } - // Combine all batches - let schema = batches[0].schema(); - concat_batches(&schema, &batches).map_err(|e| GraphError::ExecutionError { + // UNION branches may emit batches with narrower nullability. The complete + // result schema must govern concatenation regardless of which branch finishes first. + concat_batches(&result_schema, &batches).map_err(|e| GraphError::ExecutionError { message: format!("Failed to concatenate result batches: {}", e), location: snafu::Location::new(file!(), line!(), column!()), }) diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py index 73ad036..acdbce9 100644 --- a/python/python/tests/test_union.py +++ b/python/python/tests/test_union.py @@ -94,3 +94,24 @@ def test_union_query_metadata(): "UNION MATCH (c:City)-[:Y]->(p:Person) RETURN p.id AS n" ) assert set(relationships.relationship_types()) == {"X", "Y"} + + +def test_union_nullability_does_not_depend_on_first_batch(engine): + for query in [ + "RETURN null AS n UNION ALL RETURN 1 AS n", + "RETURN 1 AS n UNION ALL RETURN null AS n", + "RETURN $a AS n UNION ALL RETURN $b AS n", + ]: + for _ in range(5): + table = engine.execute(query, {"a": None, "b": 1}) + assert table.schema.field("n").nullable + assert sorted(table.column(0).to_pylist(), key=lambda n: n is None) == [ + 1, + None, + ] + stream = engine.execute_batches(query, {"a": None, "b": 1}).read_all() + assert stream.schema.field("n").nullable + assert sorted(stream.column(0).to_pylist(), key=lambda n: n is None) == [ + 1, + None, + ]