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..a25daa0 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,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 { + 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, 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..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!()), }) @@ -1099,14 +1109,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 +1127,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 +1145,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 +1369,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 +1406,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..acdbce9 --- /dev/null +++ b/python/python/tests/test_union.py @@ -0,0 +1,117 @@ +# 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"} + + +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, + ]