From f2cc84c64898185bf8ea4826b276b61ea8b20b63 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 16:15:39 -0700 Subject: [PATCH] feat(cypher): support scoped read-only CALL subqueries --- crates/lance-graph/README.md | 4 + crates/lance-graph/src/apply.rs | 527 ++++++++++++++++++ crates/lance-graph/src/ast.rs | 64 ++- crates/lance-graph/src/csr_expand.rs | 9 +- .../src/datafusion_planner/analysis.rs | 17 + .../datafusion_planner/builder/basic_ops.rs | 4 +- .../datafusion_planner/builder/call_ops.rs | 79 +++ .../src/datafusion_planner/builder/helpers.rs | 14 + .../src/datafusion_planner/builder/mod.rs | 11 + crates/lance-graph/src/lib.rs | 1 + crates/lance-graph/src/logical_plan.rs | 143 ++++- .../lance-graph/src/parameter_substitution.rs | 1 + crates/lance-graph/src/parser.rs | 16 +- crates/lance-graph/src/path_value.rs | 32 +- crates/lance-graph/src/query.rs | 19 +- crates/lance-graph/src/semantic.rs | 90 +++ docs/api-reference.md | 2 +- docs/cypher.md | 47 +- docs/execution.md | 4 + python/python/tests/test_call_subqueries.py | 475 ++++++++++++++++ python/python/tests/test_path_values.py | 4 +- 21 files changed, 1544 insertions(+), 19 deletions(-) create mode 100644 crates/lance-graph/src/apply.rs create mode 100644 crates/lance-graph/src/datafusion_planner/builder/call_ops.rs create mode 100644 python/python/tests/test_call_subqueries.py diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index c55c314..cc6e5cc 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -187,3 +187,7 @@ See the Python package docs for setup and development: ## License Apache-2.0. See the top-level LICENSE file for details. + +## Rust AST migration for CALL + +Read-only subqueries add `ReadingClause::Call(Box)` and `LogicalOperator::{Apply, CallInput}`. Update exhaustive matches and visitors to recurse into the child query, including UNION branches. A leading plain importing `WITH` declares its available outer variables; other bindings stay in their own scope. diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs new file mode 100644 index 0000000..66e9fd1 --- /dev/null +++ b/crates/lance-graph/src/apply.rs @@ -0,0 +1,527 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Read-only CALL: one independently bound child stream per outer row. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion::common::{DFSchemaRef, Result, ScalarValue}; +use datafusion::execution::context::SessionState; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::EmissionType; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + Partitioning, PlanProperties, +}; +use futures::StreamExt; + +/// The child is deliberately unoptimized until its imports have concrete values. +/// It is not an ordinary input: planning it before binding would give global top-k +/// and aggregate semantics instead of evaluating them once per outer row. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct ApplyNode { + input: LogicalPlan, + inner: LogicalPlan, + /// (outer column name, private placeholder name). + imports: Vec<(String, String)>, + schema: DFSchemaRef, +} + +impl ApplyNode { + pub(crate) fn try_new( + input: LogicalPlan, + inner: LogicalPlan, + imports: Vec<(String, String)>, + ) -> Result { + let schema = Arc::new(input.schema().join(inner.schema())?); + Ok(Self { + input, + inner, + imports, + schema, + }) + } +} + +impl PartialOrd for ApplyNode { + fn partial_cmp(&self, other: &Self) -> Option { + self.input + .partial_cmp(&other.input) + .and_then(|order| match order { + Ordering::Equal => self.inner.partial_cmp(&other.inner), + order => Some(order), + }) + .filter(|order| *order != Ordering::Equal || self == other) + } +} + +impl UserDefinedLogicalNodeCore for ApplyNode { + fn name(&self) -> &str { + "Apply" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "Apply: imports={:?}, child={}", + self.imports + .iter() + .map(|(name, _)| name) + .collect::>(), + self.inner.display_indent() + ) + } + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + Self::try_new(inputs.remove(0), self.inner.clone(), self.imports.clone()) + } +} + +pub(crate) struct ApplyExec { + input: Arc, + inner: LogicalPlan, + imports: Vec<(String, String)>, + state: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl fmt::Debug for ApplyExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("ApplyExec") + .field("imports", &self.imports) + .finish_non_exhaustive() + } +} + +impl ApplyExec { + fn new( + input: Arc, + inner: LogicalPlan, + imports: Vec<(String, String)>, + state: Arc, + schema: SchemaRef, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + inner, + imports, + state, + schema, + properties, + } + } +} + +impl DisplayAs for ApplyExec { + fn fmt_as(&self, _format: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "ApplyExec: per-row CALL, {} imported columns", + self.imports.len() + ) + } +} + +impl ExecutionPlan for ApplyExec { + fn name(&self) -> &str { + "ApplyExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new( + children.remove(0), + self.inner.clone(), + self.imports.clone(), + self.state.clone(), + self.schema.clone(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.input.execute(partition, context.clone())?; + let indices = self + .imports + .iter() + .map(|(column, name)| Ok((input.schema().index_of(column)?, name.clone()))) + .collect::>>()?; + let state = ApplyStreamState { + input, + batch: None, + next_row: 0, + iterations_since_yield: 0, + current: vec![], + child: None, + inner: self.inner.clone(), + imports: indices, + session: self.state.clone(), + context, + schema: self.schema.clone(), + }; + let stream = futures::stream::try_unfold(state, |mut state| async move { + loop { + // Bound all ready work, including empty outer and child batches. + if state.iterations_since_yield == 64 { + crate::csr_expand::cooperative_yield().await; + state.iterations_since_yield = 0; + } + state.iterations_since_yield += 1; + if let Some(child) = &mut state.child { + if let Some(batch) = child.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + let mut columns = state + .current + .iter() + .map(|value| value.to_array_of_size(batch.num_rows())) + .collect::>>()?; + columns.extend_from_slice(batch.columns()); + let output = RecordBatch::try_new(state.schema.clone(), columns)?; + return Ok(Some((output, state))); + } + state.child = None; + } + if state + .batch + .as_ref() + .is_none_or(|batch| state.next_row >= batch.num_rows()) + { + let Some(batch) = state.input.next().await else { + return Ok(None); + }; + state.batch = Some(batch?); + state.next_row = 0; + if state + .batch + .as_ref() + .is_some_and(|batch| batch.num_rows() == 0) + { + continue; + } + } + let batch = state.batch.as_ref().expect("a nonempty outer batch"); + state.current = batch + .columns() + .iter() + .map(|array| ScalarValue::try_from_array(array, state.next_row)) + .collect::>>()?; + state.next_row += 1; + let parameters: HashMap<_, _> = state + .imports + .iter() + .map(|(index, name)| (name.clone(), state.current[*index].clone())) + .collect(); + let bound = state.inner.clone().with_param_values(parameters)?; + let physical = state.session.create_physical_plan(&bound).await?; + state.child = Some(execute_stream(physical, state.context.clone())?); + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } +} + +struct ApplyStreamState { + input: SendableRecordBatchStream, + batch: Option, + next_row: usize, + iterations_since_yield: usize, + current: Vec, + child: Option, + inner: LogicalPlan, + imports: Vec<(usize, String)>, + session: Arc, + context: Arc, + schema: SchemaRef, +} + +pub(crate) fn plan_apply( + node: &ApplyNode, + inputs: &[Arc], + state: &SessionState, +) -> Arc { + Arc::new(ApplyExec::new( + inputs[0].clone(), + node.inner.clone(), + node.imports.clone(), + Arc::new(state.clone()), + Arc::new(node.schema.as_arrow().clone()), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{CypherQuery, GraphConfig}; + use arrow_array::{Int64Array, StringArray}; + + fn graph() -> (GraphConfig, HashMap) { + let config = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let people = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as arrow_array::ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec!["A", "B", "C", "D"])), + ), + ("age", Arc::new(Int64Array::from(vec![20, 30, 40, 50]))), + ]) + .unwrap(); + let edges = RecordBatch::try_from_iter(vec![ + ( + "src", + Arc::new(Int64Array::from(vec![10, 10, 10, 20])) as arrow_array::ArrayRef, + ), + ("dst", Arc::new(Int64Array::from(vec![20, 20, 30, 40]))), + ]) + .unwrap(); + ( + config, + HashMap::from([("Person".into(), people), ("KNOWS".into(), edges)]), + ) + } + + async fn run(text: &str) -> crate::error::Result { + let (config, data) = graph(); + CypherQuery::new(text)? + .with_config(config) + .execute(data, None) + .await + } + + fn ints(batch: &RecordBatch, name: &str) -> Vec { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[tokio::test] + async fn call_empty_memory_batches_yield_before_end_of_stream() { + use datafusion::datasource::{MemTable, TableProvider}; + use datafusion::execution::context::SessionContext; + use datafusion::logical_expr::{lit, LogicalPlanBuilder}; + use futures::{Future, TryStreamExt}; + use std::task::{Context, Poll}; + + let batch = RecordBatch::try_from_iter(vec![( + "x", + Arc::new(Int64Array::from(vec![7])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let batches = vec![RecordBatch::new_empty(batch.schema()); 65]; + let table = MemTable::try_new(batch.schema(), vec![batches]).unwrap(); + let state = SessionContext::new().state(); + let input = table.scan(&state, None, &[], None).await.unwrap(); + let inner = LogicalPlanBuilder::empty(true) + .project(vec![lit(1_i64).alias("n")]) + .unwrap() + .build() + .unwrap(); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + batch.schema().field(0).clone(), + inner.schema().as_arrow().field(0).clone(), + ])); + let plan = ApplyExec::new(input, inner, vec![], Arc::new(state.clone()), schema); + let mut stream = plan.execute(0, state.task_ctx()).unwrap(); + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + // Isolate Apply's own bound from Tokio's task-wide scheduling budget. + let mut next = Box::pin(tokio::task::unconstrained(stream.next())); + assert!(matches!(next.as_mut().poll(&mut cx), Poll::Pending)); + drop(next); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert!(batches.is_empty()); + } + + #[tokio::test] + async fn call_per_row_top_k_and_entity_export() { + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 } RETURN p.id AS id, friend.id AS friend ORDER BY id").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 20]); + assert_eq!(ints(&result, "friend"), [30, 40]); + } + + #[tokio::test] + async fn call_aggregate_keeps_empty_child_and_outer_multiplicity() { + let result = run("MATCH (p:Person) UNWIND [1, 1] AS duplicate CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 10, 20, 20, 30, 30, 40, 40]); + assert_eq!(ints(&result, "n"), [3, 3, 1, 1, 0, 0, 0, 0]); + } + + #[tokio::test] + async fn call_nested_and_sibling_imports_have_separate_bindings() { + let (config, data) = graph(); + let result = CypherQuery::new("UNWIND [1, 2] AS x CALL { WITH x CALL { WITH x RETURN x + $__call_1_0 AS y } RETURN y + 1 AS z } CALL { WITH z RETURN z * 2 AS answer } RETURN x, answer ORDER BY x") + .unwrap().with_parameter("__call_1_0", 10).with_config(config).execute(data, None).await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 2]); + assert_eq!(ints(&result, "answer"), [24, 26]); + } + + #[tokio::test] + async fn call_empty_outer_never_evaluates_child() { + let result = + run("MATCH (p:Person) WHERE p.id < 0 CALL { RETURN 1 / 0 AS invalid } RETURN invalid") + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + } + + #[tokio::test] + async fn call_inner_names_are_lexical_and_exports_explicit() { + let result = run("UNWIND [1, 2] AS x CALL { UNWIND [10, 20] AS x RETURN x AS y } RETURN x, y ORDER BY x, y").await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 1, 2, 2]); + assert_eq!(ints(&result, "y"), [10, 20, 10, 20]); + for (query, message) in [ + ( + "UNWIND [1] AS x CALL { RETURN x AS y } RETURN y", + "Undefined variable", + ), + ( + "UNWIND [1] AS x CALL { RETURN 2 AS x } RETURN x", + "collides", + ), + ("CALL { RETURN 1 + 2 } RETURN 1", "need an alias"), + ( + "CALL { WITH missing RETURN missing AS x } RETURN x", + "undefined variable", + ), + ( + "UNWIND [1] AS x CALL { WITH x + 1 AS y RETURN y } RETURN y", + "importing WITH", + ), + ( + "CALL { UNWIND [1] AS local RETURN 2 AS y } RETURN local", + "Undefined variable", + ), + ] { + let error = run(query).await.unwrap_err().to_string(); + assert!(error.contains(message), "{query}: {error}"); + } + } + + #[tokio::test] + async fn call_union_branches_bind_independently_and_keep_outer_duplicates() { + let result = run("UNWIND [1, 1, 2] AS x CALL { WITH x RETURN x AS n UNION WITH x RETURN x AS n } RETURN x, n ORDER BY x").await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 1, 2]); + assert_eq!(ints(&result, "n"), [1, 1, 2]); + let result = run("WITH 2 AS x, 3 AS y CALL { WITH x RETURN x AS n UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n").await.unwrap(); + assert_eq!(ints(&result, "n"), [2, 3]); + let error = run("WITH 2 AS x CALL { WITH x RETURN x AS n UNION RETURN x AS n } RETURN n") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("Undefined variable"), "{error}"); + } + + #[tokio::test] + async fn call_union_node_exports_and_empty_aggregate_branches() { + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend UNION WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend } RETURN p.id AS id, friend.id AS friend ORDER BY id, friend").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 10, 20]); + assert_eq!(ints(&result, "friend"), [20, 30, 40]); + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n UNION ALL WITH p MATCH (p)-[:KNOWS]->(f:Person) WHERE f.id < 0 RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id, n").await.unwrap(); + assert_eq!(ints(&result, "n"), [0, 3, 0, 1, 0, 0, 0, 0]); + } + + #[tokio::test] + async fn call_union_uses_coerced_schema_for_numbers_and_nulls() { + let batch = run("CALL { RETURN 1 AS n UNION ALL RETURN 2.5 AS n } RETURN n ORDER BY n") + .await + .unwrap(); + assert_eq!( + batch.schema().field(0).data_type(), + &arrow_schema::DataType::Float64 + ); + for text in [ + "RETURN null AS n UNION ALL RETURN 1 AS n", + "RETURN 1 AS n UNION ALL RETURN null AS n", + "CALL { RETURN null AS n UNION ALL RETURN 1 AS n } RETURN n", + ] { + let batch = run(text).await.unwrap(); + assert_eq!( + batch.schema().field(0).data_type(), + &arrow_schema::DataType::Int64 + ); + assert!(batch.schema().field(0).is_nullable()); + assert_eq!(batch.column(0).null_count(), 1); + } + } + + #[test] + fn call_bypasses_snapshot_templates_in_every_reading_stage() { + let (config, _) = graph(); + let catalog = Arc::new(lance_graph_catalog::InMemoryCatalog::new()); + for text in [ + "CALL { RETURN 1 AS n } RETURN n", + "WITH 1 AS x CALL { WITH x RETURN x AS n } RETURN n", + "WITH 1 AS x WITH x AS y CALL { WITH y RETURN y AS n } RETURN n", + "RETURN 1 AS n UNION CALL { RETURN 2 AS n } RETURN n", + "RETURN 1 AS n UNION RETURN 2 AS n LIMIT $take", + "MATCH path = (p:Person)-[:KNOWS]->(f:Person) RETURN path", + "MATCH path = shortestPath((p:Person)-[:KNOWS*..3]->(f:Person)) RETURN relationships(path)", + ] { + let mut query = CypherQuery::new(text).unwrap().with_config(config.clone()); + assert!(!query.prepare_cached_plan(catalog.clone()).unwrap()); + } + } + + #[test] + fn call_parser_requires_returning_read_only_body() { + for query in [ + "CALL db.labels() RETURN 1", + "CALL { CREATE (n:Person) RETURN n } RETURN n", + "CALL { MATCH (n:Person) } RETURN n", + "CALL { RETURN 1 AS x } YIELD x RETURN x", + ] { + assert!(CypherQuery::new(query).is_err(), "{query}"); + } + } +} diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index c09925e..c3930b7 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -80,9 +80,14 @@ impl CypherQuery { ) } - /// Visit each independently scoped query without cloning its AST. + /// Visit query scopes, including UNION branches and nested CALL bodies. pub(crate) fn for_each_query(&self, visit: &mut dyn FnMut(&Self)) { visit(self); + for clause in self.all_reading_clauses() { + if let ReadingClause::Call(query) = clause { + query.for_each_query(visit); + } + } for branch in &self.union_branches { branch.query.for_each_query(visit); } @@ -93,6 +98,13 @@ impl CypherQuery { let mut labels = Vec::new(); // Iterate all match clauses directly for clause in self.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + for label in inner.get_node_labels() { + if !labels.contains(&label) { + labels.push(label); + } + } + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { match pattern { @@ -135,6 +147,13 @@ impl CypherQuery { pub fn get_relationship_types(&self) -> Vec { let mut types = Vec::new(); for clause in self.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + for kind in inner.get_relationship_types() { + if !types.contains(&kind) { + types.push(kind); + } + } + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { self.collect_relationship_types_from_pattern(pattern, &mut types); @@ -173,6 +192,49 @@ impl CypherQuery { pub enum ReadingClause { Match(MatchClause), Unwind(UnwindClause), + /// Read-only subquery, with an optional explicit importing WITH. + Call(Box), +} + +impl CypherQuery { + /// Every outer variable needed by this CALL's independently scoped UNION branches. + pub(crate) fn call_union_imports(&self) -> crate::error::Result> { + let mut names = self.call_imports()?; + for branch in &self.union_branches { + names.extend(branch.query.call_union_imports()?); + } + names.sort(); + names.dedup(); + Ok(names) + } + + /// Names explicitly imported by a leading plain WITH in a CALL body. + pub(crate) fn call_imports(&self) -> crate::error::Result> { + let Some(clause) = self + .with_clause + .as_ref() + .filter(|_| self.reading_clauses.is_empty()) + else { + return Ok(vec![]); + }; + if clause.distinct + || clause.order_by.is_some() + || clause.limit.is_some() + || clause.where_clause.is_some() + { + return Err(crate::error::GraphError::PlanError { + message: "CALL importing WITH must contain only unaliased variable names; use a second WITH for expressions or filters".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + clause.items.iter().map(|item| match (&item.expression, &item.alias) { + (ValueExpression::Variable(name), None) => Ok(name.to_lowercase()), + _ => Err(crate::error::GraphError::PlanError { + message: "CALL importing WITH must contain only unaliased variable names; use a second WITH for expressions".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + }).collect() + } } /// A MATCH clause containing graph patterns diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index f0ec4d7..0e8a6b3 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -1614,6 +1614,13 @@ impl ExtensionPlanner for CsrExtensionPlanner { physical_inputs: &[Arc], session_state: &SessionState, ) -> datafusion::common::Result>> { + if let Some(node) = node.as_any().downcast_ref::() { + return Ok(Some(crate::apply::plan_apply( + node, + physical_inputs, + session_state, + ))); + } if let Some(node) = node .as_any() .downcast_ref::() @@ -1651,7 +1658,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { } /// Query planner that adds [`CsrExtensionPlanner`]; install it on any -/// `SessionContext` that executes plans built with an [`AdjacencyCatalog`]. +/// `SessionContext` that executes plans built with an [`AdjacencyCatalog`] or CALL subqueries. #[derive(Debug, Default)] pub struct CsrQueryPlanner; diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index b4371b9..477231b 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -40,6 +40,7 @@ pub struct PlanningContext<'a> { pub(crate) path_sources: crate::path_inventory::SourceBatches, pub(crate) relationship_instance_idx: HashMap, /// Element rows of the list comprehensions being built, innermost last. + pub(crate) call_input: Option, pub(crate) comprehension_elements: Vec, } @@ -50,6 +51,7 @@ impl<'a> PlanningContext<'a> { path_sources: Default::default(), relationship_instance_idx: HashMap::new(), comprehension_elements: Vec::new(), + call_input: None, } } @@ -268,6 +270,21 @@ fn analyze_operator( } analyze_operator(element, analysis, rel_counter)?; } + LogicalOperator::Apply { + input, + inner, + labels, + .. + } => { + analyze_operator(input, analysis, rel_counter)?; + analysis + .required_datasets + .extend(analyze(inner)?.required_datasets); + analysis.var_to_label.extend(labels.clone()); + } + LogicalOperator::CallInput { labels } => { + analysis.var_to_label.extend(labels.clone()); + } LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } Ok(()) 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 07373b6..cd4b6f5 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -56,6 +56,7 @@ impl DataFusionPlanner { Self::collect_labels(input, &mut analysis.var_to_label); let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); + branch.call_input = ctx.call_input.clone(); branch.path_sources = ctx.path_sources.clone(); branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); @@ -123,6 +124,7 @@ impl DataFusionPlanner { Self::collect_labels(input, &mut analysis.var_to_label); let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); + branch.call_input = ctx.call_input.clone(); branch.path_sources = ctx.path_sources.clone(); branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); @@ -157,7 +159,7 @@ impl DataFusionPlanner { }); } } - // Consumers need the final output schema before collecting UNION results. + // A CALL must know its child's final schema before physical planning per row. // 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![ diff --git a/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs new file mode 100644 index 0000000..d93e67d --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +use super::super::{analysis, DataFusionPlanner, PlanningContext}; +use crate::apply::ApplyNode; +use crate::error::Result; +use crate::logical_plan::LogicalOperator; +use datafusion::logical_expr::{ + col, expr::Placeholder, Expr, Extension, LogicalPlan, LogicalPlanBuilder, +}; +use std::sync::Arc; + +impl DataFusionPlanner { + pub(crate) fn build_apply( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + inner: &LogicalOperator, + names: &[String], + id: usize, + ) -> Result { + let input = self.build_operator(ctx, input)?; + // A scalar named x wins over entity-shaped names such as x__property. + let entity_prefixes: Vec<_> = names + .iter() + .filter(|name| { + !input + .schema() + .fields() + .iter() + .any(|field| field.name() == *name) + }) + .map(|name| format!("{name}__")) + .collect(); + let mut imports = Vec::new(); + let mut values = Vec::new(); + for (index, (_, field)) in input.schema().iter().enumerate() { + if names.contains(field.name()) + || entity_prefixes + .iter() + .any(|prefix| field.name().starts_with(prefix)) + { + // User parameters have already become literals. The monotonically assigned + // CALL id also separates nested and sibling import placeholders. + let parameter = format!("__call_{id}_{index}"); + imports.push((field.name().clone(), parameter.clone())); + values.push(Expr::Placeholder(Placeholder::new_with_field( + format!("${parameter}"), + Some(field.clone()), + ))); + } + } + // DataFusion reconstructs VALUES with column1, column2, ... when binding + // placeholders. Keep those names and restore the imported names in a projection. + let seed = + if imports.is_empty() { + LogicalPlanBuilder::empty(true).build() + } else { + LogicalPlanBuilder::values(vec![values]).and_then(|builder| { + builder + .project(imports.iter().enumerate().map(|(index, (name, _))| { + col(format!("column{}", index + 1)).alias(name) + }))? + .build() + }) + } + .map_err(|error| self.plan_error("CALL imports", error))?; + let analysis = analysis::analyze(inner)?; + let mut child_ctx = PlanningContext::new(&analysis); + child_ctx.call_input = Some(seed); + child_ctx.path_sources = ctx.path_sources.clone(); + let child = self.build_operator(&mut child_ctx, inner)?; + let node = ApplyNode::try_new(input, child, imports) + .map_err(|error| self.plan_error("CALL output", error))?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } +} diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 869d3bb..d2d3e38 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -122,6 +122,11 @@ impl DataFusionPlanner { LogicalOperator::ListComprehension { input, .. } => { Self::collect_variables(input, vars) } + LogicalOperator::Apply { input, labels, .. } => { + Self::collect_variables(input, vars); + vars.extend(labels.keys().cloned()); + } + LogicalOperator::CallInput { labels } => vars.extend(labels.keys().cloned()), LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } @@ -195,6 +200,15 @@ impl DataFusionPlanner { LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), + LogicalOperator::Apply { + input, + labels: exports, + .. + } => { + Self::collect_labels(input, labels); + labels.extend(exports.clone()); + } + LogicalOperator::CallInput { labels: imports } => labels.extend(imports.clone()), LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index acad424..6ba1518 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -14,6 +14,7 @@ mod aggregate_ops; mod basic_ops; +mod call_ops; mod comprehension_ops; mod expand_ops; mod helpers; @@ -35,6 +36,16 @@ impl DataFusionPlanner { op: &LogicalOperator, ) -> Result { match op { + LogicalOperator::Apply { + input, + inner, + imports, + id, + .. + } => self.build_apply(ctx, input, inner, imports, *id), + LogicalOperator::CallInput { .. } => ctx.call_input.clone().ok_or_else(|| { + self.plan_error("CALL input", "import rows must be built inside CALL") + }), LogicalOperator::SingleRow => LogicalPlanBuilder::empty(true) .build() .map_err(|e| self.plan_error("Failed to build singleton row", e)), diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index e024968..4dbdeae 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -35,6 +35,7 @@ //! # } //! ``` +mod apply; pub mod ast; pub mod case_insensitive; pub mod config; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index b319dba..fd5343a 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -188,6 +188,17 @@ pub enum LogicalOperator { output: String, }, + /// Evaluate a read-only subquery once per input row, extending it by each result. + Apply { + input: Box, + inner: Box, + imports: Vec, + labels: HashMap, + id: usize, + }, + /// The single imported row of a CALL body; its values are bound at execution. + CallInput { labels: HashMap }, + /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. ComprehensionElement, } @@ -233,6 +244,7 @@ pub struct LogicalPlanner<'a> { node_lists: HashMap, /// List comprehensions planned so far (output columns `__list_`). comprehensions: usize, + calls: usize, /// The node label of each path variable's nodes, for `nodes(path)`. path_labels: HashMap, rich_paths: bool, @@ -252,6 +264,7 @@ impl<'a> LogicalPlanner<'a> { exists_marks: 0, node_lists: HashMap::new(), comprehensions: 0, + calls: 0, path_labels: HashMap::new(), rich_paths: false, seed_source: None, @@ -409,9 +422,30 @@ impl<'a> LogicalPlanner<'a> { /// Convert a Cypher AST to a logical plan pub fn plan(&mut self, query: &CypherQuery) -> Result { self.rich_paths = crate::path_value::needs_values(query, self.config); + self.plan_query(query, None, false) + } + + fn plan_query( + &mut self, + query: &CypherQuery, + input: Option, + preserve_entities: bool, + ) -> Result { + let seed = input.clone(); + let imports = ( + self.variables.clone(), + self.node_lists.clone(), + self.path_labels.clone(), + ); + if preserve_entities { + let names = query.call_imports()?; + self.variables.retain(|name, _| names.contains(name)); + self.node_lists.retain(|name, _| names.contains(name)); + self.path_labels.retain(|name, _| names.contains(name)); + } // Plan main MATCH clauses; the WHERE clause guides where traversal starts. let where_expr = query.where_clause.as_ref().map(|w| &w.expression); - let mut plan = self.plan_reading_clauses(None, &query.reading_clauses, where_expr)?; + let mut plan = self.plan_reading_clauses(input, &query.reading_clauses, where_expr)?; // Apply WHERE clause if present (before WITH) if let Some(where_clause) = &query.where_clause { @@ -445,7 +479,20 @@ impl<'a> LogicalPlanner<'a> { } // Apply RETURN clause - plan = self.plan_return_clause(&query.return_clause, plan)?; + plan = if preserve_entities { + self.plan_with_clause( + &WithClause { + distinct: query.return_clause.distinct, + items: query.return_clause.items.clone(), + order_by: None, + limit: None, + where_clause: None, + }, + plan, + )? + } else { + self.plan_return_clause(&query.return_clause, plan)? + }; // Apply ORDER BY if present if let Some(order_by) = &query.order_by { @@ -475,7 +522,32 @@ impl<'a> LogicalPlanner<'a> { let mut planner = LogicalPlanner::new(self.config) .with_fan_out(self.fan_out) .with_adjacency(self.adjacency); - let right = planner.plan(&branch.query)?; + planner.calls = self.calls; + planner.rich_paths = self.rich_paths; + if preserve_entities { + planner.variables = imports.0.clone(); + planner.node_lists = imports.1.clone(); + planner.path_labels = imports.2.clone(); + } + let right = planner.plan_query(&branch.query, seed.clone(), preserve_entities)?; + self.calls = planner.calls; + if preserve_entities { + let node_bindings = |variables: &HashMap| { + variables + .iter() + .filter(|(_, label)| self.config.get_node_mapping(label).is_some()) + .map(|(name, label)| (name.clone(), label.clone())) + .collect::>() + }; + if node_bindings(&self.variables) != node_bindings(&planner.variables) + || self.node_lists != planner.node_lists + { + return Err(GraphError::PlanError { + message: "CALL UNION node and node-list outputs must retain the same labels in every branch".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } plan = LogicalOperator::QueryUnion { left: Box::new(plan), right: Box::new(right), @@ -515,6 +587,7 @@ impl<'a> LogicalPlanner<'a> { filter: Option<&BooleanExpression>, ) -> Result { match clause { + ReadingClause::Call(query) => self.plan_call(base, query), ReadingClause::Match(match_clause) if match_clause.optional => { self.plan_optional_match(base, match_clause, None) } @@ -531,6 +604,59 @@ impl<'a> LogicalPlanner<'a> { } } + fn plan_call( + &mut self, + input: Option, + query: &CypherQuery, + ) -> Result { + let imports = query.call_union_imports()?; + self.calls += 1; + let id = self.calls; + let mut inner = Self::new(self.config) + .with_fan_out(self.fan_out) + .with_adjacency(self.adjacency); + inner.calls = self.calls; + inner.rich_paths = self.rich_paths; + for name in &imports { + if let Some(label) = self.variables.get(name) { + inner.variables.insert(name.clone(), label.clone()); + } + if let Some(label) = self.node_lists.get(name) { + inner.node_lists.insert(name.clone(), label.clone()); + } + if let Some(label) = self.path_labels.get(name) { + inner.path_labels.insert(name.clone(), label.clone()); + } + } + let input_labels = inner.variables.clone(); + let plan = inner.plan_query( + query, + Some(LogicalOperator::CallInput { + labels: input_labels, + }), + true, + )?; + self.calls = inner.calls; + let labels = inner.variables.clone(); + self.variables.extend(inner.variables); + self.node_lists.extend(inner.node_lists); + for item in &query.return_clause.items { + if let ValueExpression::Variable(name) = &item.expression { + if let Some(label) = inner.path_labels.get(name) { + self.path_labels + .insert(item.alias.as_ref().unwrap_or(name).clone(), label.clone()); + } + } + } + Ok(LogicalOperator::Apply { + input: Box::new(input.unwrap_or(LogicalOperator::SingleRow)), + inner: Box::new(plan), + imports, + labels, + id, + }) + } + /// Plan an UNWIND clause fn plan_unwind_clause_with_base( &mut self, @@ -1598,10 +1724,13 @@ impl<'a> LogicalPlanner<'a> { #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { match plan { - LogicalOperator::SingleRow => Err(GraphError::PlanError { - message: "a singleton row binds no pattern variable".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }), + LogicalOperator::Apply { input, .. } => self.extract_variable_from_plan(input), + LogicalOperator::CallInput { .. } | LogicalOperator::SingleRow => { + Err(GraphError::PlanError { + message: "a singleton row binds no pattern variable".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } LogicalOperator::ScanByLabel { variable, .. } => Ok(variable.clone()), LogicalOperator::Unwind { alias, .. } => Ok(alias.clone()), LogicalOperator::Expand { diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index 785f18f..e4676d6 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -80,6 +80,7 @@ fn substitute_in_reading_clause( parameters: &HashMap, ) -> Result<()> { match clause { + ReadingClause::Call(query) => substitute_parameters(query, parameters)?, ReadingClause::Match(match_clause) => { for pattern in &mut match_clause.patterns { substitute_in_graph_pattern(pattern, parameters)?; diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 43a6cfc..716b351 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -52,8 +52,10 @@ fn normalize_inline_properties(query: &mut CypherQuery) { .chain(&mut query.post_with_reading_clauses) .chain(stages) { - if let ReadingClause::Match(m) = clause { - move_variable_properties_to_where(m, &mut anonymous); + match clause { + ReadingClause::Match(m) => move_variable_properties_to_where(m, &mut anonymous), + ReadingClause::Call(inner) => normalize_inline_properties(inner), + ReadingClause::Unwind(_) => {} } } for branch in &mut query.union_branches { @@ -179,9 +181,17 @@ fn reading_clause(input: &str) -> IResult<&str, ReadingClause> { alt(( map(match_clause, ReadingClause::Match), map(unwind_clause, ReadingClause::Unwind), + map(call_clause, |query| ReadingClause::Call(Box::new(query))), ))(input) } +fn call_clause(input: &str) -> IResult<&str, CypherQuery> { + let (input, _) = multispace0(input)?; + let (input, _) = tag_no_case("CALL")(input)?; + let (input, _) = multispace0(input)?; + delimited(char('{'), cypher_query, preceded(multispace0, char('}')))(input) +} + // Parse a MATCH clause fn match_clause(input: &str) -> IResult<&str, MatchClause> { let (input, _) = multispace0(input)?; @@ -2614,7 +2624,7 @@ mod tests { .iter() .map(|c| match c { ReadingClause::Match(m) => m, - ReadingClause::Unwind(_) => unreachable!(), + ReadingClause::Unwind(_) | ReadingClause::Call(_) => unreachable!(), }) .collect(); assert_eq!(clauses.len(), 3); diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index af3ff5c..a1b35be 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -202,7 +202,18 @@ pub(crate) fn needs_values( let mut needed = query .union_branches .iter() - .any(|branch| needs_values(&branch.query, config)); + .any(|branch| needs_values(&branch.query, config)) + || clauses.iter().any(|clause| match clause { + ReadingClause::Call(inner) => { + needs_values(inner, config) + || inner + .call_union_imports() + .unwrap_or_default() + .iter() + .any(|name| names.contains(name)) + } + _ => false, + }); let mut inspect = |expression: &V| { let mut value = expression.clone(); crate::ast::transform_value(&mut value, &mut |v| match v { @@ -236,6 +247,7 @@ pub(crate) fn needs_values( }; for clause in clauses { match clause { + ReadingClause::Call(_) => {} ReadingClause::Unwind(u) => inspect(&u.expression), ReadingClause::Match(m) => { if let Some(w) = &m.where_clause { @@ -296,7 +308,9 @@ mod tests { fn direct_shortest_queries_keep_the_csr_route() { let config = crate::config::GraphConfig::builder() .node_table("Person", "id") + .node_table("Place", "id") .rel_table("R", "Person", "Person", "src", "dst") + .rel_table("LIVES", "Person", "Place", "src", "dst") .build() .unwrap(); let prefix = "MATCH p = shortestPath((a:Person)-[:R*]-(b:Person)) "; @@ -310,6 +324,14 @@ mod tests { ), ("WITH p AS q RETURN q", true), ("WITH p AS q RETURN relationships(q)", true), + ( + "WITH p AS q CALL { WITH q RETURN relationships(q) AS edges } RETURN edges", + true, + ), + ( + "WITH p AS q CALL { WITH q RETURN q AS result UNION WITH q RETURN q AS result } RETURN nodes(result)", + true, + ), ( "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r)", true, @@ -339,5 +361,13 @@ mod tests { .unwrap(); assert!(needs_values(&query, &config)); } + for text in [ + "CALL { MATCH p = shortestPath((a:Person)-[:LIVES*0..]->(b:Place)) RETURN length(p) AS n } RETURN n", + "RETURN 0 AS n UNION MATCH p = shortestPath((a:Person)-[:LIVES*0..]->(b:Place)) RETURN length(p) AS n", + "CALL { RETURN 0 AS n UNION MATCH p = shortestPath((a:Person)-[:LIVES*0..]-(b:Person)) RETURN length(p) AS n } RETURN n", + ] { + let query = crate::parser::parse_cypher_query(text).unwrap(); + assert!(needs_values(&query, &config), "{text}"); + } } } diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 905033c..632b7ae 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -255,13 +255,16 @@ impl CypherQuery { .iter() .filter_map(|stage| stage.with_clause.limit.as_ref()), ) - .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); + .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))) + || query + .all_reading_clauses() + .any(|clause| matches!(clause, ReadingClause::Call(_))); }); if shape_sensitive || self.scalar_parameters().is_none() || crate::path_value::needs_values(&self.ast, self.require_config()?) { - // Parameterized row counts alter plan shape. + // CALL binds its own child templates; row counts alter plan shape. return Ok(false); } let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; @@ -1270,6 +1273,7 @@ impl CypherQuery { } } }); + labels.sort(); labels.dedup(); labels @@ -1288,6 +1292,7 @@ impl CypherQuery { } } }); + types.sort(); types.dedup(); types @@ -1300,6 +1305,15 @@ impl CypherQuery { self.ast.for_each_query(&mut |ast| { for clause in ast.all_reading_clauses() { match clause { + ReadingClause::Call(inner) => { + variables.extend(inner.call_imports().unwrap_or_default()); + variables.extend(inner.return_clause.items.iter().filter_map(|item| { + item.alias.clone().or_else(|| match &item.expression { + crate::ast::ValueExpression::Variable(name) => Some(name.clone()), + _ => None, + }) + })); + } ReadingClause::Match(match_clause) => { for pattern in &match_clause.patterns { self.collect_variables_from_pattern(pattern, &mut variables); @@ -1311,6 +1325,7 @@ impl CypherQuery { } } }); + variables.sort(); variables.dedup(); variables diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 372deae..9a8533c 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -27,6 +27,7 @@ pub struct SemanticAnalyzer { current_scope: ScopeType, catalog: Option>, expression_types: HashMap, + call_body: bool, } /// Information about a variable in the query @@ -78,6 +79,7 @@ impl SemanticAnalyzer { current_scope: ScopeType::Match, catalog: None, expression_types: HashMap::new(), + call_body: false, } } @@ -96,6 +98,17 @@ impl SemanticAnalyzer { query: &CypherQuery, parameters: &HashMap, ) -> Result { + // UNION siblings receive the same available CALL imports, never variables + // created by the preceding branch. Each branch opts in with its own WITH. + let (imports, import_types) = if self.call_body { + let available = (self.variables.clone(), self.expression_types.clone()); + let names = query.call_imports()?; + self.variables.retain(|name, _| names.contains(name)); + self.expression_types.retain(|name, _| names.contains(name)); + available + } else { + (HashMap::new(), HashMap::new()) + }; // Clone the query to perform parameter substitution let mut analyzed_query = query.clone(); @@ -109,6 +122,11 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &analyzed_query.reading_clauses { match clause { + ReadingClause::Call(query) => { + if let Err(error) = self.analyze_call(query) { + errors.push(format!("CALL clause error: {error}")); + } + } ReadingClause::Match(match_clause) => { if let Err(e) = self.analyze_match_clause(match_clause) { errors.push(format!("MATCH clause error: {}", e)); @@ -142,6 +160,11 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &analyzed_query.post_with_reading_clauses { match clause { + ReadingClause::Call(query) => { + if let Err(error) = self.analyze_call(query) { + errors.push(format!("CALL clause error: {error}")); + } + } ReadingClause::Match(match_clause) => { if let Err(e) = self.analyze_match_clause(match_clause) { errors.push(format!("Post-WITH MATCH clause error: {}", e)); @@ -172,6 +195,7 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &stage.reading_clauses { let result = match clause { + ReadingClause::Call(query) => self.analyze_call(query), ReadingClause::Match(match_clause) => self.analyze_match_clause(match_clause), ReadingClause::Unwind(unwind_clause) => { self.analyze_unwind_clause(unwind_clause) @@ -203,6 +227,14 @@ impl SemanticAnalyzer { } } + if self.call_body + && query.return_clause.items.iter().any(|item| { + item.alias.is_none() && !matches!(item.expression, ValueExpression::Variable(_)) + }) + { + errors.push("CALL RETURN expressions need an alias".into()); + } + // Phase 7: Schema validation self.validate_schema(&mut warnings); @@ -212,6 +244,9 @@ impl SemanticAnalyzer { for branch in &mut analyzed_query.union_branches { let mut analyzer = SemanticAnalyzer::new(self.config.clone()); analyzer.catalog = self.catalog.clone(); + analyzer.call_body = self.call_body; + analyzer.variables = imports.clone(); + analyzer.expression_types = import_types.clone(); let analyzed = analyzer.analyze(&branch.query, parameters)?; branch.query = analyzed.ast; errors.extend( @@ -231,6 +266,61 @@ impl SemanticAnalyzer { }) } + fn analyze_call(&mut self, query: &CypherQuery) -> Result<()> { + let error = |message| GraphError::PlanError { + message, + location: snafu::Location::new(file!(), line!(), column!()), + }; + let mut inner = Self::new(self.config.clone()); + inner.catalog = self.catalog.clone(); + inner.call_body = true; + for name in query.call_union_imports()? { + let variable = self + .variables + .get(&name) + .ok_or_else(|| error(format!("CALL imports undefined variable '{name}'")))?; + inner.variables.insert(name.clone(), variable.clone()); + if let Some(value_type) = self.expression_types.get(&name) { + inner.expression_types.insert(name, value_type.clone()); + } + } + // Outer analysis already substituted all user parameters recursively. + let analyzed = inner.analyze(query, &HashMap::new())?; + if !analyzed.errors.is_empty() { + return Err(error(analyzed.errors.join("\n"))); + } + let mut exports = HashSet::new(); + for item in &query.return_clause.items { + let name = match (&item.alias, &item.expression) { + (Some(alias), _) => alias.to_lowercase(), + (None, ValueExpression::Variable(name)) => name.to_lowercase(), + _ => return Err(error("CALL RETURN expressions need an alias".into())), + }; + if self.variables.contains_key(&name) || !exports.insert(name.clone()) { + return Err(error(format!( + "CALL export '{name}' collides with an existing variable" + ))); + } + let value_type = inner.expression_type(&item.expression); + let (variable_type, labels) = match &value_type { + ExpressionType::Entity(kind, labels) => (kind.clone(), labels.clone()), + _ => (VariableType::Property, vec![]), + }; + self.variables.insert( + name.clone(), + VariableInfo { + name: name.clone(), + variable_type, + labels, + properties: HashSet::new(), + defined_in: self.current_scope.clone(), + }, + ); + self.expression_types.insert(name, value_type); + } + Ok(()) + } + /// Analyze MATCH clause and discover variables fn analyze_match_clause(&mut self, match_clause: &MatchClause) -> Result<()> { for pattern in &match_clause.patterns { diff --git a/docs/api-reference.md b/docs/api-reference.md index a20b93f..c5cc2ef 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -129,7 +129,7 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are -bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT` and unsupported template +bound before optimization on every call. Containers, rich path values, `CALL` subqueries, parameterized `SKIP`/`LIMIT` and unsupported template shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. See [Repeated queries](lance.md#repeated-queries) for the scope and limits. diff --git a/docs/cypher.md b/docs/cypher.md index 0d97ead..55ef59e 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -210,12 +210,13 @@ used by IC14. The tables below specify its supported surface. | `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. | +| `CALL { ... RETURN ... }` | Read-only subquery; import outer variables with a leading plain `WITH`. | | `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`), `CALL`, and `FOREACH`. + Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), procedure calls (`CALL name(...)` / `YIELD`), 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 @@ -232,6 +233,50 @@ column counts, names or types raise `ValueError`. This follows the existing homo engine.execute("RETURN 1 AS n UNION RETURN 1 AS n").to_pylist() # [{'n': 1}] ``` +### Read-only subqueries + +Use `CALL { ... }` when a match, aggregate or limit should run separately for each incoming row: + +```python +print(engine.execute(""" + MATCH (p:Person) + CALL { + WITH p + MATCH (p)-[:KNOWS]->(f:Person) + RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 + } + RETURN p.name AS person, friend.name AS friend ORDER BY person +""").to_pylist()) +# [{'person': 'Alice', 'friend': 'Bob'}, {'person': 'Bob', 'friend': 'David'}, +# {'person': 'Carol', 'friend': 'David'}] +``` + +The leading importing `WITH` accepts only existing variable names, without aliases, expressions, +`DISTINCT`, filters or pagination. Use a second `WITH` to transform or filter them. A body without an +importing `WITH` has its own independent scope. Calls can nest and can appear after `WITH` or `UNWIND`. + +A CALL body can combine branches with `UNION` or `UNION ALL`. Each branch needs its own importing +`WITH` and may import a different subset of the outer variables. Branches cannot see variables created +by sibling branches. Deduplication happens within each incoming row, so repeated outer rows remain +repeated. Node and node-list exports must retain the same labels in every branch; mixed-label entity +outputs are not represented by the current binding metadata. + +Every body must end with `RETURN`. Give expressions an `AS` alias; bare variables keep their names. +Returned names must be distinct and must not already exist outside the call. Returning a node or +relationship retains its properties, so later clauses can use the returned entity. Names that the body +only uses internally do not become visible outside it. + +Each child result row is combined with its incoming row. Duplicate results remain duplicates. An empty +child removes that incoming row; an aggregate such as `count()` without grouping still returns a row +with zero. An empty outer input never executes the child. `ORDER BY` and `LIMIT` inside the body apply +separately to each incoming row. + +CALL currently binds, optimizes and physically plans a child for every incoming row. It streams the +child's batches before advancing that row, but operators such as sorting and aggregation can still +buffer their inputs. CALL queries bypass the query-template cache. For large outer inputs, measure the +query on representative data before choosing this form. Procedures, `YIELD`, and writes are unsupported. + + ### Patterns ```cypher diff --git a/docs/execution.md b/docs/execution.md index f96bd1a..3c4fa0f 100644 --- a/docs/execution.md +++ b/docs/execution.md @@ -99,3 +99,7 @@ results stay bounded, while joins, sorts, aggregations and the CSR itself hold d sets no memory limit or spill-to-disk policy, and never distributes work across machines: distributed execution is [out of scope](index.md#what-it-is-not) for the project. + +## Read-only subqueries + +`CALL { ... RETURN ... }` evaluates its child query once per incoming row. A leading plain `WITH` imports outer variables; each body and UNION branch has its own scope. The operator streams each child result before advancing to the next outer row and yields after bounded polling work, including empty batches. Dropping or expiring the reader releases its owned child stream. diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py new file mode 100644 index 0000000..87668e5 --- /dev/null +++ b/python/python/tests/test_call_subqueries.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Read-only CALL uses explicit imports and evaluates its body per outer row.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + + +@pytest.fixture(params=[False, True]) +def graph(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + data = { + "Person": pa.table( + { + "id": [10, 20, 30, 40], + "name": ["A", "B", "C", "D"], + "age": [20, 30, 40, 50], + } + ), + "KNOWS": pa.table({"src": [10, 10, 10, 20], "dst": [20, 20, 30, 40]}), + } + return CypherEngine(config, data, use_csr=request.param), config, data + + +def test_call_top_k_preserves_node_alias_and_properties(graph): + engine, config, data = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 } " + "RETURN p.id AS id, friend.name AS name ORDER BY id" + ) + expected = [{"id": 10, "name": "C"}, {"id": 20, "name": "D"}] + assert engine.execute(query).to_pylist() == expected + assert engine.execute_batches(query).read_all().to_pylist() == expected + assert CypherQuery(query).with_config(config).execute(data).to_pylist() == expected + assert "Apply" in engine.explain(query) + + +def test_call_aggregates_and_duplicates(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate CALL { WITH p " + "MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n } " + "RETURN p.id AS id, n ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [ + {"id": person, "n": count} + for person, count in [(10, 3), (20, 1), (30, 0), (40, 0)] + for _ in range(2) + ] + query = ( + "MATCH (p:Person) WHERE p.id = 10 CALL { WITH p " + "MATCH (p)-[:KNOWS]->(f:Person) RETURN f.id AS id } RETURN id ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [{"id": 20}, {"id": 20}, {"id": 30}] + + +def test_call_nested_parameters_are_fresh_and_do_not_collide(graph): + engine, _, _ = graph + query = ( + "UNWIND [1, 2] AS x CALL { WITH x CALL { WITH x " + "RETURN x + $__call_1_0 AS y } RETURN y + 1 AS z } " + "CALL { WITH z RETURN z * 2 AS answer } RETURN x, answer ORDER BY x" + ) + for value in [10, 20, 10]: + assert engine.execute(query, {"__call_1_0": value}).to_pylist() == [ + {"x": x, "answer": (x + value + 1) * 2} for x in [1, 2] + ] + + +def test_call_nullable_scalar_import_and_second_with(graph): + engine, _, _ = graph + query = ( + "UNWIND [null, 2] AS x CALL { WITH x WITH coalesce(x, 7) AS value " + "RETURN value + 1 AS answer } RETURN answer ORDER BY answer" + ) + assert engine.execute(query).to_pylist() == [{"answer": 3}, {"answer": 8}] + + +def test_call_export_can_start_outer_traversal(graph): + engine, _, _ = graph + query = ( + "CALL { MATCH (p:Person) WHERE p.id = 10 RETURN p AS person } " + "MATCH (person)-[:KNOWS]->(f:Person) RETURN f.name AS name ORDER BY name" + ) + assert engine.execute(query).to_pylist() == [ + {"name": "B"}, + {"name": "B"}, + {"name": "C"}, + ] + + +def test_call_empty_outer_and_early_reader_close(graph): + engine, _, _ = graph + assert ( + engine.execute( + "MATCH (p:Person) WHERE p.id < 0 " + "CALL { RETURN 1 / 0 AS invalid } RETURN invalid" + ).num_rows + == 0 + ) + reader = engine.execute_batches( + "MATCH (p:Person) CALL { WITH p UNWIND range(1, 10000) AS n RETURN n } RETURN n" + ) + assert reader.read_next_batch().num_rows > 0 + reader.close() + assert engine.execute("RETURN 1 AS ok").to_pylist() == [{"ok": 1}] + + +@pytest.mark.parametrize( + ("query", "message"), + [ + ("UNWIND [1] AS x CALL { RETURN x AS y } RETURN y", "Undefined variable"), + ("UNWIND [1] AS x CALL { RETURN 2 AS x } RETURN x", "collides"), + ("CALL { RETURN 1 + 2 } RETURN 1", "need an alias"), + ("CALL { WITH missing RETURN missing AS x } RETURN x", "undefined variable"), + ( + "UNWIND [1] AS x CALL { WITH x + 1 AS y RETURN y } RETURN y", + "importing WITH", + ), + ( + "CALL { UNWIND [1] AS local RETURN 2 AS y } RETURN local", + "Undefined variable", + ), + ], +) +def test_call_scope_errors(graph, query, message): + engine, _, _ = graph + with pytest.raises(ValueError, match=message): + engine.execute(query) + + +def test_call_optional_entity_import_and_relationship_export(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "CALL { WITH f RETURN f.name AS name } " + "RETURN p.id AS id, name ORDER BY id, name" + ) + assert engine.execute(query).to_pylist() == [ + {"id": 10, "name": "B"}, + {"id": 10, "name": "B"}, + {"id": 10, "name": "C"}, + {"id": 20, "name": "D"}, + {"id": 30, "name": None}, + {"id": 40, "name": None}, + ] + query = ( + "CALL { MATCH (p:Person)-[r:KNOWS]->(f:Person) RETURN r AS edge } " + "CALL { WITH edge RETURN edge.dst AS destination } " + "RETURN destination ORDER BY destination" + ) + assert engine.execute(query).to_pylist() == [ + {"destination": 20}, + {"destination": 20}, + {"destination": 30}, + {"destination": 40}, + ] + + +def test_call_preserves_container_imports(graph): + engine, _, _ = graph + assert engine.execute( + "UNWIND [{values: [1, 2]}, {values: [3]}] AS item " + "CALL { WITH item UNWIND item.values AS n RETURN sum(n) AS total } " + "RETURN total ORDER BY total" + ).to_pylist() == [{"total": 3}, {"total": 3}] + + +def test_call_scalar_import_does_not_import_entity_shaped_names(graph): + engine, _, _ = graph + assert engine.execute( + "WITH 1 AS x, 2 AS x__other CALL { WITH x RETURN x AS y } RETURN y" + ).to_pylist() == [{"y": 1}] + + +def test_call_parameterized_child_limit_is_per_row_and_fresh(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f.id AS friend ORDER BY friend LIMIT $take } " + "RETURN p.id AS id, friend ORDER BY id, friend" + ) + for take, expected in [ + (1, [{"id": 10, "friend": 20}, {"id": 20, "friend": 40}]), + (0, []), + ( + 2, + [ + {"id": 10, "friend": 20}, + {"id": 10, "friend": 20}, + {"id": 20, "friend": 40}, + ], + ), + ]: + assert engine.execute(query, {"take": take}).to_pylist() == expected + + +@pytest.mark.parametrize(("union", "repeats"), [("UNION", 1), ("UNION ALL", 2)]) +def test_call_union_distinct_is_per_outer_row(graph, union, repeats): + engine, _, _ = graph + query = ( + "UNWIND [1, 1, 2] AS x CALL { WITH x RETURN x + $value AS n " + f"{union} WITH x RETURN x + $value AS n }} RETURN x, n ORDER BY x" + ) + for value in [10, 20, 10]: + expected = [{"x": x, "n": x + value} for x in [1, 1, 2] for _ in range(repeats)] + assert engine.execute(query, {"value": value}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"value": value}).read_all().to_pylist() + == expected + ) + + +def test_call_union_branches_import_different_variables_without_leaking(graph): + engine, _, _ = graph + assert engine.execute( + "WITH 2 AS x, 3 AS y CALL { WITH x RETURN x AS n " + "UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 2}, {"n": 3}] + assert engine.execute( + "WITH 3 AS y CALL { RETURN 2 AS n " + "UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 2}, {"n": 3}] + for query in [ + "WITH 2 AS x CALL { WITH x RETURN x AS n UNION RETURN x AS n } RETURN n", + "WITH 2 AS x CALL { RETURN x AS n UNION WITH x RETURN x AS n } RETURN n", + "WITH 2 AS x CALL { WITH x RETURN x AS n UNION WITH x RETURN x + 1 } RETURN n", + ]: + with pytest.raises(ValueError): + engine.execute(query) + + +@pytest.mark.parametrize("union", ["UNION", "UNION ALL"]) +def test_call_union_nodes_and_per_branch_top_k(graph, union): + engine, _, _ = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 " + f"{union} WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age ASC LIMIT 1 } " + "RETURN p.id AS id, friend.name AS name ORDER BY id, name" + ) + expected = [ + {"id": 10, "name": "B"}, + {"id": 10, "name": "C"}, + {"id": 20, "name": "D"}, + ] + if union == "UNION ALL": + expected.append({"id": 20, "name": "D"}) + assert engine.execute(query).to_pylist() == expected + + +def test_call_union_aggregates_and_empty_branches(graph): + engine, _, _ = graph + assert engine.execute( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN count(f) AS n UNION ALL WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "WHERE f.id < 0 RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id, n" + ).to_pylist() == [ + {"id": person, "n": n} + for person, count in [(10, 3), (20, 1), (30, 0), (40, 0)] + for n in [0, count] + ] + assert engine.execute( + "UNWIND [1, 2] AS x CALL { WITH x UNWIND [] AS empty RETURN x AS n " + "UNION ALL WITH x RETURN x AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 1}, {"n": 2}] + + +def test_call_in_later_union_branch_uses_fresh_parameters(graph): + engine, _, _ = graph + query = "RETURN 0 AS n UNION ALL CALL { RETURN $value AS n } RETURN n" + for value in [1, 2, 1]: + assert sorted( + row["n"] for row in engine.execute(query, {"value": value}).to_pylist() + ) == [0, value] + + +def test_call_union_schema_validation_uses_each_branch_catalog(graph): + engine, _, _ = graph + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute( + "MATCH (p:Person) CALL { WITH p RETURN p.age AS n " + "UNION ALL WITH p RETURN p.name + 1 AS n } RETURN n" + ) + + +def test_call_union_rejects_mixed_node_label_exports(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("City", "id") + .build() + ) + engine = CypherEngine( + config, {"Person": pa.table({"id": [1]}), "City": pa.table({"id": [1]})} + ) + with pytest.raises(ValueError, match="same labels"): + engine.execute( + "CALL { MATCH (p:Person) RETURN p AS entity " + "UNION MATCH (c:City) RETURN c AS entity } RETURN entity.id" + ) + + +def test_call_union_coerces_numeric_and_null_output_schema(graph): + engine, _, _ = graph + query = ( + "UNWIND [1, 2] AS x CALL { WITH x RETURN x AS n UNION ALL " + "RETURN $value AS n } RETURN n ORDER BY n" + ) + for value, expected in [(2.5, [1.0, 2.0, 2.5, 2.5]), (None, [None, None, 1, 2])]: + table = engine.execute(query, {"value": value}) + expected_type = pa.float64() if value is not None else pa.int64() + assert table.schema.field("n").type == expected_type + assert table.column(0).to_pylist() == expected + reader = engine.execute_batches(query, {"value": value}) + assert reader.schema.field("n").type == expected_type + assert reader.read_all().column(0).to_pylist() == expected + + +def test_call_imported_shortest_path_keeps_relationship_records(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) CALL { WITH path " + "RETURN [edge IN relationships(path) | endNode(edge).id] AS targets, " + "reduce(total = 0, edge IN relationships(path) | " + "total + endNode(edge).id) AS total } " + "RETURN targets, total ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40], "total": 60}, + {"targets": [20, 40], "total": 60}, + ] + + +def test_call_union_shares_path_occurrence_identity(graph): + engine, _, _ = graph + query = ( + "CALL { MATCH path = (a:Person {id: 10})-[:KNOWS]->(b:Person) " + "RETURN path UNION MATCH path = (a:Person {id: 10})-[:KNOWS]->(b:Person) " + "RETURN path } WITH collect(path) AS paths UNWIND paths AS path " + "RETURN length(path) AS hops, [n IN nodes(path) | n.name] AS names, " + "[r IN relationships(path) | endNode(r).id] AS targets ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"hops": 1, "names": ["A", "B"], "targets": [20]}, + {"hops": 1, "names": ["A", "B"], "targets": [20]}, + {"hops": 1, "names": ["A", "C"], "targets": [30]}, + ] + + +def test_call_keeps_path_alias_across_nested_union(graph): + engine, _, _ = graph + query = ( + "MATCH path = (a:Person {id: 20})-[:KNOWS]->(b:Person) " + "CALL { WITH path CALL { WITH path RETURN path AS value " + "UNION ALL WITH path RETURN path AS value LIMIT 0 } RETURN value AS result } " + "RETURN length(result) AS hops, [n IN nodes(result) | labels(n)] AS labels" + ) + assert engine.execute(query).to_pylist() == [ + {"hops": 1, "labels": [["Person"], ["Person"]]} + ] + + +def test_call_empty_children_respect_query_deadline(graph): + engine, _, _ = graph + with pytest.raises(TimeoutError): + engine.execute( + "UNWIND range(1, 1000000) AS x " + "CALL { WITH x UNWIND [] AS n RETURN n } RETURN n", + timeout=0.02, + ) + assert engine.execute("RETURN 1 AS ok").to_pylist() == [{"ok": 1}] + + +def test_rich_path_parameters_bind_freshly(graph): + engine, _, _ = graph + query = ( + "MATCH path = (a:Person {id: $start})-[:KNOWS]->(b:Person) " + "RETURN nodes(path) AS ids ORDER BY ids" + ) + for start, expected in [ + (10, [[10, 20], [10, 20], [10, 30]]), + (20, [[20, 40]]), + (10, [[10, 20], [10, 20], [10, 30]]), + ]: + assert engine.execute(query, {"start": start}).column(0).to_pylist() == expected + + +def test_rich_path_reader_retains_snapshot_after_refresh(graph, tmp_path): + import lance + + original, config, data = graph + paths = {name: str(tmp_path / f"{name}.lance") for name in data} + datasets = { + name: lance.write_dataset(table, paths[name]) for name, table in data.items() + } + engine = CypherEngine(config, datasets, use_csr=bool(original.csr_stats())) + query = ( + "MATCH path = (a:Person {id: $start})-[:KNOWS]->(b:Person) " + "RETURN nodes(path) AS ids ORDER BY ids" + ) + old = engine.execute_batches(query, {"start": 20}) + lance.write_dataset( + pa.table({"src": [20], "dst": [30]}), paths["KNOWS"], mode="append" + ) + engine.refresh() + assert engine.execute(query, {"start": 20}).column(0).to_pylist() == [ + [20, 30], + [20, 40], + ] + assert old.read_all().column(0).to_pylist() == [[20, 40]] + assert engine.execute(query, {"start": 10}).column(0).to_pylist() == [ + [10, 20], + [10, 20], + [10, 30], + ] + + +def test_lambda_types_shadow_and_restore_outer_bindings(graph): + engine, _, _ = graph + for initial, expected in [(0, 3), (0.5, 3.5), (None, None), (0, 3)]: + rows = engine.execute( + "WITH 'outer' AS total RETURN total AS outer, " + "reduce(total = $initial, item IN [1, 2] | total + item) AS value", + {"initial": initial}, + ).to_pylist() + assert rows == [{"outer": "outer", "value": expected}] + assert engine.execute( + "WITH 'outer' AS item RETURN all(item IN [1, 2] WHERE item > 0) AS valid, " + "item AS outer" + ).to_pylist() == [{"valid": True, "outer": "outer"}] + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute( + "RETURN reduce(total = 0, item IN ['text'] | total + item) AS value" + ) + + +def test_shortest_path_alias_import_demands_relationship_values(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) WITH path AS carried " + "CALL { WITH carried RETURN relationships(carried) AS edges } " + "RETURN [edge IN edges | endNode(edge).id] AS targets ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40]}, + {"targets": [20, 40]}, + ] + + +def test_shortest_path_alias_survives_call_union(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) WITH path AS carried " + "CALL { WITH carried RETURN carried AS result " + "UNION WITH carried RETURN carried AS result } " + "RETURN [edge IN relationships(result) | endNode(edge).id] AS targets " + "ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40]}, + {"targets": [20, 40]}, + ] diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index d7c7779..512e46f 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -556,7 +556,7 @@ def cross_label_graph(request): return config, tables, request.param -@pytest.mark.parametrize("scope", ["direct", "alias", "union"]) +@pytest.mark.parametrize("scope", ["direct", "alias", "call", "union", "call_union"]) def test_cross_label_shortest_scopes_rebind_each_parameter(cross_label_graph, scope): config, tables, use_csr = cross_label_graph engine = CypherEngine(config, tables, use_csr=use_csr) @@ -565,7 +565,9 @@ def test_cross_label_shortest_scopes_rebind_each_parameter(cross_label_graph, sc queries = { "direct": body, "alias": prefix + "WITH p AS q RETURN length(q) AS hops, nodes(q) AS ids", + "call": "CALL { " + body + " } RETURN hops, ids", "union": body + " UNION ALL " + body, + "call_union": "CALL { " + body + " UNION ALL " + body + " } RETURN hops, ids", } query = queries[scope] copies = 2 if "union" in scope else 1