diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index 591531c..63b2a8f 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -26,6 +26,15 @@ matches and visitors to traverse the fold input and body while respecting iterat scope. Prefer parsing Cypher text to construct these expressions; the parser resolves local variable references so they cannot accidentally bind to graph variables with the same name. +## Rust AST migration for path values + +Named ordinary paths add `PathPattern.variable`; set it to `None` in existing unnamed-path struct +literals, or `Some(name)` for `p = (a)-[:R]->(b)`. Shortest-path names still use +`ShortestPathSpec.variable`. Exhaustive matches and visitors must handle `LogicalOperator::PathMatch` +by traversing its input and path pattern. It retains ordered node and relationship records, so a path +value must not be replaced by its length. Public `nodes(path)` still returns integer keys; the Cypher +guide describes the separate Arrow record returned by a path value. + ## Quick Start ```rust diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index d6834ad..114ea88 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -179,6 +179,9 @@ pub struct NodePattern { /// A path pattern connecting nodes through relationships #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct PathPattern { + /// Named ordinary path (`p = (a)-[:R]->(b)`). + #[serde(default)] + pub variable: Option, /// Starting node pub start_node: NodePattern, /// Relationships and intermediate nodes diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index c539f56..0dc28d8 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -247,6 +247,23 @@ impl Adjacency { }) } + /// Inspect one occurrence so cancelled work also stops through long tombstone runs. + fn next_candidate(&self, key: i64, offset: &mut usize) -> Option> { + let base = self.base_neighbors(key); + let position = *offset; + let neighbor = if position < base.len() { + base[position] + } else { + *self.added.get(&key)?.get(position - base.len())? + }; + *offset += 1; + let removed = self + .removed + .get(&key) + .is_some_and(|r| r.binary_search(&position).is_ok()); + Some((!removed).then_some(neighbor)) + } + /// Live edges: base + appended - removed. pub fn num_edges(&self) -> u64 { self.base.csr.num_edges() + self.added_count - self.removed_count @@ -827,6 +844,21 @@ impl Neighbors { ) } + /// Inspect one raw edge, including tombstones and suppressed reverse self-loops. + pub(crate) fn next_candidate( + &self, + key: i64, + position: &mut NeighborPosition, + ) -> Option> { + if let Some(candidate) = self.forward.next_candidate(key, &mut position.forward) { + return Some(candidate); + } + self.reverse + .as_ref()? + .next_candidate(key, &mut position.reverse) + .map(|candidate| candidate.filter(|&n| n as i64 != key)) + } + /// Edges visited over all keys (self-loops count once per direction in this estimate). pub fn num_edges(&self) -> u64 { self.forward.num_edges() + self.reverse.as_ref().map_or(0, |r| r.num_edges()) @@ -1569,6 +1601,13 @@ impl ExtensionPlanner for CsrExtensionPlanner { physical_inputs: &[Arc], session_state: &SessionState, ) -> datafusion::common::Result>> { + if let Some(node) = node + .as_any() + .downcast_ref::() + { + return crate::path_match::plan_path_match(node, physical_inputs, session_state) + .map(Some); + } if let Some(node) = node .as_any() .downcast_ref::() diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index f1acce5..777b81b 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -37,6 +37,7 @@ pub struct RelationshipInstance { /// Planning context that tracks state during plan building pub struct PlanningContext<'a> { pub analysis: &'a QueryAnalysis, + 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) comprehension_elements: Vec, @@ -46,6 +47,7 @@ impl<'a> PlanningContext<'a> { pub fn new(analysis: &'a QueryAnalysis) -> Self { Self { analysis, + path_sources: Default::default(), relationship_instance_idx: HashMap::new(), comprehension_elements: Vec::new(), } @@ -220,6 +222,18 @@ fn analyze_operator( analysis.required_datasets.insert(label.clone()); } } + LogicalOperator::PathMatch { input, path } => { + analyze_operator(input, analysis, rel_counter)?; + for node in + std::iter::once(&path.start_node).chain(path.segments.iter().map(|s| &s.end_node)) + { + if let (Some(variable), Some(label)) = (&node.variable, node.labels.first()) { + analysis + .var_to_label + .insert(variable.clone(), label.clone()); + } + } + } LogicalOperator::ShortestPath { input, .. } => { analyze_operator(input, 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 5688010..5bbee1a 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -27,6 +27,7 @@ impl DataFusionPlanner { let mut analysis = ctx.analysis.clone(); Self::collect_labels(input, &mut analysis.var_to_label); let mut branch = PlanningContext::new(&analysis); + branch.path_sources = ctx.path_sources.clone(); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); plans.push(self.build_operator(&mut branch, input)?); // Branches consume distinct relationship instances, in analysis order. diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 6385b6a..5d641f9 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -95,6 +95,15 @@ impl DataFusionPlanner { LogicalOperator::OptionalMatch { optional, .. } => { Self::collect_variables(optional, vars); } + LogicalOperator::PathMatch { input, path } => { + Self::collect_variables(input, vars); + vars.extend( + std::iter::once(&path.start_node) + .chain(path.segments.iter().map(|s| &s.end_node)) + .filter_map(|n| n.variable.clone()), + ); + vars.extend(path.variable.clone()); + } LogicalOperator::ShortestPath { input, path_variable, @@ -171,6 +180,16 @@ impl DataFusionPlanner { LogicalOperator::OptionalMatch { optional, .. } => { Self::collect_labels(optional, labels) } + LogicalOperator::PathMatch { input, path } => { + Self::collect_labels(input, labels); + for node in std::iter::once(&path.start_node) + .chain(path.segments.iter().map(|s| &s.end_node)) + { + if let (Some(variable), Some(label)) = (&node.variable, node.labels.first()) { + labels.insert(variable.clone(), label.clone()); + } + } + } LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, 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..4a8986d 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -18,6 +18,7 @@ mod comprehension_ops; mod expand_ops; mod helpers; mod join_builder; +mod path_ops; use super::DataFusionPlanner; use crate::error::Result; @@ -137,6 +138,7 @@ impl DataFusionPlanner { optional, keys, } => self.build_optional_match(ctx, input, optional, keys), + LogicalOperator::PathMatch { input, path } => self.build_path_match(ctx, input, path), LogicalOperator::ShortestPath { input, source_variable, diff --git a/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs new file mode 100644 index 0000000..7cf18c2 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +use crate::ast::PathPattern; +use crate::case_insensitive::qualify_column; +use crate::datafusion_planner::{analysis::PlanningContext, DataFusionPlanner}; +use crate::error::{GraphError, Result}; +use crate::logical_plan::LogicalOperator; +use crate::path_inventory::{EdgeSource, NodeSource}; +use crate::path_match::{Binding, PathMatchNode, PathSpec, Step}; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use std::sync::Arc; + +impl DataFusionPlanner { + pub(crate) fn build_path_match( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + path: &PathPattern, + ) -> Result { + let error = |message: String| GraphError::UnsupportedFeature { + feature: message, + location: snafu::Location::new(file!(), line!(), column!()), + }; + let catalog = self.get_catalog()?; + let input = self.build_operator(ctx, input)?; + let mut node_labels = Vec::::new(); + let mut edge_labels = Vec::::new(); + for segment in &path.segments { + if segment.relationship.types.is_empty() { + return Err(error("named paths need explicit relationship types".into())); + } + for kind in &segment.relationship.types { + if edge_labels.iter().any(|t| t.eq_ignore_ascii_case(kind)) { + continue; + } + let mapping = self.get_relationship_mapping(kind)?; + for label in [&mapping.source_label, &mapping.target_label] { + let label = label.as_ref().ok_or_else(|| { + error("named paths need configured endpoint labels".into()) + })?; + if !node_labels.iter().any(|l| l.eq_ignore_ascii_case(label)) { + node_labels.push(label.clone()); + } + } + edge_labels.push(mapping.relationship_type.clone()); + } + } + let mut plans = vec![input]; + let mut nodes = Vec::new(); + for label in &node_labels { + let mapping = self + .config + .get_node_mapping(label) + .ok_or_else(|| error(format!("missing node mapping {label}")))?; + let source = catalog + .node_source(label) + .ok_or_else(|| error(format!("missing node table {label}")))?; + let schema = source.schema(); + let field = schema + .fields() + .iter() + .find(|f| f.name().eq_ignore_ascii_case(&mapping.id_field)) + .ok_or_else(|| error(format!("missing node ID {}", mapping.id_field)))?; + if !crate::csr_expand::is_supported_id_type(field.data_type()) { + return Err(error(format!( + "path node ID '{}' has unsupported type {}", + field.name(), + field.data_type() + ))); + } + nodes.push(NodeSource { + label: mapping.label.clone(), + id: mapping.id_field.clone(), + schema, + }); + plans.push( + LogicalPlanBuilder::scan(label.to_lowercase(), source, None) + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to scan path nodes", e))?, + ); + } + let mut edges = Vec::new(); + for label in &edge_labels { + let mapping = self.get_relationship_mapping(label)?; + let source = catalog + .relationship_source(label) + .ok_or_else(|| error(format!("missing relationship table {label}")))?; + let schema = source.schema(); + for id in [&mapping.source_id_field, &mapping.target_id_field] { + let field = schema + .fields() + .iter() + .find(|f| f.name().eq_ignore_ascii_case(id)) + .ok_or_else(|| error(format!("missing relationship ID {id}")))?; + if !crate::csr_expand::is_supported_id_type(field.data_type()) { + return Err(error(format!( + "path relationship ID '{id}' has unsupported type {}", + field.data_type() + ))); + } + } + edges.push(EdgeSource { + label: mapping.relationship_type.clone(), + source_label: node_labels + .iter() + .position(|l| Some(l) == mapping.source_label.as_ref()) + .unwrap(), + target_label: node_labels + .iter() + .position(|l| Some(l) == mapping.target_label.as_ref()) + .unwrap(), + source_id: mapping.source_id_field.clone(), + target_id: mapping.target_id_field.clone(), + schema, + }); + plans.push( + LogicalPlanBuilder::scan(label.to_lowercase(), source, None) + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to scan path edges", e))?, + ); + } + let mut bindings = Vec::new(); + for node in + std::iter::once(&path.start_node).chain(path.segments.iter().map(|s| &s.end_node)) + { + let variable = node + .variable + .clone() + .ok_or_else(|| error("named path endpoint was not assigned a variable".into()))?; + let label = ctx + .analysis + .var_to_label + .get(&variable) + .or_else(|| node.labels.first()) + .ok_or_else(|| { + error(format!("path variable '{variable}' needs a resolved label")) + })?; + let label = node_labels + .iter() + .position(|l| l.eq_ignore_ascii_case(label)) + .ok_or_else(|| { + error(format!( + "path label '{label}' is incompatible with its relationship types" + )) + })?; + bindings.push(Binding { + id_column: qualify_column(&variable, &nodes[label].id), + variable, + label, + properties: node.properties.clone(), + }); + } + let shortest = path.shortest.as_ref().map(|s| s.all); + let mut steps = Vec::new(); + for segment in &path.segments { + let rel = &segment.relationship; + let min = rel.length.as_ref().and_then(|l| l.min).unwrap_or(1); + let max = rel + .length + .as_ref() + .map(|l| { + l.max.unwrap_or(if shortest.is_some() { + u32::MAX + } else { + crate::MAX_VARIABLE_LENGTH_HOPS + }) + }) + .unwrap_or(1); + if min > max || (shortest.is_none() && max > crate::MAX_VARIABLE_LENGTH_HOPS) { + return Err(error(format!("invalid named path range {min}..{max}; ordinary variable-length paths allow at most {} hops", crate::MAX_VARIABLE_LENGTH_HOPS))); + } + steps.push(Step { + types: rel + .types + .iter() + .map(|t| { + edge_labels + .iter() + .position(|l| l.eq_ignore_ascii_case(t)) + .unwrap() + }) + .collect(), + direction: rel.direction.clone(), + min, + max, + relationship: rel.variable.clone(), + properties: rel.properties.clone(), + }); + } + let variable = path + .variable + .as_ref() + .or_else(|| path.shortest.as_ref().and_then(|s| s.variable.as_ref())) + .cloned() + .ok_or_else(|| error("rich path needs a variable".into()))?; + let node = PathMatchNode::try_new( + plans, + PathSpec { + variable, + sources: ctx.path_sources.clone(), + nodes, + edges, + bindings, + steps, + shortest, + }, + ) + .map_err(|e| self.plan_error("Failed to plan path value", e))?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } +} diff --git a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs index e4f47a2..171e787 100644 --- a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs +++ b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs @@ -124,7 +124,14 @@ impl ScalarUDFImpl for CypherField { let Some(Some(ScalarValue::Utf8(Some(name)))) = args.scalar_arguments.get(1) else { return plan_err!("cypher_field takes a field name literal"); }; - let data_type = match args.arg_fields[0].data_type() { + let input_type = args.arg_fields[0].data_type(); + let input_type = match input_type { + DataType::Struct(fields) if crate::path_value::entity_kind(input_type).is_some() => { + fields[0].data_type() + } + other => other, + }; + let data_type = match input_type { DataType::Struct(fields) => match struct_field(fields, name) { Some(i) => fields[i].data_type().clone(), None => return plan_err!("no key '{name}' in this map"), @@ -141,13 +148,23 @@ impl ScalarUDFImpl for CypherField { return exec_err!("cypher_field takes a field name literal"); }; let base = args.args[0].to_array(args.number_rows)?; + let parent_nulls = base.logical_nulls(); + let base = if crate::path_value::entity_kind(base.data_type()).is_some() { + base.as_struct().column(0).clone() + } else { + base + }; let result = match base.data_type() { DataType::Struct(fields) => { let i = struct_field(fields, name).expect("checked when planning"); let column = base.as_struct().column(i).clone(); // A null map gives null keys, whatever its child arrays hold. let nulls = arrow::buffer::NullBuffer::union( - base.logical_nulls().as_ref(), + arrow::buffer::NullBuffer::union( + parent_nulls.as_ref(), + base.logical_nulls().as_ref(), + ) + .as_ref(), column.logical_nulls().as_ref(), ); let data = column.to_data().into_builder().nulls(nulls).build()?; diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 919cfc6..6b42030 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -243,7 +243,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { / lit(1000_i64), ]), // A path's length and node ids are columns of its variable (see shortest_path). - "length" | "nodes" => match &args[0] { + "length" | "nodes" | "relationships" | "__path_node_records" => match &args[0] { VE::Variable(path) => col(qualify_column(path, &name.to_lowercase())), _ => Expr::Literal(datafusion::scalar::ScalarValue::Null, None), }, @@ -253,7 +253,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { "left" | "right" | "substring" | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" | "last" | "tail" | "keys" | "__entity_keys" | "labels" | "__predicate_all" | "__predicate_any" | "__predicate_none" - | "__predicate_single" => { + | "__predicate_single" | "startnode" | "endnode" => { super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) } "range" => { @@ -551,6 +551,29 @@ pub(crate) fn bind_value( field: p.property.clone(), }; } + ValueExpression::ScalarFunction { name, args } + if matches!( + name.to_lowercase().as_str(), + "length" | "nodes" | "relationships" | "__path_node_records" + ) => + { + if let [base] = args.as_slice() { + let value = match base { + ValueExpression::Variable(var) => has(&var.to_lowercase()), + _ => true, + }; + if value { + *e = ValueExpression::Field { + base: Box::new(base.clone()), + field: if name == "__path_node_records" { + "node_records".into() + } else { + name.to_lowercase() + }, + }; + } + } + } ValueExpression::ScalarFunction { name, args } if name.eq_ignore_ascii_case("keys") => { if let [ValueExpression::Variable(var)] = args.as_slice() { let properties = entity_properties(var); diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index a11afea..26476fc 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -18,7 +18,7 @@ pub mod analysis; mod builder; mod config_helpers; mod cypher_udf; -mod expression; +pub(crate) mod expression; mod join_ops; mod list_folds; mod math_functions; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index dbf283c..d96a16b 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -52,7 +52,10 @@ impl ScalarUDFImpl for CypherScalar { | "__predicate_single" => list_field(t) .is_some_and(|f| matches!(f.data_type(), DataType::Boolean | DataType::Null)), "keys" | "__entity_keys" => matches!(t, DataType::Struct(_)), - "labels" => false, + "labels" => crate::path_value::entity_kind(t) == Some("node"), + "startnode" | "endnode" => { + crate::path_value::entity_kind(t) == Some("relationship") + } "size" | "reverse" => is_string(t) || list_field(t).is_some(), "last" | "tail" => list_field(t).is_some(), "left" | "right" | "substring" if i > 0 => t.is_integer(), @@ -66,6 +69,10 @@ impl ScalarUDFImpl for CypherScalar { "__predicate_all" | "__predicate_any" | "__predicate_none" | "__predicate_single" => { DataType::Boolean } + "startnode" | "endnode" => match &types[0] { + DataType::Struct(fields) => fields[2].data_type().clone(), + _ => DataType::Null, + }, "size" => DataType::Int64, "range" => DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), "split" | "keys" | "__entity_keys" | "labels" => { @@ -91,6 +98,36 @@ impl ScalarUDFImpl for CypherScalar { args.number_rows, ))); } + if matches!(self.name.as_str(), "startnode" | "endnode") { + let input = args.args[0].to_array(args.number_rows)?; + let column = input + .as_struct() + .column(if self.name == "startnode" { 2 } else { 3 }); + let nulls = arrow::buffer::NullBuffer::union( + input.logical_nulls().as_ref(), + column.logical_nulls().as_ref(), + ); + return Ok(ColumnarValue::Array(arrow::array::make_array( + column.to_data().into_builder().nulls(nulls).build()?, + ))); + } + if self.name == "labels" { + let input = args.args[0].to_array(args.number_rows)?; + let labels = input + .as_struct() + .column(1) + .as_struct() + .column(0) + .as_string::(); + let mut result = ListBuilder::new(StringBuilder::new()); + for row in 0..input.len() { + if input.is_valid(row) { + result.values().append_value(labels.value(row)); + } + result.append(input.is_valid(row)); + } + return Ok(ColumnarValue::Array(Arc::new(result.finish()))); + } if let Some(kind) = self.name.strip_prefix("__predicate_") { let input = args.args[0].to_array(args.number_rows)?; let mut result = arrow::array::BooleanBuilder::with_capacity(input.len()); @@ -117,13 +154,18 @@ impl ScalarUDFImpl for CypherScalar { } if matches!(self.name.as_str(), "keys" | "__entity_keys") { let input = args.args[0].to_array(args.number_rows)?; - let map = input.as_struct(); + let typed_entity = crate::path_value::entity_kind(input.data_type()).is_some(); + let map = if typed_entity { + input.as_struct().column(0).as_struct() + } else { + input.as_struct() + }; let nulls: Vec<_> = map .columns() .iter() .map(|column| column.logical_nulls()) .collect(); - let entity = self.name == "__entity_keys"; + let entity = self.name == "__entity_keys" || typed_entity; let mut result = ListBuilder::new(StringBuilder::new()); for row in 0..input.len() { let missing = |i: usize| nulls[i].as_ref().is_some_and(|nulls| nulls.is_null(row)); diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index 6d89ced..e024968 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -48,6 +48,9 @@ pub mod lance_vector_search; pub mod logical_plan; pub mod parameter_substitution; pub mod parser; +mod path_inventory; +mod path_match; +mod path_value; mod pattern_alternatives; mod pattern_order; pub mod query; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 05c8284..9f80f30 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -137,6 +137,12 @@ pub enum LogicalOperator { /// same variables, possibly with different labels; columns a branch lacks are null. Union { inputs: Vec }, + /// A named path retaining actual relationship occurrences and typed node records. + PathMatch { + input: Box, + path: PathPattern, + }, + /// `shortestPath` / `allShortestPaths` between two variables bound in `input`; adds /// `__length` and `__nodes` (node ids, source first). ShortestPath { @@ -223,6 +229,7 @@ pub struct LogicalPlanner<'a> { comprehensions: usize, /// The node label of each path variable's nodes, for `nodes(path)`. path_labels: HashMap, + rich_paths: bool, /// The rows a projection or WHERE started from, while its EXISTS marks and /// comprehensions are being added: each is seeded from these, not from the rows with /// the previous ones joined on, so planned copies grow linearly, not exponentially. @@ -240,6 +247,7 @@ impl<'a> LogicalPlanner<'a> { node_lists: HashMap::new(), comprehensions: 0, path_labels: HashMap::new(), + rich_paths: false, seed_source: None, } } @@ -259,6 +267,12 @@ impl<'a> LogicalPlanner<'a> { } fn retain_variable_alternation(&self, patterns: &[GraphPattern]) -> bool { + if patterns + .iter() + .all(|p| matches!(p, GraphPattern::Path(path) if path.variable.is_some())) + { + return true; + } let Some(catalog) = self.adjacency else { return false; }; @@ -324,6 +338,9 @@ impl<'a> LogicalPlanner<'a> { /// 100). `None` keeps the written order: start already right, no estimates, or a /// variable-length segment. fn cheaper_orientation(&self, path: &PathPattern) -> Option { + if path.variable.is_some() || path.shortest.is_some() { + return None; + } let bound = |n: &NodePattern| { n.variable .as_ref() @@ -379,11 +396,13 @@ impl<'a> LogicalPlanner<'a> { start_node: nodes[0].clone(), segments, shortest: path.shortest.clone(), + variable: path.variable.clone(), } } /// 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); // 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)?; @@ -519,15 +538,15 @@ impl<'a> LogicalPlanner<'a> { }) } - /// Node and path variables used as values: a node is its id (so `a = b`, `a IN list`, - /// `collect(a)`, `count(DISTINCT a)` and `a IS NULL` work on ids) and a path is its - /// length (so `path IS NULL` is true when no path was found). Bare items of WITH and - /// RETURN, which carry whole nodes, are not resolved. + /// Resolve node variables to ids in scalar expressions while retaining path values. + /// Path functions and null checks are lowered according to the path representation. + /// Bare WITH and RETURN items, which carry whole nodes, are not resolved. fn resolve_value(&self, expr: &mut ValueExpression) { match expr { ValueExpression::Variable(var) => { let property = match self.variables.get(var).map(String::as_str) { - Some("Path") => Some("length".to_string()), + Some("Path") if !self.rich_paths => Some("length".to_string()), + Some("Path") => None, Some(label) => self .config .get_node_mapping(label) @@ -576,7 +595,10 @@ impl<'a> LogicalPlanner<'a> { && matches!(args.as_slice(), [ValueExpression::Variable(_)]) => {} // Path functions take the path variable itself. ValueExpression::ScalarFunction { name, args } - if matches!(name.to_lowercase().as_str(), "length" | "nodes") => {} + if matches!( + name.to_lowercase().as_str(), + "length" | "nodes" | "relationships" | "__path_node_records" + ) => {} ValueExpression::ScalarFunction { args, .. } | ValueExpression::AggregateFunction { args, .. } => { args.iter_mut().for_each(|a| self.resolve_value(a)) @@ -750,9 +772,23 @@ impl<'a> LogicalPlanner<'a> { } } } + GraphPattern::Path(path) if path.variable.is_some() => { + plan = Some(self.plan_named_path(plan, path)?) + } GraphPattern::Path(path) if path.shortest.is_some() => { plan = Some(self.plan_shortest_path(plan, path)?) } + GraphPattern::Path(path) + if self.rich_paths + && path + .segments + .iter() + .any(|s| s.relationship.variable.is_some()) => + { + let mut path = path.clone(); + path.variable = Some(format!("_path_value_{}", self.variables.len())); + plan = Some(self.plan_named_path(plan, &path)?); + } GraphPattern::Path(path) => plan = Some(self.plan_path(plan, path)?), } } @@ -1064,6 +1100,13 @@ impl<'a> LogicalPlanner<'a> { } => { input = self.rewrite_exists_in_value(input, list)?; self.resolve_value(list); + if self.rich_paths { + if let ValueExpression::ScalarFunction { name, .. } = list.as_mut() { + if name.eq_ignore_ascii_case("nodes") { + *name = "__path_node_records".into(); + } + } + } let label = match list.as_ref() { ValueExpression::Variable(l) => self.node_lists.get(l).cloned(), ValueExpression::ScalarFunction { name, args } @@ -1169,6 +1212,12 @@ impl<'a> LogicalPlanner<'a> { self.path_labels.insert(path_variable.clone(), label); } let length = rel.length.as_ref(); + if self.rich_paths { + return Ok(LogicalOperator::PathMatch { + input: Box::new(plan.expect("both ends planned")), + path: path.clone(), + }); + } Ok(LogicalOperator::ShortestPath { input: Box::new(plan.expect("both end nodes are planned")), source_variable: ends[0].clone(), @@ -1481,6 +1530,54 @@ impl<'a> LogicalPlanner<'a> { Ok(plan) } + fn plan_named_path( + &mut self, + base: Option, + path: &PathPattern, + ) -> Result { + let mut path = path.clone(); + if path.start_node.variable.is_none() { + path.start_node.variable = Some(format!("_path_node_{}", self.variables.len())); + } + let start = path.start_node.variable.as_ref().unwrap(); + let base = if self.variables.contains_key(start) { + base.expect("bound variables have input rows") + } else { + let scan = self.plan_node_scan(&path.start_node)?; + match base { + Some(input) => LogicalOperator::Join { + left: Box::new(input), + right: Box::new(scan), + join_type: JoinType::Cross, + }, + None => scan, + } + }; + for segment in &mut path.segments { + if segment.end_node.variable.is_none() { + segment.end_node.variable = Some(format!("_path_node_{}", self.variables.len())); + } + let variable = segment.end_node.variable.as_ref().unwrap(); + self.validate_variable_label(variable, &segment.end_node.labels)?; + let label = self + .variables + .get(variable) + .cloned() + .or_else(|| segment.end_node.labels.first().cloned()) + .ok_or_else(|| GraphError::UnsupportedFeature { + feature: "named path boundary needs a resolved node label".into(), + location: Location::new(file!(), line!(), column!()), + })?; + self.variables.insert(variable.clone(), label); + } + self.variables + .insert(path.variable.clone().unwrap(), "Path".into()); + Ok(LogicalOperator::PathMatch { + input: Box::new(base), + path, + }) + } + /// Extract the main variable from a logical plan (for chaining) #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { @@ -1518,6 +1615,14 @@ impl<'a> LogicalPlanner<'a> { message: "a comprehension's element rows bind no pattern variable".to_string(), location: snafu::Location::new(file!(), line!(), column!()), }), + LogicalOperator::PathMatch { path, .. } => Ok(path + .segments + .last() + .unwrap() + .end_node + .variable + .clone() + .unwrap()), LogicalOperator::ShortestPath { target_variable, .. } => Ok(target_variable.clone()), @@ -1559,7 +1664,7 @@ impl<'a> LogicalPlanner<'a> { ValueExpression::Variable(var) => { match self.variables.get(var) { // if it is a node variable, expand to all properties - Some(label) if label != "Unwound" => { + Some(label) if label != "Unwound" && label != "Path" => { let mapping = self.config.get_node_mapping(label).ok_or_else(|| { GraphError::PlanError { message: format!("Node label '{}' doesn't exist", label), @@ -1652,6 +1757,9 @@ impl<'a> LogicalPlanner<'a> { if let Some(label) = self.variables.get(var).cloned() { self.variables.insert(alias.clone(), label); } + if let Some(label) = self.path_labels.get(var).cloned() { + self.path_labels.insert(alias.clone(), label); + } Some(alias.clone()) } (Some(alias), _) => Some(alias.clone()), diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 435eb26..3346936 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -204,6 +204,7 @@ fn unwind_clause(input: &str) -> IResult<&str, UnwindClause> { fn graph_pattern(input: &str) -> IResult<&str, GraphPattern> { alt(( map(shortest_path_pattern, GraphPattern::Path), + map(named_path_pattern, GraphPattern::Path), map(path_pattern, GraphPattern::Path), map(node_pattern, GraphPattern::Node), ))(input) @@ -232,6 +233,18 @@ fn shortest_path_pattern(input: &str) -> IResult<&str, PathPattern> { Ok((input, path)) } +// Named ordinary paths retain their textual node and relationship order. +fn named_path_pattern(input: &str) -> IResult<&str, PathPattern> { + let (input, variable) = delimited( + multispace0, + identifier, + tuple((multispace0, char('='), multispace0)), + )(input)?; + let (input, mut path) = path_pattern(input)?; + path.variable = Some(variable.to_string()); + Ok((input, path)) +} + // Parse a path pattern (only if there are segments) fn path_pattern(input: &str) -> IResult<&str, PathPattern> { let (input, start_node) = node_pattern(input)?; @@ -251,6 +264,7 @@ fn path_pattern(input: &str) -> IResult<&str, PathPattern> { start_node, segments, shortest: None, + variable: None, }, )) } diff --git a/crates/lance-graph/src/path_inventory.rs b/crates/lance-graph/src/path_inventory.rs new file mode 100644 index 0000000..713a1db --- /dev/null +++ b/crates/lance-graph/src/path_inventory.rs @@ -0,0 +1,317 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! One immutable inventory per rich-path operator. An edge's identity is its actual +//! scanned occurrence, shared by every hop, never reconstructed from endpoint pairs. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; + +use arrow::array::{ + new_null_array, Array, ArrayRef, AsArray, Int64Array, RecordBatch, StringArray, StructArray, + UInt32Array, +}; +use arrow::compute::{cast, concat_batches, take}; +use arrow_schema::{DataType, Field, Fields, Schema}; +use datafusion::common::{exec_err, plan_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use futures::{lock::Mutex, StreamExt}; + +/// One source scan per query, shared even by distinct named-path operators. +pub(crate) type SourceBatches = Arc>>; + +use crate::csr_expand::{cooperative_yield, is_supported_id_type}; +use crate::path_value::{entity_type, struct_array}; + +#[derive(Debug, Clone)] +pub(crate) struct NodeSource { + pub label: String, + pub id: String, + pub schema: Arc, +} + +#[derive(Debug, Clone)] +pub(crate) struct EdgeSource { + pub label: String, + pub source_label: usize, + pub target_label: usize, + pub source_id: String, + pub target_id: String, + pub schema: Arc, +} + +pub(crate) fn property_fields(schemas: impl Iterator>) -> Result { + let mut fields = BTreeMap::::new(); + for schema in schemas { + for field in schema.fields() { + let name = field.name().to_lowercase(); + if fields.get(&name).is_some_and(|t| t != field.data_type()) { + return plan_err!( + "path entities have conflicting Arrow types for property '{name}'" + ); + } + fields.insert(name, field.data_type().clone()); + } + } + Ok(fields + .into_iter() + .map(|(name, t)| Field::new(name, t, true)) + .collect()) +} + +pub(crate) fn id_column(batch: &RecordBatch, name: &str) -> Result { + let schema = batch.schema(); + let Some((i, field)) = schema + .fields() + .iter() + .enumerate() + .find(|(_, f)| f.name().eq_ignore_ascii_case(name)) + else { + return exec_err!("path ID column '{name}' is missing"); + }; + if !is_supported_id_type(field.data_type()) { + return exec_err!( + "path ID column '{name}' must be a signed integer or uint8/16/32, got {}", + field.data_type() + ); + } + Ok(cast(batch.column(i), &DataType::Int64)? + .as_primitive::() + .clone()) +} + +fn properties(batch: &RecordBatch, fields: &Fields) -> Result { + let columns = fields + .iter() + .map(|f| { + batch + .schema() + .fields() + .iter() + .position(|old| old.name().eq_ignore_ascii_case(f.name())) + .map(|i| batch.column(i).clone()) + .unwrap_or_else(|| new_null_array(f.data_type(), batch.num_rows())) + }) + .collect(); + Ok(RecordBatch::try_new( + Arc::new(Schema::new(fields.clone())), + columns, + )?) +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct EdgeRef { + pub edge: u32, + pub node: u32, + pub outgoing: bool, +} + +#[derive(Debug)] +pub(crate) struct Inventory { + pub nodes: StructArray, + pub edges: StructArray, + pub node_properties: RecordBatch, + pub edge_properties: RecordBatch, + pub node_labels: Vec, + pub edge_types: Vec, + pub endpoints: Vec>, + pub by_key: HashMap<(usize, i64), u32>, + pub neighbors: Vec>, +} + +impl Inventory { + pub async fn build( + node_sources: &[NodeSource], + edge_sources: &[EdgeSource], + plans: &[Arc], + context: Arc, + sources: &SourceBatches, + ) -> Result { + let mut batches = Vec::with_capacity(plans.len()); + let source_keys = node_sources + .iter() + .map(|s| (false, s.label.to_lowercase())) + .chain(edge_sources.iter().map(|s| (true, s.label.to_lowercase()))); + for (plan, key) in plans.iter().zip(source_keys) { + let mut sources = sources.lock().await; + if !sources.contains_key(&key) { + let mut stream = plan.execute(0, context.clone())?; + let mut read = Vec::new(); + while let Some(batch) = stream.next().await { + read.push(batch?); + cooperative_yield().await; + } + sources.insert(key.clone(), concat_batches(&plan.schema(), &read)?); + } + batches.push(sources[&key].clone()); + } + let node_fields = property_fields(node_sources.iter().map(|s| s.schema.clone()))?; + let edge_fields = property_fields(edge_sources.iter().map(|s| s.schema.clone()))?; + let node_type = entity_type(node_fields.clone(), "node", None); + let edge_type = entity_type(edge_fields.clone(), "relationship", Some(&node_type)); + let mut by_key = HashMap::new(); + let mut node_labels = Vec::new(); + let mut labels = Vec::new(); + let mut ids = Vec::new(); + let mut node_batches = Vec::new(); + for (label, source) in node_sources.iter().enumerate() { + let batch = &batches[label]; + let keys = id_column(batch, &source.id)?; + for row in 0..batch.num_rows() { + let dense = u32::try_from(ids.len()).map_err(|_| { + datafusion::common::DataFusionError::Execution( + "path inventory exceeds u32 row references".into(), + ) + })?; + if keys.is_valid(row) && by_key.insert((label, keys.value(row)), dense).is_some() { + return exec_err!( + "path values require unique node IDs within label '{}'; duplicate ID {}", + source.label, + keys.value(row) + ); + } + node_labels.push(label); + labels.push(source.label.as_str()); + ids.push(keys.is_valid(row).then(|| keys.value(row))); + if row % 4096 == 4095 { + cooperative_yield().await; + } + } + node_batches.push(properties(batch, &node_fields)?); + } + let node_properties = concat_batches(&Arc::new(Schema::new(node_fields)), &node_batches)?; + let identity = StructArray::from(vec![ + ( + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(StringArray::from(labels)) as ArrayRef, + ), + // Null keys never enter the lookup; their records cannot appear in a path. + ( + Arc::new(Field::new("id", DataType::Int64, false)), + Arc::new(Int64Array::from_iter_values( + ids.iter().map(|v| v.unwrap_or_default()), + )) as ArrayRef, + ), + ]); + let nodes = struct_array( + &node_type, + vec![ + Arc::new(StructArray::from(node_properties.clone())), + Arc::new(identity), + ], + ); + let mut neighbors = vec![Vec::new(); nodes.len()]; + let mut edge_batches = Vec::new(); + let mut edge_types = Vec::new(); + let mut occurrence_ids = Vec::new(); + let mut endpoints = Vec::new(); + let mut labels = Vec::new(); + for (kind, source) in edge_sources.iter().enumerate() { + let batch = &batches[node_sources.len() + kind]; + let sources = id_column(batch, &source.source_id)?; + let targets = id_column(batch, &source.target_id)?; + for row in 0..batch.num_rows() { + let edge = u32::try_from(endpoints.len()).map_err(|_| { + datafusion::common::DataFusionError::Execution( + "path inventory exceeds u32 row references".into(), + ) + })?; + let ends = if sources.is_valid(row) && targets.is_valid(row) { + by_key + .get(&(source.source_label, sources.value(row))) + .zip(by_key.get(&(source.target_label, targets.value(row)))) + .map(|(&a, &b)| (a, b)) + } else { + None + }; + if let Some((a, b)) = ends { + neighbors[a as usize].push(EdgeRef { + edge, + node: b, + outgoing: true, + }); + if a != b { + neighbors[b as usize].push(EdgeRef { + edge, + node: a, + outgoing: false, + }); + } + } + endpoints.push(ends); + edge_types.push(kind); + occurrence_ids.push(row as i64); + labels.push(source.label.as_str()); + if row % 4096 == 4095 { + cooperative_yield().await; + } + } + edge_batches.push(properties(batch, &edge_fields)?); + } + let edge_properties = concat_batches(&Arc::new(Schema::new(edge_fields)), &edge_batches)?; + let identity = StructArray::from(vec![ + ( + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(StringArray::from(labels)) as ArrayRef, + ), + ( + Arc::new(Field::new("id", DataType::Int64, false)), + Arc::new(Int64Array::from(occurrence_ids)) as ArrayRef, + ), + ]); + let starts = UInt32Array::from_iter(endpoints.iter().map(|p| p.map(|(a, _)| a))); + let ends = UInt32Array::from_iter(endpoints.iter().map(|p| p.map(|(_, b)| b))); + let edges = struct_array( + &edge_type, + vec![ + Arc::new(StructArray::from(edge_properties.clone())), + Arc::new(identity), + take(&nodes, &starts, None)?, + take(&nodes, &ends, None)?, + ], + ); + Ok(Self { + nodes, + edges, + node_properties, + edge_properties, + node_labels, + edge_types, + endpoints, + by_key, + neighbors, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + #[tokio::test] + async fn duplicate_node_keys_do_not_choose_an_arbitrary_endpoint() { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![1, 1])) as ArrayRef, + )]) + .unwrap(); + let source = NodeSource { + label: "Person".into(), + id: "id".into(), + schema: batch.schema(), + }; + let ctx = SessionContext::new(); + let plan = ctx + .read_batch(batch) + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let error = Inventory::build(&[source], &[], &[plan], ctx.task_ctx(), &Default::default()) + .await + .unwrap_err(); + assert!(error.to_string().contains("duplicate ID 1")); + } +} diff --git a/crates/lance-graph/src/path_match.rs b/crates/lance-graph/src/path_match.rs new file mode 100644 index 0000000..7d99585 --- /dev/null +++ b/crates/lance-graph/src/path_match.rs @@ -0,0 +1,999 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Rich paths use an immutable edge inventory, resumable trail enumeration, and the +//! existing bidirectional BFS on a bipartite node/edge graph for shortest paths. + +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array, RecordBatch, UInt32Array}; +use arrow::compute::take; +use arrow_schema::{Field, Schema, SchemaRef}; +use datafusion::common::{DFSchema, DFSchemaRef, Result}; +use datafusion::execution::session_state::SessionState; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{col, Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; +use futures::{lock::Mutex, StreamExt}; + +use crate::ast::{PropertyValue, RelationshipDirection}; +use crate::csr_expand::{cooperative_yield, Adjacency, Neighbors}; +use crate::path_inventory::{id_column, property_fields, EdgeSource, Inventory, NodeSource}; +use crate::path_value::{entity_type, path_type, paths_array, PathRows}; + +type Parents = HashMap, ahash::RandomState>; + +#[derive(Debug, Clone)] +pub(crate) struct Binding { + pub variable: String, + pub label: usize, + pub id_column: String, + pub properties: HashMap, +} + +#[derive(Debug, Clone)] +pub(crate) struct Step { + pub types: Vec, + pub direction: RelationshipDirection, + pub min: u32, + pub max: u32, + pub relationship: Option, + pub properties: HashMap, +} + +#[derive(Debug)] +pub(crate) struct PathSpec { + pub variable: String, + pub sources: crate::path_inventory::SourceBatches, + pub nodes: Vec, + pub edges: Vec, + pub bindings: Vec, + pub steps: Vec, + /// Some(all): exactly one shortest-path segment, with both endpoints bound. + pub shortest: Option, +} + +#[derive(Clone)] +pub(crate) struct PathMatchNode { + inputs: Vec, + spec: Arc, + schema: DFSchemaRef, + outputs: Vec, +} + +#[derive(Debug, Clone)] +enum Output { + RelationshipValue { step: usize }, + Node { boundary: usize, property: usize }, + Relationship { step: usize, property: usize }, +} + +// DataFusion formats extension nodes during planning; never print cached graph batches. +impl fmt::Debug for PathMatchNode { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PathMatchNode") + .field("variable", &self.spec.variable) + .field("sources", &(self.inputs.len() - 1)) + .finish() + } +} + +impl PathMatchNode { + pub fn try_new(inputs: Vec, spec: PathSpec) -> Result { + let node_fields = property_fields(spec.nodes.iter().map(|s| s.schema.clone()))?; + let edge_fields = property_fields(spec.edges.iter().map(|s| s.schema.clone()))?; + let node_type = entity_type(node_fields.clone(), "node", None); + let edge_type = entity_type(edge_fields.clone(), "relationship", Some(&node_type)); + let mut fields = Vec::new(); + let mut outputs = Vec::new(); + let input = inputs[0].schema(); + for (boundary, binding) in spec.bindings.iter().enumerate().skip(1) { + if input.has_column_with_unqualified_name(&binding.id_column) { + continue; + } + for field in spec.nodes[binding.label].schema.fields() { + let name = crate::case_insensitive::qualify_column(&binding.variable, field.name()); + if fields.iter().any(|f: &Field| f.name() == &name) { + continue; + } + let property = node_fields + .iter() + .position(|f| f.name().eq_ignore_ascii_case(field.name())) + .expect("unified property"); + fields.push(Field::new(name, field.data_type().clone(), true)); + outputs.push(Output::Node { boundary, property }); + } + } + for (step, s) in spec.steps.iter().enumerate() { + if let Some(variable) = &s.relationship { + if s.min != 1 || s.max != 1 { + return datafusion::common::plan_err!("a relationship variable in a named variable-length path is not supported; use relationships(path)"); + } + fields.push(Field::new(variable.to_lowercase(), edge_type.clone(), true)); + outputs.push(Output::RelationshipValue { step }); + for field in &edge_fields { + fields.push(Field::new( + crate::case_insensitive::qualify_column(variable, field.name()), + field.data_type().clone(), + true, + )); + outputs.push(Output::Relationship { + step, + property: edge_fields + .iter() + .position(|f| f.name() == field.name()) + .unwrap(), + }); + } + } + } + fields.push(Field::new( + spec.variable.to_lowercase(), + path_type(node_type, edge_type), + true, + )); + let schema = Arc::new(input.join(&DFSchema::try_from(Schema::new(fields))?)?); + Ok(Self { + inputs, + spec: Arc::new(spec), + schema, + outputs, + }) + } +} +impl PartialEq for PathMatchNode { + fn eq(&self, other: &Self) -> bool { + self.inputs == other.inputs && Arc::ptr_eq(&self.spec, &other.spec) + } +} +impl Eq for PathMatchNode {} +impl Hash for PathMatchNode { + fn hash(&self, state: &mut H) { + self.inputs.hash(state); + Arc::as_ptr(&self.spec).hash(state); + } +} +impl PartialOrd for PathMatchNode { + fn partial_cmp(&self, other: &Self) -> Option { + Arc::as_ptr(&self.spec).partial_cmp(&Arc::as_ptr(&other.spec)) + } +} +impl UserDefinedLogicalNodeCore for PathMatchNode { + fn name(&self) -> &str { + "PathMatch" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + self.inputs.iter().collect() + } + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + fn expressions(&self) -> Vec { + vec![] + } + fn prevent_predicate_push_down_columns(&self) -> HashSet { + self.schema + .fields() + .iter() + .skip(self.inputs[0].schema().fields().len()) + .map(|f| f.name().clone()) + .collect() + } + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "PathMatch: {}, {}", + self.spec.variable, + if self.spec.shortest.is_some() { + "shortest edge inventory" + } else { + "trail edge inventory" + } + ) + } + fn with_exprs_and_inputs(&self, _: Vec, inputs: Vec) -> Result { + Ok(Self { + inputs, + ..self.clone() + }) + } +} + +#[derive(Debug)] +struct Ready { + inventory: Inventory, + nodes_allowed: Vec, + edges_allowed: Vec, + shortest: Option<(Neighbors, Neighbors)>, +} + +struct PathMatchExec { + inputs: Vec>, + spec: Arc, + outputs: Vec, + schema: SchemaRef, + properties: Arc, + predicates: Vec>>, + ready: Arc>>>, +} + +fn predicate(properties: &HashMap) -> Option { + properties + .iter() + .map(|(key, value)| { + col(key.to_lowercase()).eq(crate::datafusion_planner::expression::to_df_value_expr( + &crate::ast::ValueExpression::Literal(value.clone()), + )) + }) + .reduce(Expr::and) +} + +pub(crate) fn plan_path_match( + node: &PathMatchNode, + inputs: &[Arc], + state: &SessionState, +) -> Result> { + let mut plans = vec![inputs[0].clone()]; + plans.extend( + inputs[1..] + .iter() + .map(|p| Arc::new(CoalescePartitionsExec::new(p.clone())) as Arc), + ); + let nodes = DFSchema::try_from(Schema::new(property_fields( + node.spec.nodes.iter().map(|s| s.schema.clone()), + )?))?; + let edges = DFSchema::try_from(Schema::new(property_fields( + node.spec.edges.iter().map(|s| s.schema.clone()), + )?))?; + let predicates = node + .spec + .bindings + .iter() + .map(|b| (&b.properties, &nodes)) + .chain(node.spec.steps.iter().map(|s| (&s.properties, &edges))) + .map(|(properties, schema)| { + predicate(properties) + .map(|e| state.create_physical_expr(e, schema)) + .transpose() + }) + .collect::>>()?; + let schema = Arc::new(node.schema.as_arrow().clone()); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + inputs[0].output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Arc::new(PathMatchExec { + inputs: plans, + spec: node.spec.clone(), + outputs: node.outputs.clone(), + schema, + properties, + predicates, + ready: Arc::new(Mutex::new(None)), + })) +} + +impl fmt::Debug for PathMatchExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PathMatchExec") + .field("variable", &self.spec.variable) + .field("sources", &(self.inputs.len() - 1)) + .finish() + } +} + +impl DisplayAs for PathMatchExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PathMatchExec: {} edge inventory", self.spec.variable) + } +} +impl ExecutionPlan for PathMatchExec { + fn name(&self) -> &str { + "PathMatchExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + self.inputs.iter().collect() + } + // Inventory builds consume one stream per source; keep the coalescers through optimization. + fn required_input_distribution(&self) -> Vec { + use datafusion::physical_plan::Distribution; + std::iter::once(Distribution::UnspecifiedDistribution) + .chain((1..self.inputs.len()).map(|_| Distribution::SinglePartition)) + .collect() + } + + fn with_new_children( + self: Arc, + inputs: Vec>, + ) -> Result> { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(self.schema.clone()), + inputs[0].output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + inputs, + spec: self.spec.clone(), + outputs: self.outputs.clone(), + schema: self.schema.clone(), + properties, + predicates: self.predicates.clone(), + ready: self.ready.clone(), + })) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.inputs[0].execute(partition, context.clone())?; + let (cell, plans, spec, outputs, schema, predicates) = ( + self.ready.clone(), + self.inputs[1..].to_vec(), + self.spec.clone(), + self.outputs.clone(), + self.schema.clone(), + self.predicates.clone(), + ); + let out_schema = schema.clone(); + let batch_size = context.session_config().batch_size(); + let stream = futures::stream::try_unfold( + (input, None::), + move |(mut input, mut pending)| { + let (cell, plans, spec, outputs, schema, predicates, context) = ( + cell.clone(), + plans.clone(), + spec.clone(), + outputs.clone(), + schema.clone(), + predicates.clone(), + context.clone(), + ); + async move { + loop { + if pending.is_none() { + let Some(batch) = input.next().await else { + return Ok(None); + }; + let batch = batch?; + let graph = { + let mut ready = cell.lock().await; + if ready.is_none() { + *ready = Some(Arc::new( + Ready::build(&spec, &plans, &predicates, context.clone()) + .await?, + )); + } + ready.as_ref().expect("built").clone() + }; + pending = Some(Pending::new(batch, graph, &spec)?); + } + let p = pending.as_mut().unwrap(); + let (rows, paths) = p.next(&spec, batch_size).await?; + if rows.is_empty() { + pending = None; + continue; + } + let indices = UInt32Array::from(rows); + let mut columns: Vec = p + .batch + .columns() + .iter() + .map(|c| take(c, &indices, None)) + .collect::>()?; + for output in &outputs { + let (array, positions) = match output { + Output::RelationshipValue { step } => { + let positions = UInt32Array::from_iter_values( + paths.iter().map(|p| p.edges[p.boundaries[*step]]), + ); + columns.push(take(&p.graph.inventory.edges, &positions, None)?); + continue; + } + Output::Node { boundary, property } => ( + p.graph.inventory.node_properties.column(*property), + UInt32Array::from_iter_values( + paths.iter().map(|p| p.nodes[p.boundaries[*boundary]]), + ), + ), + Output::Relationship { step, property } => ( + p.graph.inventory.edge_properties.column(*property), + UInt32Array::from_iter_values( + paths.iter().map(|p| p.edges[p.boundaries[*step]]), + ), + ), + }; + columns.push(take(array, &positions, None)?); + } + columns.push(paths_array( + &paths, + &p.graph.inventory.nodes, + &p.graph.inventory.edges, + )?); + let batch = RecordBatch::try_new(schema, columns)?; + return Ok(Some((batch, (input, pending)))); + } + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) + } +} + +impl Ready { + async fn build( + spec: &PathSpec, + plans: &[Arc], + predicates: &[Option>], + context: Arc, + ) -> Result { + let inventory = + Inventory::build(&spec.nodes, &spec.edges, plans, context, &spec.sources).await?; + let mask = + |p: &Option>, batch: &RecordBatch| -> Result { + match p { + Some(p) => Ok(p + .evaluate(batch)? + .into_array(batch.num_rows())? + .as_boolean() + .clone()), + None => Ok(BooleanArray::from(vec![true; batch.num_rows()])), + } + }; + let nodes_allowed = predicates[..spec.bindings.len()] + .iter() + .map(|p| mask(p, &inventory.node_properties)) + .collect::>>()?; + let edges_allowed = predicates[spec.bindings.len()..] + .iter() + .map(|p| mask(p, &inventory.edge_properties)) + .collect::>>()?; + let shortest = if spec.shortest.is_some() { + let step = &spec.steps[0]; + let mut from = Vec::new(); + let mut to = Vec::new(); + for (edge, ends) in inventory.endpoints.iter().enumerate() { + if let Some((a, b)) = ends { + let vertex = inventory.nodes.len() as i64 + edge as i64; + let (a, b) = (*a as i64, *b as i64); + match step.direction { + RelationshipDirection::Outgoing => { + from.extend([a, vertex]); + to.extend([vertex, b]); + } + RelationshipDirection::Incoming => { + from.extend([b, vertex]); + to.extend([vertex, a]); + } + RelationshipDirection::Undirected => { + from.extend([a, vertex, b, vertex]); + to.extend([vertex, b, vertex, a]); + } + } + } + if edge % 1024 == 1023 { + cooperative_yield().await; + } + } + let batch = RecordBatch::try_from_iter([ + ("src", Arc::new(Int64Array::from(from)) as ArrayRef), + ("dst", Arc::new(Int64Array::from(to)) as ArrayRef), + ])?; + let stream = || { + let chunks: Vec<_> = (0..batch.num_rows()) + .step_by(1024) + .map(|start| batch.slice(start, (batch.num_rows() - start).min(1024))) + .collect(); + Box::pin(RecordBatchStreamAdapter::new( + batch.schema(), + futures::stream::iter(chunks).then(|batch| async move { + cooperative_yield().await; + Ok(batch) + }), + )) as SendableRecordBatchStream + }; + Some(( + Neighbors::directed(Arc::new( + Adjacency::build(stream(), "src", "dst") + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?, + )), + Neighbors::directed(Arc::new( + Adjacency::build(stream(), "dst", "src") + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?, + )), + )) + } else { + None + }; + Ok(Self { + inventory, + nodes_allowed, + edges_allowed, + shortest, + }) + } +} + +struct Pending { + batch: RecordBatch, + graph: Arc, + keys: Vec>, + row: usize, + walk: Option, + shortest: Option, +} +impl Pending { + fn new(batch: RecordBatch, graph: Arc, spec: &PathSpec) -> Result { + let keys = spec + .bindings + .iter() + .map(|b| { + batch + .schema() + .index_of(&b.id_column) + .ok() + .map(|_| id_column(&batch, &b.id_column)) + .transpose() + }) + .collect::>>()?; + Ok(Self { + batch, + graph, + keys, + row: 0, + walk: None, + shortest: None, + }) + } + fn bound(&self, spec: &PathSpec) -> Option>> { + self.keys + .iter() + .zip(&spec.bindings) + .map(|(keys, binding)| match keys { + None => Some(None), + Some(keys) if keys.is_null(self.row) => None, + Some(keys) => self + .graph + .inventory + .by_key + .get(&(binding.label, keys.value(self.row))) + .copied() + .map(Some), + }) + .collect() + } + async fn next(&mut self, spec: &PathSpec, size: usize) -> Result<(Vec, Vec)> { + let mut rows = Vec::new(); + let mut paths = Vec::new(); + while self.row < self.batch.num_rows() && rows.len() < size { + if self.walk.is_none() && self.shortest.is_none() { + let Some(bound) = self.bound(spec) else { + self.row += 1; + continue; + }; + let source = bound[0].expect("source bound by planner"); + if !allowed(&self.graph.nodes_allowed[0], source) { + self.row += 1; + continue; + } + if let Some(all) = spec.shortest { + let target = bound[1].expect("shortest target bound"); + let (forward, reverse) = + self.graph.shortest.as_ref().expect("shortest adjacency"); + let step = &spec.steps[0]; + let max = step.max.saturating_mul(2); + let (depths, parents) = crate::shortest_path::search_bidirectional( + forward, + reverse, + source as i64, + target as i64, + max, + ) + .await; + if depths + .get(&(target as i64)) + .is_none_or(|&depth| depth / 2 < step.min) + { + self.row += 1; + continue; + } + self.shortest = Some(ShortestPaths::new( + target as i64, + parents, + all, + self.graph.inventory.nodes.len(), + )); + } else { + self.walk = Some(Walk::new(source, bound)); + } + } + let path = if let Some(search) = &mut self.shortest { + search.next().await + } else { + self.walk.as_mut().unwrap().next(spec, &self.graph).await + }; + if let Some(path) = path { + rows.push(self.row as u32); + paths.push(path); + } else { + self.walk = None; + self.shortest = None; + self.row += 1; + } + cooperative_yield().await; + } + Ok((rows, paths)) + } +} + +fn allowed(mask: &BooleanArray, row: u32) -> bool { + mask.is_valid(row as usize) && mask.value(row as usize) +} + +struct Frame { + stage: usize, + hops: u32, + node: u32, + next: usize, + advanced: bool, + entered: bool, +} +struct Walk { + frames: Vec, + path: PathRows, + bound: Vec>, +} +impl Walk { + fn new(source: u32, bound: Vec>) -> Self { + Self { + frames: vec![Frame { + stage: 0, + hops: 0, + node: source, + next: 0, + advanced: false, + entered: false, + }], + path: PathRows { + nodes: vec![source], + edges: vec![], + boundaries: vec![0], + }, + bound, + } + } + fn pop(&mut self) { + let frame = self.frames.pop().unwrap(); + if frame.entered { + self.path.nodes.pop(); + self.path.edges.pop(); + } else if frame.stage > 0 { + self.path.boundaries.pop(); + } + } + async fn next(&mut self, spec: &PathSpec, graph: &Ready) -> Option { + let mut work = 0_usize; + while let Some(frame) = self.frames.last_mut() { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + if frame.stage == spec.steps.len() { + if !frame.advanced { + frame.advanced = true; + return Some(self.path.clone()); + } + self.pop(); + continue; + } + let step = &spec.steps[frame.stage]; + let binding = &spec.bindings[frame.stage + 1]; + if !frame.advanced { + frame.advanced = true; + let same = spec.bindings[..frame.stage + 1] + .iter() + .position(|b| b.variable == binding.variable) + .map(|i| self.path.nodes[self.path.boundaries[i]]); + if frame.hops >= step.min + && graph.inventory.node_labels[frame.node as usize] == binding.label + && self.bound[frame.stage + 1].is_none_or(|n| n == frame.node) + && same.is_none_or(|n| n == frame.node) + && allowed(&graph.nodes_allowed[frame.stage + 1], frame.node) + { + let next = Frame { + stage: frame.stage + 1, + hops: 0, + node: frame.node, + next: 0, + advanced: false, + entered: false, + }; + self.path.boundaries.push(self.path.nodes.len() - 1); + self.frames.push(next); + continue; + } + } + if frame.hops == step.max { + self.pop(); + continue; + } + let Some(edge) = graph.inventory.neighbors[frame.node as usize].get(frame.next) else { + self.pop(); + continue; + }; + frame.next += 1; + let direction = match step.direction { + RelationshipDirection::Outgoing => edge.outgoing, + RelationshipDirection::Incoming => !edge.outgoing || edge.node == frame.node, + RelationshipDirection::Undirected => true, + }; + if !direction + || !step + .types + .contains(&graph.inventory.edge_types[edge.edge as usize]) + || self.path.edges.contains(&edge.edge) + || !allowed(&graph.edges_allowed[frame.stage], edge.edge) + { + continue; + } + let next = Frame { + stage: frame.stage, + hops: frame.hops + 1, + node: edge.node, + next: 0, + advanced: false, + entered: true, + }; + self.path.nodes.push(edge.node); + self.path.edges.push(edge.edge); + self.frames.push(next); + } + None + } +} + +/// Depth-first predecessor enumeration keeps only one in-progress shortest path. +struct ShortestPaths { + parents: Parents, + stack: Vec<(i64, usize)>, + all: bool, + done: bool, + nodes: usize, +} +impl ShortestPaths { + fn new(target: i64, parents: Parents, all: bool, nodes: usize) -> Self { + Self { + parents, + stack: vec![(target, 0)], + all, + done: false, + nodes, + } + } + async fn next(&mut self) -> Option { + if self.done { + return None; + } + let mut work = 0_usize; + while let Some((node, next)) = self.stack.last_mut() { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + match self.parents.get(node) { + None if *next == 0 => { + *next = 1; + let nodes: Vec<_> = self + .stack + .iter() + .rev() + .step_by(2) + .map(|(n, _)| *n as u32) + .collect(); + let edges = self + .stack + .iter() + .rev() + .skip(1) + .step_by(2) + .map(|(n, _)| (*n as usize - self.nodes) as u32) + .collect(); + let end = nodes.len() - 1; + self.done = !self.all; + return Some(PathRows { + nodes, + edges, + boundaries: vec![0, end], + }); + } + Some(parents) if *next < parents.len() => { + let parent = parents[*next]; + *next += 1; + self.stack.push((parent, 0)); + } + _ => { + self.stack.pop(); + } + } + } + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + async fn graph( + edges: &[(i64, i64)], + max: u32, + shortest: Option, + ) -> (PathSpec, Arc) { + let n = edges + .iter() + .flat_map(|&(a, b)| [a, b]) + .max() + .unwrap_or(3) + .max(3); + let nodes = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(1..=n)) as ArrayRef, + )]) + .unwrap(); + let edge_batch = RecordBatch::try_from_iter([ + ( + "src", + Arc::new(Int64Array::from_iter_values(edges.iter().map(|e| e.0))) as ArrayRef, + ), + ( + "dst", + Arc::new(Int64Array::from_iter_values(edges.iter().map(|e| e.1))) as ArrayRef, + ), + ( + "weight", + Arc::new(Int64Array::from_iter_values(0..edges.len() as i64)) as ArrayRef, + ), + ]) + .unwrap(); + let binding = |variable: &str| Binding { + variable: variable.into(), + label: 0, + id_column: format!("{variable}__id"), + properties: HashMap::new(), + }; + let spec = PathSpec { + variable: "p".into(), + sources: Default::default(), + nodes: vec![NodeSource { + label: "Person".into(), + id: "id".into(), + schema: nodes.schema(), + }], + edges: vec![EdgeSource { + label: "R".into(), + source_label: 0, + target_label: 0, + source_id: "src".into(), + target_id: "dst".into(), + schema: edge_batch.schema(), + }], + bindings: vec![binding("a"), binding("b")], + steps: vec![Step { + types: vec![0], + direction: RelationshipDirection::Outgoing, + min: 0, + max, + relationship: None, + properties: HashMap::new(), + }], + shortest, + }; + let ctx = SessionContext::new(); + let mut plans = Vec::new(); + for batch in [nodes, edge_batch] { + plans.push( + ctx.read_batch(batch) + .unwrap() + .create_physical_plan() + .await + .unwrap(), + ); + } + let ready = Ready::build(&spec, &plans, &[None, None, None], ctx.task_ctx()) + .await + .unwrap(); + (spec, Arc::new(ready)) + } + + #[tokio::test] + async fn named_walk_keeps_edge_occurrences_and_never_reuses_one() { + let (spec, graph) = graph(&[(1, 2), (1, 2), (2, 3), (3, 3)], 20, None).await; + let mut walk = Walk::new(0, vec![Some(0), None]); + let mut paths = Vec::new(); + while let Some(path) = walk.next(&spec, &graph).await { + paths.push(path); + } + assert_eq!(paths.len(), 7); // identity, two one-hop, two two-hop, two with self-loop + assert!(paths + .iter() + .all(|p| p.edges.iter().collect::>().len() == p.edges.len())); + let two: Vec<_> = paths + .iter() + .filter(|p| p.edges.len() == 2) + .map(|p| p.edges.clone()) + .collect(); + assert_eq!(two, vec![vec![0, 2], vec![1, 2]]); + let array = paths_array(&paths, &graph.inventory.nodes, &graph.inventory.edges).unwrap(); + let rels = array.as_struct().column(3).as_list::(); + assert_eq!(rels.value(0).len(), 0); + let first = rels.value(1); + assert_eq!( + first + .as_struct() + .column(1) + .as_struct() + .column(1) + .as_primitive::() + .value(0), + 0 + ); + } + + #[tokio::test] + async fn shortest_values_preserve_parallel_edges_and_bounded_output() { + let (spec, graph) = graph( + &[(1, 2), (1, 2), (2, 3), (2, 3), (3, 3)], + u32::MAX, + Some(true), + ) + .await; + let batch = RecordBatch::try_from_iter([ + ("a__id", Arc::new(Int64Array::from(vec![1])) as ArrayRef), + ("b__id", Arc::new(Int64Array::from(vec![3])) as ArrayRef), + ]) + .unwrap(); + let mut pending = Pending::new(batch, graph, &spec).unwrap(); + let mut edges = Vec::new(); + loop { + let (rows, paths) = pending.next(&spec, 1).await.unwrap(); + assert!(rows.len() <= 1); + if rows.is_empty() { + break; + } + assert_eq!(paths[0].nodes, vec![0, 1, 2]); + edges.push(paths[0].edges.clone()); + } + edges.sort(); + assert_eq!(edges, vec![vec![0, 2], vec![0, 3], vec![1, 2], vec![1, 3]]); + } + + #[tokio::test] + async fn rejected_hub_yields_before_scanning_all_edges() { + let (spec, graph) = graph(&(2..100_002).map(|n| (1, n)).collect::>(), 1, None).await; + let mut walk = Walk::new(0, vec![Some(0), Some(0)]); + assert!(walk.next(&spec, &graph).await.is_some()); // zero hop + let mut next = Box::pin(walk.next(&spec, &graph)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(walk.frames[0].next > 0 && walk.frames[0].next <= 1024); + } +} diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs new file mode 100644 index 0000000..16705c6 --- /dev/null +++ b/crates/lance-graph/src/path_value.rs @@ -0,0 +1,389 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Arrow values for paths and their entities. Properties live in a marked child so +//! user property names cannot collide with identity or endpoint metadata. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, Int64Array, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::take; +use arrow_schema::{DataType, Field, Fields}; +use datafusion::common::Result; + +pub(crate) const PROPERTIES: &str = "__lance_graph_properties"; +pub(crate) const IDENTITY: &str = "__lance_graph_identity"; +pub(crate) const START: &str = "__lance_graph_start"; +pub(crate) const END: &str = "__lance_graph_end"; +const ENTITY: &str = "lance_graph.entity"; + +pub(crate) fn entity_kind(data_type: &DataType) -> Option<&str> { + let DataType::Struct(fields) = data_type else { + return None; + }; + fields.first()?.metadata().get(ENTITY).map(String::as_str) +} + +pub(crate) fn entity_type( + properties: Fields, + kind: &str, + node_type: Option<&DataType>, +) -> DataType { + let mut fields = vec![ + Field::new(PROPERTIES, DataType::Struct(properties), false) + .with_metadata(HashMap::from([(ENTITY.into(), kind.into())])), + Field::new( + IDENTITY, + DataType::Struct( + vec![ + Field::new("label", DataType::Utf8, false), + Field::new("id", DataType::Int64, false), + ] + .into(), + ), + false, + ), + ]; + if let Some(node_type) = node_type { + fields.push(Field::new(START, node_type.clone(), true)); + fields.push(Field::new(END, node_type.clone(), true)); + } + DataType::Struct(fields.into()) +} + +pub(crate) fn path_type(node_type: DataType, relationship_type: DataType) -> DataType { + let list = |t| DataType::List(Arc::new(Field::new("item", t, true))); + DataType::Struct( + vec![ + Field::new("length", DataType::Int64, false), + Field::new("nodes", list(DataType::Int64), false), + Field::new("node_records", list(node_type), false), + Field::new("relationships", list(relationship_type), false), + ] + .into(), + ) +} + +pub(crate) fn struct_array(data_type: &DataType, columns: Vec) -> StructArray { + let DataType::Struct(fields) = data_type else { + unreachable!("entity/path struct") + }; + StructArray::new(fields.clone(), columns, None) +} + +/// Ordered row references, independent of property equality or edge endpoints. +#[derive(Debug, Clone)] +pub(crate) struct PathRows { + pub nodes: Vec, + pub edges: Vec, + /// Node positions at the written segment boundaries (variable segments may have many hops). + pub boundaries: Vec, +} + +pub(crate) fn paths_array( + paths: &[PathRows], + nodes: &StructArray, + edges: &StructArray, +) -> Result { + let node_rows = + UInt32Array::from_iter_values(paths.iter().flat_map(|p| p.nodes.iter().copied())); + let edge_rows = + UInt32Array::from_iter_values(paths.iter().flat_map(|p| p.edges.iter().copied())); + let selected_nodes = take(nodes, &node_rows, None)?; + let selected_edges = take(edges, &edge_rows, None)?; + let ids = selected_nodes + .as_struct() + .column(1) + .as_struct() + .column(1) + .clone(); + let list = |values: ArrayRef, lengths: Vec| -> ArrayRef { + Arc::new(ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::from_lengths(lengths), + values, + None, + )) + }; + let node_lengths: Vec<_> = paths.iter().map(|p| p.nodes.len()).collect(); + let edge_lengths: Vec<_> = paths.iter().map(|p| p.edges.len()).collect(); + let data_type = path_type(nodes.data_type().clone(), edges.data_type().clone()); + Ok(Arc::new(struct_array( + &data_type, + vec![ + Arc::new(Int64Array::from_iter_values( + paths.iter().map(|p| p.edges.len() as i64), + )), + list(ids, node_lengths.clone()), + list(selected_nodes, node_lengths), + list(selected_edges, edge_lengths), + ], + ))) +} + +/// Null tests need only the compact path's validity, not its entity inventory. +fn mask_path_null_checks(predicate: &mut crate::ast::BooleanExpression, names: &HashSet) { + use crate::ast::{BooleanExpression as B, PropertyValue, ValueExpression as V}; + match predicate { + B::IsNull(value) | B::IsNotNull(value) => { + if matches!(&*value, V::Variable(name) if names.contains(&name.to_lowercase())) { + *value = V::Literal(PropertyValue::Null); + } + } + B::And(left, right) | B::Or(left, right) => { + mask_path_null_checks(left, names); + mask_path_null_checks(right, names); + } + B::Not(inner) => mask_path_null_checks(inner, names), + _ => {} + } +} + +/// Homogeneous length/nodes queries, including WITH aliases, keep their CSR plan. +/// Returned path values, relationship access and containers need actual path records. +pub(crate) fn needs_values( + query: &crate::ast::CypherQuery, + config: &crate::config::GraphConfig, +) -> bool { + use crate::ast::{GraphPattern, ReadingClause, ValueExpression as V}; + let clauses: Vec<_> = query + .reading_clauses + .iter() + .chain(&query.post_with_reading_clauses) + .chain( + query + .further_with_stages + .iter() + .flat_map(|s| &s.reading_clauses), + ) + .collect(); + let mut names = std::collections::HashSet::new(); + for clause in &clauses { + if let ReadingClause::Match(m) = clause { + for pattern in &m.patterns { + if let GraphPattern::Path(path) = pattern { + if path.variable.is_some() { + return true; + } + if path.shortest.is_some() { + // Numeric keys alone cannot distinguish nodes in different labels, + // even when both endpoints have the same label on an undirected path. + let crosses_labels = path.segments.iter().any(|segment| { + segment.relationship.types.iter().any(|kind| { + config + .get_relationship_mapping(kind) + .is_some_and(|mapping| { + mapping + .source_label + .as_deref() + .zip(mapping.target_label.as_deref()) + .is_some_and(|(source, target)| { + !source.eq_ignore_ascii_case(target) + }) + }) + }) + }); + let different_ends = path + .start_node + .labels + .first() + .zip(path.segments.last().and_then(|s| s.end_node.labels.first())) + .is_some_and(|(source, target)| !source.eq_ignore_ascii_case(target)); + if crosses_labels || different_ends { + return true; + } + if let Some(name) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) + { + names.insert(name.to_lowercase()); + } + } + } + } + } + } + // A bare WITH only renames/carries a path; inspect its eventual consumers instead. + for with in query + .with_clause + .iter() + .chain(query.further_with_stages.iter().map(|s| &s.with_clause)) + { + for item in &with.items { + if let (V::Variable(name), Some(alias)) = (&item.expression, &item.alias) { + if names.contains(&name.to_lowercase()) { + names.insert(alias.to_lowercase()); + } + } + } + } + let mut needed = false; + let mut inspect = |expression: &V| { + let mut value = expression.clone(); + crate::ast::transform_value(&mut value, &mut |v| match v { + V::Boolean(predicate) => mask_path_null_checks(predicate, &names), + V::Case { whens, .. } => { + for (predicate, _) in whens { + mask_path_null_checks(predicate, &names); + } + } + V::ListComprehension { + filter: Some(predicate), + .. + } + | V::PatternComprehension { + filter: Some(predicate), + .. + } => { + mask_path_null_checks(predicate, &names); + } + V::ScalarFunction { name, args } + if matches!(name.to_lowercase().as_str(), "length" | "nodes") => + { + args.clear() + } + V::AggregateFunction { args, .. } + if args.iter().any( + |v| matches!(v, V::Variable(name) if names.contains(&name.to_lowercase())), + ) => + { + needed = true; + } + V::ScalarFunction { name, .. } + if matches!( + name.to_lowercase().as_str(), + "relationships" | "startnode" | "endnode" + ) => + { + needed = true + } + _ => {} + }); + crate::ast::transform_value(&mut value, &mut |v| { + if matches!(v, V::Variable(name) if names.contains(&name.to_lowercase())) { + needed = true; + } + }); + }; + for clause in clauses { + match clause { + ReadingClause::Unwind(u) => inspect(&u.expression), + ReadingClause::Match(m) => { + if let Some(w) = &m.where_clause { + inspect(&V::Boolean(Box::new(w.expression.clone()))); + } + } + } + } + for with in query + .with_clause + .iter() + .chain(query.further_with_stages.iter().map(|s| &s.with_clause)) + { + for item in &with.items { + if !matches!(&item.expression, V::Variable(name) if names.contains(&name.to_lowercase())) + { + inspect(&item.expression); + } + } + if let Some(order) = &with.order_by { + for item in &order.items { + inspect(&item.expression); + } + } + if let Some(filter) = &with.where_clause { + inspect(&V::Boolean(Box::new(filter.expression.clone()))); + } + } + for item in &query.return_clause.items { + inspect(&item.expression); + } + if let Some(order) = &query.order_by { + for item in &order.items { + inspect(&item.expression); + } + } + for filter in query + .where_clause + .iter() + .chain(&query.post_with_where_clause) + .chain( + query + .further_with_stages + .iter() + .filter_map(|s| s.where_clause.as_ref()), + ) + { + inspect(&V::Boolean(Box::new(filter.expression.clone()))); + } + needed +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn direct_shortest_queries_keep_the_csr_route() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let prefix = "MATCH p = shortestPath((a:Person)-[:R*]-(b:Person)) "; + for (suffix, rich) in [ + ("RETURN length(p), nodes(p)", false), + ("RETURN p IS NULL, p IS NOT NULL", false), + ( + "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END", + false, + ), + ( + "RETURN CASE WHEN p IS NOT NULL THEN length(p) ELSE -1 END", + false, + ), + ( + "WITH p AS q WHERE q IS NOT NULL OR NOT q IS NULL RETURN length(q)", + false, + ), + ("RETURN p, p IS NULL", true), + ("RETURN relationships(p) IS NULL", true), + ("RETURN [n IN nodes(p) | n.id]", false), + ("WITH p RETURN nodes(p)", false), + ( + "WITH p AS q WITH q AS r RETURN [n IN nodes(r) | n.id]", + false, + ), + ("WITH p AS q RETURN q", true), + ("WITH p AS q RETURN relationships(q)", true), + ( + "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r)", + true, + ), + ("RETURN p", true), + ("RETURN relationships(p)", true), + ("RETURN length(head(collect(p)))", true), + ("WITH collect(p) AS ps UNWIND ps AS p RETURN nodes(p)", true), + ] { + let query = crate::parser::parse_cypher_query(&format!("{prefix}{suffix}")).unwrap(); + assert_eq!(needs_values(&query, &config), rich, "{suffix}"); + } + } + + #[test] + fn cross_label_shortest_paths_need_typed_identity() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + .unwrap(); + for ends in ["(b:Place)", "(b:Person)", "(b)"] { + let query = crate::parser::parse_cypher_query(&format!( + "MATCH p = shortestPath((a:Person)-[:LIVES*0..]-{ends}) RETURN nodes(p)" + )) + .unwrap(); + assert!(needs_values(&query, &config)); + } + } +} diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index 2d4aade..818b9dd 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -48,6 +48,7 @@ pub(crate) fn reorder_patterns( filter: Option<&BooleanExpression>, config: &GraphConfig, ) -> Option> { + if patterns.iter().any(|p| matches!(p, GraphPattern::Path(path) if path.variable.is_some() || path.shortest.is_some())) { return None; } let graph = PatternGraph::from_patterns(patterns)?; if graph.edges.is_empty() || !graph.is_connected() { return None; @@ -432,6 +433,7 @@ impl PatternGraph { start_node, segments, shortest: None, + variable: None, })); } // The next path starts at a bound node that still has an unused edge. diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index bb2d176..1798c9b 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -263,7 +263,11 @@ impl SemanticAnalyzer { // Register end node self.register_node_variable(&segment.end_node)?; } - if let Some(path_var) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) { + if let Some(path_var) = path + .variable + .as_ref() + .or_else(|| path.shortest.as_ref().and_then(|s| s.variable.as_ref())) + { self.register_projection_alias(path_var); if let Some(info) = self.variables.get_mut(&path_var.to_lowercase()) { info.variable_type = VariableType::Path; @@ -440,7 +444,7 @@ impl SemanticAnalyzer { match function_name.as_str() { "tolower" | "lower" | "toupper" | "upper" | "abs" | "round" | "floor" | "ceil" | "sqrt" | "trim" | "ltrim" | "rtrim" | "reverse" | "last" - | "tail" | "keys" | "labels" => { + | "tail" | "keys" | "labels" | "relationships" | "startnode" | "endnode" => { if args.len() != 1 { return Err(GraphError::PlanError { message: format!( @@ -1193,6 +1197,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1274,6 +1279,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1398,6 +1404,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1461,6 +1468,7 @@ mod tests { }, ], shortest: None, + variable: None, }; // Custom config that knows both relationship types to avoid warnings muddying the assertion diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 9c37dec..024889c 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -349,6 +349,7 @@ impl ExecutionPlan for ShortestPathExec { &spec, &schema, ) + .await } }); Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) @@ -386,7 +387,7 @@ async fn build_adjacency( /// Breadth-first search from `source` until every node in `targets` is reached (or the /// frontier empties, or `max_hops`). Returns each reached node's depth and, for every node /// at depth d > 0, its predecessors at depth d - 1. -fn search( +async fn search( adjacency: &Neighbors, source: i64, targets: &HashSet, @@ -400,7 +401,7 @@ fn search( let mut level = 0; while !frontier.is_empty() && !remaining.is_empty() && level < max_hops { level += 1; - let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents); + let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents).await; // Finish the level before stopping, so every shortest path is recorded. for n in &next { remaining.remove(n); @@ -411,7 +412,7 @@ fn search( } /// Finish the whole level so parallel edges and alternative shortest paths survive. -fn expand_level( +async fn expand_level( adjacency: &Neighbors, frontier: &[i64], level: u32, @@ -419,8 +420,17 @@ fn expand_level( parents: &mut FastMap>, ) -> Vec { let mut next = Vec::new(); + let mut work = 0_usize; for &node in frontier { - adjacency.for_each_neighbor(node, |n| { + let mut position = crate::csr_expand::NeighborPosition::default(); + while let Some(candidate) = adjacency.next_candidate(node, &mut position) { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + let Some(n) = candidate else { + continue; + }; let n = n as i64; match depth.get(&n) { Some(&d) if d == level => parents.entry(n).or_default().push(node), @@ -431,14 +441,14 @@ fn expand_level( next.push(n); } } - }); + } } next } /// Search a single pair using complete layers from the smaller frontier. At the /// first intersection every meeting node lies on the same depth cut of a shortest path. -fn search_bidirectional( +pub(crate) async fn search_bidirectional( forward: &Neighbors, reverse: &Neighbors, source: i64, @@ -459,7 +469,7 @@ fn search_bidirectional( while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops { if frontier.len() <= reverse_frontier.len() { level += 1; - frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents); + frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents).await; } else { reverse_level += 1; reverse_frontier = expand_level( @@ -468,7 +478,8 @@ fn search_bidirectional( reverse_level, &mut reverse_depth, &mut reverse_parents, - ); + ) + .await; } let mut pending: Vec<_> = frontier .iter() @@ -480,7 +491,12 @@ fn search_bidirectional( } // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths. let mut visited = HashSet::new(); + let mut work = 0_usize; while let Some(node) = pending.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } if !visited.insert(node) { continue; } @@ -499,10 +515,15 @@ fn search_bidirectional( /// Shortest paths from the search source to `target`, as node ids from source to target: /// all of them, or the first found. -fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { +async fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { let mut out = Vec::new(); let mut stack = vec![vec![target]]; + let mut work = 0_usize; while let Some(path) = stack.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } let head = *path.last().expect("non-empty"); match parents.get(&head) { None => { @@ -526,7 +547,7 @@ fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec = by_source - .iter() - .map(|(&s, ts)| { - let result = match reverse_adjacency.filter(|_| ts.len() == 1) { - Some(reverse) => search_bidirectional( + let mut searches = FastMap::default(); + for (&source, targets) in &by_source { + let result = match reverse_adjacency.filter(|_| targets.len() == 1) { + Some(reverse) => { + search_bidirectional( adjacency, reverse, - s, - *ts.iter().next().expect("one target"), + source, + *targets.iter().next().unwrap(), spec.max_hops, - ), - None => search(adjacency, s, ts, spec.max_hops), - }; - (s, result) - }) - .collect(); + ) + .await + } + None => search(adjacency, source, targets, spec.max_hops).await, + }; + searches.insert(source, result); + } let mut rows = Vec::new(); let mut lengths = Vec::new(); let mut nodes = ListBuilder::new(Int64Builder::new()); @@ -580,7 +602,7 @@ fn search_batch( if d < spec.min_hops || d > spec.max_hops { continue; } - for path in paths_to(t, parents, spec.all) { + for path in paths_to(t, parents, spec.all).await { rows.push(i as u32); lengths.push(d as i64); nodes.values().append_slice(&path); @@ -628,7 +650,7 @@ fn plan_error(message: &str) -> GraphError { mod tests { use super::*; - fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { + async fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { let schema = Arc::new(Schema::new(vec![ Field::new("src", DataType::Int64, false), Field::new("dst", DataType::Int64, false), @@ -652,40 +674,33 @@ mod tests { futures::stream::once(async move { Ok(batch) }), )) as SendableRecordBatchStream }; - let rt = tokio::runtime::Builder::new_current_thread() - .build() - .unwrap(); - let out = rt - .block_on(Adjacency::build(stream(), "src", "dst")) - .unwrap(); - let inc = rt - .block_on(Adjacency::build(stream(), "dst", "src")) - .unwrap(); + let out = Adjacency::build(stream(), "src", "dst").await.unwrap(); + let inc = Adjacency::build(stream(), "dst", "src").await.unwrap(); (Arc::new(out), Arc::new(inc)) } - fn neighbors(edges: &[(i64, i64)]) -> Neighbors { - let (out, inc) = adjacencies(edges); + async fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + let (out, inc) = adjacencies(edges).await; Neighbors::undirected(out, inc) } - #[test] - fn finds_every_shortest_path_and_only_shortest_ones() { + #[tokio::test] + async fn finds_every_shortest_path_and_only_shortest_ones() { // Two 2-hop routes 1-2-4 and 1-3-4, and a longer 1-5-6-4. - let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]); - let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10); + let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]).await; + let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10).await; assert_eq!(depth[&4], 2); - let mut all = paths_to(4, &parents, true); + let mut all = paths_to(4, &parents, true).await; all.sort(); assert_eq!(all, vec![vec![1, 2, 4], vec![1, 3, 4]]); - assert_eq!(paths_to(4, &parents, false).len(), 1); + assert_eq!(paths_to(4, &parents, false).await.len(), 1); // Beyond the hop limit, nothing is reached. - let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1); + let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1).await; assert!(!depth.contains_key(&4)); } - #[test] - fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { + #[tokio::test] + async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { // A wide first layer makes the target's side expand. Two meeting nodes, // including parallel edges on both sides of the cut, yield six paths. let (out, inc) = adjacencies(&[ @@ -700,11 +715,12 @@ mod tests { (4, 5), (5, 5), (5, 1), - ]); + ]) + .await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3); + let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3).await; assert_eq!(depth[&5], 3); - let mut all = paths_to(5, &parents, true); + let mut all = paths_to(5, &parents, true).await; all.sort(); assert_eq!( all, @@ -713,16 +729,17 @@ mod tests { .chain(vec![vec![1, 3, 4, 5]; 2]) .collect::>() ); - let one = paths_to(5, &parents, false); + let one = paths_to(5, &parents, false).await; assert_eq!(one.len(), 1); assert!(all.contains(&one[0])); assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2) + .await .0 .contains_key(&5)); } - #[test] - fn bidirectional_matches_forward_search_for_small_graphs() { + #[tokio::test] + async fn bidirectional_matches_forward_search_for_small_graphs() { // All directed graphs on three nodes, including disconnected pairs and cycles. let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]; for mask in 0..(1 << edges.len()) { @@ -731,7 +748,7 @@ mod tests { .enumerate() .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge)) .collect(); - let (out, inc) = adjacencies(&selected); + let (out, inc) = adjacencies(&selected).await; let directed = Neighbors::directed(out.clone()); let reverse = Neighbors::directed(inc.clone()); let undirected = Neighbors::undirected(out, inc); @@ -740,17 +757,18 @@ mod tests { for target in 0..3 { for max_hops in 0..=3 { let (expected_depth, expected_parents) = - search(forward, source, &HashSet::from([target]), max_hops); + search(forward, source, &HashSet::from([target]), max_hops).await; let (depth, parents) = - search_bidirectional(forward, reverse, source, target, max_hops); + search_bidirectional(forward, reverse, source, target, max_hops) + .await; assert_eq!( depth.get(&target), expected_depth.get(&target), "graph {mask}, {source}->{target}, max {max_hops}" ); if depth.contains_key(&target) { - let mut expected = paths_to(target, &expected_parents, true); - let mut actual = paths_to(target, &parents, true); + let mut expected = paths_to(target, &expected_parents, true).await; + let mut actual = paths_to(target, &parents, true).await; expected.sort(); actual.sort(); assert_eq!( @@ -765,8 +783,8 @@ mod tests { } } - #[test] - fn bidirectional_avoids_expanding_the_wide_middle_layer() { + #[tokio::test] + async fn bidirectional_avoids_expanding_the_wide_middle_layer() { let mut edges = Vec::new(); for branch in 1..=100 { edges.push((0, branch)); @@ -774,13 +792,13 @@ mod tests { edges.push((branch, 101 + (branch - 1) * 100 + leaf)); } } - let (out, inc) = adjacencies(&edges); + let (out, inc) = adjacencies(&edges).await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2); - assert_eq!(paths_to(101, &parents, true), vec![vec![0, 1, 101]]); + let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2).await; + assert_eq!(paths_to(101, &parents, true).await, vec![vec![0, 1, 101]]); assert_eq!(depth.len(), 102); // Root, first layer and stitched target. assert_eq!( - search(&forward, 0, &HashSet::from([101]), 2).0.len(), + search(&forward, 0, &HashSet::from([101]), 2).await.0.len(), 10_101 ); } diff --git a/docs/api-reference.md b/docs/api-reference.md index 4972240..e493575 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -162,6 +162,12 @@ be opened. --- +Cypher named paths can be returned as Arrow structs containing `length`, `nodes`, `node_records` and +`relationships`. They retain their values through `collect()` and `UNWIND`; `nodes(path)` returns integer +keys, and `relationships(path)` returns typed relationship records. `startNode(r)` and `endNode(r)` return +the stored endpoint node records. Their occurrence identities are query-local. See +[path values and their memory cost](cypher.md#path-values). + ## `CypherQuery` diff --git a/docs/csr.md b/docs/csr.md index a9677e7..97e2244 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -132,6 +132,14 @@ When those hold, it covers: Everything else in the query (filters, aggregations, sorts and any hop the CSR can't take) is planned by DataFusion as usual. You can mix: a three-hop query might run two hops on the CSR and one as a join. +Named path values and `relationships(path)` use a query-local inventory of actual edge occurrences because +CSR neighbors do not carry edge identity or properties. This also preserves parallel edges when a path +passes through `collect()` and `UNWIND`. The inventory materializes the participating node and relationship +tables; bounded traversal batches do not bound that inventory. See [path values](cypher.md#path-values). +Shortest-path queries within one node label that only read length or node keys keep the CSR route below. +Relationships crossing labels use typed identities, even for these queries, to distinguish equal keys +from different labels. + For `shortestPath` and `allShortestPaths`, a source with one target in an input batch searches from both ends using the outgoing and incoming CSR. It expands the smaller frontier and finishes the whole layer before joining paths, so it keeps all shortest alternatives. Several targets for the same source share diff --git a/docs/cypher.md b/docs/cypher.md index 2334fb9..a13f4b9 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -181,8 +181,8 @@ implemented; `ExecutionStrategy.LanceNative` is reserved and raises an error. The Graph Data Council (GDC), formerly known as the Linked Data Benchmark Council (LDBC), maintains the [GDC SNB Interactive](https://ldbcouncil.org/benchmarks/snb/) complex queries. lance-csr implements a large -read-only subset of openCypher. The tables below specify its supported surface; relationship-path -operations needed by the unchanged IC14 query are not yet supported. +read-only subset of openCypher, including named path values and the relationship-path operations +used by IC14. The tables below specify its supported surface. ### Clauses @@ -234,7 +234,8 @@ RETURN length(p), nodes(p) expansion for the whole range; the join fallback uses one branch per length. - **Variable-length alternation** follows each relationship type's configured endpoint labels. With suitable CSR and node indexes, a mixed path such as `[:hasType|isSubclassOf*1..3]` stays in one operator. - Undirected or named variable-length alternation remains unsupported. + Named paths also support typed variable-length alternation through their edge inventory. + Unnamed undirected variable-length alternation remains unsupported. - **Undirected variable-length paths** skip immediate backtracking (`a-b-a`). - **`shortestPath` and `allShortestPaths`** need exactly one relationship type, no relationship variable, and named end nodes. Their length is unbounded unless you give a maximum. `nodes(p)` returns the node keys. @@ -243,7 +244,55 @@ RETURN length(p), nodes(p) duplicate node sequences that use different parallel edges. A single source-target pair searches from both ends when reverse adjacency is available. Several targets for the same source in an input batch share one forward search. -- **Named paths** (`p = ...`) are only supported with `shortestPath` and `allShortestPaths`. +- **Named paths** (`p = (a)-[:R]->(b)`) retain the ordered node and relationship records. Named + variable-length paths use each relationship at most once, including self-loops and parallel edges. + Use `relationships(p)` for variable-length relationship lists; binding a relationship variable inside + a named variable-length segment is not supported. + +### Path values + +A named path can pass through `WITH`, `collect()` and `UNWIND` without losing its edges. Parallel edges +with identical endpoints remain distinct, including when their properties are identical. + +```python +paths = engine.execute(""" + MATCH path = (a:Person {id: 1})-[:KNOWS]->(b:Person) + WITH collect(path) AS paths + UNWIND paths AS path + RETURN [n IN nodes(path) | n.name] AS names, + [r IN relationships(path) | endNode(r).id] AS targets + ORDER BY names +""") +print(paths.to_pylist()) +# [{'names': ['Alice', 'Bob'], 'targets': [2]}, +# {'names': ['Alice', 'Carol'], 'targets': [3]}] +``` + +`nodes(path)` still returns integer keys when returned directly. A property comprehension over it uses +stored node records when the path is a value, so equal keys in different labels remain different nodes. +`startNode(r)` and `endNode(r)` return the relationship's stored source and target, even when the path +traversed that relationship backwards. Property access, `keys()` and node `labels()` work on these values. +A zero-hop path has one node and an empty relationship list. All path functions return null for a null path. + +`RETURN path` produces an Arrow struct with `length`, `nodes`, `node_records` and `relationships` fields. +Entity records store user properties in `__lance_graph_properties` and `{label, id}` identity metadata +in `__lance_graph_identity`. Relationship records also carry the node records `__lance_graph_start` and +`__lance_graph_end`. These metadata fields are distinct from user properties, even if a property has the +same name. `keys(entity)` lists only non-null user properties; `keys(map)` includes null-valued keys. +Prefer the Cypher property and endpoint functions when writing queries. Relationship identities identify actual occurrences within this query's inventory, not persistent +IDs that can be compared across queries or refreshes. Rich paths support the same signed integer and +`uint8`/`uint16`/`uint32` keys as the CSR. Mixed labels may share compatible properties; conflicting Arrow +property types fail explicitly. Non-null node IDs must be unique within each participating label so +relationship endpoints identify one node; duplicate IDs raise an error. Expanding a graph pattern from +a mixed-label value remains unsupported. + +A path value needs a query-local inventory of the participating node and relationship tables, including +properties. It therefore materializes those tables, even with `use_csr=True`. Traversal output is emitted +in batches and named ordinary paths keep a cursor per hop, but inventory memory grows with the graph. +Queries asking only for shortest-path length, node keys or path null checks, including homogeneous node-property +comprehensions and paths carried through `WITH` aliases, retain the usual CSR execution route when +the relationship joins the same node label. Cross-label shortest paths use typed identities even for +length or node keys, so equal numeric keys from different labels cannot become a false zero-hop path. ### Expressions and predicates @@ -277,7 +326,7 @@ RETURN length(p), nodes(p) | List predicates | `all`, `any`, `none`, `single` | | Math | `abs`, `round`, `floor`, `ceil`, `sqrt` | | Temporal | `date`, `datetime`, `date_part('month', d)`, `to_epoch_ms`, `duration.inSeconds(a, b).seconds` | -| Paths | `length(p)`, `nodes(p)` on `shortestPath` variables | +| Paths | `length(p)`, `nodes(p)`, `relationships(p)`; `startNode(r)`, `endNode(r)` on relationship values | | Vectors | `vector_distance(a, b, metric)`, `vector_similarity(a, b, metric)`; see [Vector search](vector-search.md) | The string and list functions follow the [Cypher v9 reference](https://s3.amazonaws.com/artifacts.opencypher.org/openCypher9.pdf), @@ -340,7 +389,7 @@ print(result.to_pylist()) # [{'labels': ['Person'], 'properties': 3}] ``` Any other function raises `ValueError` with `Cypher function 'name' is not implemented`. This is a -bounded function set, not a claim of full Cypher conformance. `relationships(p)` is not supported. +bounded function set, not a claim of full Cypher conformance. ### Ordered reduce and list predicates diff --git a/python/python/tests/queries/NOTICE.txt b/python/python/tests/queries/NOTICE.txt new file mode 100644 index 0000000..69d3211 --- /dev/null +++ b/python/python/tests/queries/NOTICE.txt @@ -0,0 +1,13 @@ + Copyright [2020-]2025 Linked Data Benchmark Council + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/python/python/tests/queries/README.md b/python/python/tests/queries/README.md new file mode 100644 index 0000000..280f6db --- /dev/null +++ b/python/python/tests/queries/README.md @@ -0,0 +1,11 @@ +# Query fixtures + +`ic14.cypher` is the unchanged query string from `run_query14` in +[graph-benchmark-ldbc-complex](https://github.com/hyperdimensionalcomputingai/graph-benchmark-ldbc-complex/blob/c185123edd58b37b417bc4deb2fec643fdbd5e61/neo4j/query.py). +The fixture covers collected path values, relationship endpoints, correlated pattern +comprehensions and ordered folds using a small generated graph, including parallel edges. + +The source adapts the Apache-2.0 LDBC SNB Interactive reference implementation. +The upstream copyright and attribution notice is retained in `NOTICE.txt`; the +Apache-2.0 license is at the repository root. No benchmark dataset or expected +SF1 answers are distributed with this fixture. diff --git a/python/python/tests/queries/ic14.cypher b/python/python/tests/queries/ic14.cypher new file mode 100644 index 0000000..684741c --- /dev/null +++ b/python/python/tests/queries/ic14.cypher @@ -0,0 +1,26 @@ + + MATCH path = allShortestPaths((person1:Person {ID: $person1Id})-[:knows*0..]-(person2:Person {ID: $person2Id})) + WITH collect(path) AS paths + UNWIND paths AS path + WITH path, relationships(path) AS rels_in_path + WITH + [n IN nodes(path) | n.ID] AS personIdsInPath, + [r IN rels_in_path | + reduce(w = 0.0, v IN [ + (a:Person)<-[:commentHasCreator]-(:Comment)-[:replyOfPost]->(:Post)-[:postHasCreator]->(b:Person) + WHERE (a.ID = startNode(r).ID AND b.ID = endNode(r).ID) OR (a.ID = endNode(r).ID AND b.ID = startNode(r).ID) + | 1.0] | w + v) + ] AS weight1, + [r IN rels_in_path | + reduce(w = 0.0, v IN [ + (a:Person)<-[:commentHasCreator]-(:Comment)-[:replyOfComment]->(:Comment)-[:commentHasCreator]->(b:Person) + WHERE (a.ID = startNode(r).ID AND b.ID = endNode(r).ID) OR (a.ID = endNode(r).ID AND b.ID = startNode(r).ID) + | 0.5] | w + v) + ] AS weight2 + WITH + personIdsInPath, + reduce(w = 0.0, v IN weight1 | w + v) AS w1, + reduce(w = 0.0, v IN weight2 | w + v) AS w2 + RETURN personIdsInPath, (w1 + w2) AS pathWeight + ORDER BY pathWeight DESC + \ No newline at end of file diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 922a369..bece713 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -137,3 +137,32 @@ def test_shortest_paths_shares_a_source_across_targets(use_csr): (2, 2, [5, 1, 2]), (3, 2, [5, 1, 3]), ] + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize(("target", "hops", "missing"), [(3, 2, False), (6, -1, True)]) +@pytest.mark.parametrize( + "projection", + [ + "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END AS hops, " + "p IS NULL AS missing", + "RETURN CASE WHEN p IS NOT NULL THEN length(p) ELSE -1 END AS hops, " + "p IS NULL AS missing", + "WITH p AS q RETURN CASE q IS NULL WHEN true THEN -1 " + "ELSE length(q) END AS hops, q IS NULL AS missing", + "WITH p AS q WHERE q IS NOT NULL OR q IS NULL " + "RETURN coalesce(length(q), -1) AS hops, q IS NULL AS missing", + ], +) +def test_shortest_path_null_checks_keep_compact_plan( + use_csr, target, hops, missing, projection +): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) + query = ( + f"MATCH (a:Person {{id: 5}}), (b:Person {{id: {target}}}) " + "OPTIONAL MATCH p = shortestPath((a)-[:KNOWS*1..5]-(b)) " + projection + ) + assert engine.execute(query).to_pylist() == [{"hops": hops, "missing": missing}] + plan = engine.explain(query) + assert "ShortestPathExec" in plan + assert "PathMatchExec" not in plan diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py new file mode 100644 index 0000000..b5f3232 --- /dev/null +++ b/python/python/tests/test_path_values.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Actual edge records survive path values, including duplicate endpoint pairs.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["scan", "csr"]) +def engine(request): + people = pa.table({"id": [1, 2, 3, 4], "name": ["a", "b", "c", "d"]}) + edges = pa.table( + {"src": [1, 1, 2, 3, 3], "dst": [2, 2, 3, 3, 99], "weight": [5, 7, 11, 13, 17]} + ) + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, {"Person": people, "KNOWS": edges}, use_csr=request.param + ) + + +def test_named_path_records(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person)" + "-[:KNOWS]->(c:Person {id: 3}) " + "RETURN length(p) AS length, nodes(p) AS ids, " + "[n IN nodes(p) | n.name] AS names, " + "[r IN relationships(p) | r.weight] AS weights" + ).to_pylist() + assert sorted(rows, key=lambda r: r["weights"]) == [ + {"length": 2, "ids": [1, 2, 3], "names": ["a", "b", "c"], "weights": [5, 11]}, + {"length": 2, "ids": [1, 2, 3], "names": ["a", "b", "c"], "weights": [7, 11]}, + ] + + +@pytest.mark.parametrize("search", ["shortestPath", "allShortestPaths"]) +def test_shortest_path_survives_collect_unwind(engine, search): + rows = engine.execute( + f"MATCH p = {search}((a:Person {{id: 1}})-[:KNOWS*0..]->(b:Person {{id: 3}})) " + "WITH collect(p) AS paths UNWIND paths AS p " + "RETURN nodes(p) AS ids, [r IN relationships(p) | r.weight] AS weights, " + "[n IN nodes(p) | n.name] AS names" + ).to_pylist() + expected = [[5, 11], [7, 11]] + assert len(rows) == (2 if search == "allShortestPaths" else 1) + assert all(r["ids"] == [1, 2, 3] and r["names"] == ["a", "b", "c"] for r in rows) + assert all(r["weights"] in expected for r in rows) + if search == "allShortestPaths": + assert sorted(r["weights"] for r in rows) == expected + + +def test_reverse_path_keeps_stored_relationship_endpoints(engine): + rows = engine.execute( + "MATCH p = (b:Person {id: 2})<-[r:KNOWS]-(a:Person {id: 1}) " + "UNWIND relationships(p) AS edge RETURN nodes(p) AS ids, " + "edge.weight AS weight, startNode(edge).id AS source, " + "endNode(edge).id AS target" + ).to_pylist() + assert sorted(rows, key=lambda r: r["weight"]) == [ + {"ids": [2, 1], "weight": 5, "source": 1, "target": 2}, + {"ids": [2, 1], "weight": 7, "source": 1, "target": 2}, + ] + + +@pytest.mark.parametrize("shortest", [False, True]) +def test_zero_path_records(engine, shortest): + pattern = "(a:Person {id: 3})-[:KNOWS*0..0]->(b:Person {id: 3})" + if shortest: + pattern = f"shortestPath({pattern})" + rows = engine.execute( + f"MATCH p = {pattern} RETURN p, length(p) AS length, " + "nodes(p) AS ids, relationships(p) AS edges" + ).to_pylist() + assert len(rows) == 1 + assert rows[0]["length"] == 0 + assert rows[0]["ids"] == [3] + assert rows[0]["edges"] == [] + assert rows[0]["p"]["node_records"][0]["__lance_graph_properties"]["name"] == "c" + + +def test_named_trails_do_not_reuse_a_relationship(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 3})-[:KNOWS*1..3]->(b:Person) " + "RETURN nodes(p) AS ids, [r IN relationships(p) | r.weight] AS weights" + ).to_pylist() + assert rows == [{"ids": [3, 3], "weights": [13]}] + + +def test_optional_path_null_values(engine): + rows = engine.execute( + "MATCH (a:Person {id: 1}), (b:Person {id: 4}) " + "OPTIONAL MATCH p = shortestPath((a)-[:KNOWS*0..]->(b)) " + "RETURN p, relationships(p) AS edges, nodes(p) AS ids" + ).to_pylist() + assert rows == [{"p": None, "edges": None, "ids": None}] + + +def test_path_entities_have_properties_and_labels(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person) " + "UNWIND relationships(p) AS r " + "RETURN keys(r) AS properties, labels(startNode(r)) AS labels" + ).to_pylist() + assert rows == [{"properties": ["dst", "src", "weight"], "labels": ["Person"]}] * 2 + + +def test_fast_shortest_plan_retained(engine): + query = ( + "MATCH p = shortestPath((a:Person {id: 1})-[:KNOWS*0..]->(b:Person {id: 3})) " + "RETURN [n IN nodes(p) | n.id] AS ids" + ) + assert "PathMatch" not in engine.explain(query) + assert engine.execute(query).to_pylist() == [{"ids": [1, 2, 3]}] + + +def test_mixed_label_path_identity(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1], "name": ["person"]}), + "Place": pa.table({"id": [1], "name": ["place"]}), + "LIVES": pa.table({"src": [1], "dst": [1]}), + }, + ) + rows = engine.execute( + "MATCH p = (a:Person)-[:LIVES]->(b:Place) RETURN nodes(p) AS ids, " + "[n IN nodes(p) | n.name] AS names, [n IN nodes(p) | labels(n)] AS labels" + ).to_pylist() + assert rows == [ + {"ids": [1, 1], "names": ["person", "place"], "labels": [["Person"], ["Place"]]} + ] + + +def test_named_variable_alternation_tracks_labels(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .rel_table("INSIDE", "Place", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1], "name": ["person"]}), + "Place": pa.table({"id": [1, 2], "name": ["place", "country"]}), + "LIVES": pa.table({"src": [1], "dst": [1]}), + "INSIDE": pa.table({"src": [1], "dst": [2]}), + }, + use_csr=True, + ) + query = ( + "MATCH p = (a:Person)-[:LIVES|INSIDE*0..20]->(b:Place) " + "RETURN nodes(p) AS ids, [n IN nodes(p) | n.name] AS names ORDER BY ids" + ) + assert engine.execute(query).to_pylist() == [ + {"ids": [1, 1], "names": ["person", "place"]}, + {"ids": [1, 1, 2], "names": ["person", "place", "country"]}, + ] + assert engine.explain(query).count("PathMatch:") == 1 + + +def test_identical_parallel_occurrences_are_distinct(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2]}), + "R": pa.table({"src": [1, 1], "dst": [2, 2], "weight": [5, 5]}), + }, + ) + rows = engine.execute( + "MATCH p = allShortestPaths((a:Person {id: 1})" + "-[:R*]->(b:Person {id: 2})) RETURN p" + ).to_pylist() + relationships = [r["p"]["relationships"][0] for r in rows] + assert len(relationships) == 2 + assert len({r["__lance_graph_identity"]["id"] for r in relationships}) == 2 + assert [r["__lance_graph_properties"]["weight"] for r in relationships] == [5, 5] + + +@pytest.mark.parametrize("bad_type", [pa.string(), pa.uint64()]) +def test_rich_path_rejects_unsupported_ids(bad_type): + values = ["1", "2"] if pa.types.is_string(bad_type) else [1, 2] + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": pa.array(values, type=bad_type)}), + "R": pa.table( + { + "src": pa.array(values[:1], type=bad_type), + "dst": pa.array(values[1:], type=bad_type), + } + ), + }, + ) + with pytest.raises(NotImplementedError, match="unsupported type"): + engine.execute("MATCH p = (a:Person)-[:R]->(b:Person) RETURN p") + + +def test_mixed_path_rejects_conflicting_property_types(): + config = ( + GraphConfig.builder() + .node_table("A", "id") + .node_table("B", "id") + .rel_table("R", "A", "B", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "A": pa.table({"id": [1], "name": ["name"]}), + "B": pa.table({"id": [1], "name": [99]}), + "R": pa.table({"src": [1], "dst": [1]}), + }, + ) + with pytest.raises(ValueError, match="conflicting Arrow types"): + engine.execute("MATCH p = (a:A)-[:R]->(b:B) RETURN p") + + +def test_path_inventory_keeps_lance_snapshot(tmp_path): + import lance + + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + people_path, edges_path = ( + str(tmp_path / "people.lance"), + str(tmp_path / "edges.lance"), + ) + lance.write_dataset( + pa.table({"id": [1, 2], "name": ["old", "target"]}), people_path + ) + lance.write_dataset(pa.table({"src": [1], "dst": [2], "weight": [5]}), edges_path) + engine = CypherEngine( + config, + {"Person": lance.dataset(people_path), "R": lance.dataset(edges_path)}, + use_csr=True, + ) + query = ( + "MATCH p = (a:Person {id: 1})-[:R]->(b:Person) " + "RETURN [n IN nodes(p) | n.name] AS names, " + "[r IN relationships(p) | r.weight] AS weights" + ) + reader = engine.execute_batches(query) + lance.write_dataset( + pa.table({"src": [1], "dst": [2], "weight": [7]}), edges_path, mode="overwrite" + ) + engine.refresh() + assert reader.read_all().to_pylist() == [ + {"names": ["old", "target"], "weights": [5]} + ] + assert engine.execute(query).to_pylist() == [ + {"names": ["old", "target"], "weights": [7]} + ] + + +def test_bound_relationship_endpoints_survive_collect(engine): + rows = engine.execute( + "MATCH (a:Person {id: 1})-[r:KNOWS]->(b:Person {id: 2}) " + "WITH collect(r) AS rs UNWIND rs AS edge " + "RETURN edge.weight AS weight, keys(edge) AS properties, " + "startNode(edge).name AS source, endNode(edge).name AS target, " + "labels(endNode(edge)) AS labels ORDER BY weight" + ).to_pylist() + assert rows == [ + { + "weight": weight, + "properties": ["dst", "src", "weight"], + "source": "a", + "target": "b", + "labels": ["Person"], + } + for weight in [5, 7] + ] + + +def test_entity_metadata_does_not_shadow_properties(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table( + {"id": [1, 2], "nullable": pa.array([None, None], pa.string())} + ), + "R": pa.table({"src": [1], "dst": [2], "__lance_graph_identity": [42]}), + }, + ) + rows = engine.execute( + "MATCH p = (a:Person)-[:R]->(b:Person) UNWIND relationships(p) AS r " + "RETURN r.__lance_graph_identity AS user_value, keys(r) AS properties, " + "keys(startNode(r)) AS node_properties, keys({nullable: null}) AS map_keys" + ).to_pylist() + assert rows == [ + { + "user_value": 42, + "properties": ["__lance_graph_identity", "dst", "src"], + "node_properties": ["id"], + "map_keys": ["nullable"], + } + ] + + +def test_multiple_path_operators_share_occurrence_identity(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person {id: 2}) " + "MATCH q = (a)-[:KNOWS]->(b) " + "RETURN head(relationships(p)) AS first, head(relationships(q)) AS second" + ).to_pylist() + assert len(rows) == 4 + ids = {} + for row in rows: + for edge in row.values(): + weight = edge["__lance_graph_properties"]["weight"] + identity = edge["__lance_graph_identity"] + assert ids.setdefault(weight, identity) == identity + assert ids[5] != ids[7] + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_unmodified_ic14_with_real_relationship_values(use_csr): + from pathlib import Path + + # Exact query text from the benchmark's Neo4j IC14, including collect(path), + # relationship endpoints, nested correlated pattern comprehensions, and reduce. + query = (Path(__file__).parent / "queries" / "ic14.cypher").read_text() + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Comment", "id") + .node_table("Post", "id") + .rel_table("knows", "Person", "Person", "src", "dst") + .rel_table("commentHasCreator", "Comment", "Person", "src", "dst") + .rel_table("postHasCreator", "Post", "Person", "src", "dst") + .rel_table("replyOfPost", "Comment", "Post", "src", "dst") + .rel_table("replyOfComment", "Comment", "Comment", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2, 3, 4]}), + "Comment": pa.table({"id": [10, 11, 12, 13]}), + "Post": pa.table({"id": [20, 21]}), + "knows": pa.table({"src": [1, 1, 2, 1, 4], "dst": [2, 2, 3, 4, 3]}), + "commentHasCreator": pa.table( + {"src": [10, 11, 12, 13], "dst": [1, 2, 2, 4]} + ), + "postHasCreator": pa.table({"src": [20, 21], "dst": [2, 3]}), + "replyOfPost": pa.table({"src": [10, 11, 13], "dst": [20, 21, 21]}), + "replyOfComment": pa.table({"src": [12], "dst": [10]}), + }, + use_csr=use_csr, + ) + rows = engine.execute(query, {"person1Id": 1, "person2Id": 3}).to_pylist() + assert [(r["personidsinpath"], r["pathweight"]) for r in rows] == [ + ([1, 2, 3], 2.5), + ([1, 2, 3], 2.5), + ([1, 4, 3], 1.0), + ] + + +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +@pytest.mark.parametrize("shortest", [False, True]) +def test_rich_inventory_reads_every_partition(tmp_path, storage, shortest): + import lance + + n = 20_000 + people = pa.table({"id": pa.array(range(n + 1), pa.int64())}) + edges = pa.table( + { + "src": pa.array(range(n), pa.int64()), + "dst": pa.array(range(1, n + 1), pa.int64()), + } + ) + datasets = {"Person": people, "R": edges} + if storage == "lance": + datasets = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in datasets.items() + } + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine(config, datasets, use_csr=True) + pattern = f"(a:Person {{id: {n - 10}}})-[:R*1..5]->(b:Person {{id: {n - 5}}})" + if shortest: + pattern = f"allShortestPaths({pattern})" + rows = engine.execute(f"MATCH p = {pattern} RETURN p").to_pylist() + assert len(rows) == 1 + assert rows[0]["p"]["nodes"] == list(range(n - 10, n - 4)) + assert len(rows[0]["p"]["relationships"]) == 5 + + +@pytest.mark.parametrize( + ("suffix", "rich"), + [ + ("WITH p RETURN nodes(p) AS ids", False), + ("WITH p AS q WITH q AS r RETURN nodes(r) AS ids", False), + ("WITH p AS q RETURN [n IN nodes(q) | n.id] AS ids", False), + ( + "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r) AS ids", + True, + ), + ], +) +def test_path_pass_through_only_materializes_consumed_values(engine, suffix, rich): + query = ( + "MATCH p = allShortestPaths((a:Person {id: 1})-[:KNOWS*]->(b:Person {id: 3})) " + + suffix + ) + assert engine.execute(query).to_pylist() == [{"ids": [1, 2, 3]}] * 2 + assert ("PathMatchExec" in engine.explain(query)) is rich + + +@pytest.mark.parametrize( + "suffix", ["RETURN relationships(q) AS edges", "RETURN q AS path"] +) +def test_path_alias_retains_actual_relationship_values(engine, suffix): + query = ( + "MATCH p = allShortestPaths((a:Person {id: 1})-[:KNOWS*]->(b:Person {id: 3})) " + "WITH p AS q " + suffix + ) + rows = engine.execute(query).to_pylist() + edges = [row.get("edges", row.get("path", {}).get("relationships")) for row in rows] + assert sorted( + [r["__lance_graph_properties"]["weight"] for r in path] for path in edges + ) == [ + [5, 11], + [7, 11], + ] + assert "PathMatchExec" in engine.explain(query) + + +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +def test_rich_path_rejects_ambiguous_node_ids(tmp_path, storage): + import lance + + tables = { + "Person": pa.table({"id": [1, 1, 2], "name": ["one", "duplicate", "two"]}), + "R": pa.table({"src": [1], "dst": [2]}), + } + if storage == "lance": + tables = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in tables.items() + } + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine(config, tables, use_csr=True) + with pytest.raises( + Exception, match="unique node IDs within label 'Person'; duplicate ID 1" + ): + engine.execute("MATCH p = (a:Person)-[:R]->(b:Person) RETURN p") + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("search", ["shortestPath", "allShortestPaths"]) +@pytest.mark.parametrize("minimum", [0, 1]) +def test_shortest_scalar_consumers_preserve_cross_label_identity( + use_csr, search, minimum +): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2], "name": ["one", "two"]}), + "Place": pa.table({"id": [1], "name": ["place"]}), + "LIVES": pa.table({"src": [1, 2], "dst": [1, 1]}), + }, + use_csr=use_csr, + ) + query = ( + f"MATCH p = {search}((a:Person {{id: 1}})-[:LIVES*{minimum}..]->(b:Place)) " + "WITH p AS q RETURN length(q) AS hops, nodes(q) AS ids, " + "[n IN nodes(q) | n.name] AS names, [n IN nodes(q) | labels(n)] AS labels" + ) + assert engine.execute(query).to_pylist() == [ + { + "hops": 1, + "ids": [1, 1], + "names": ["one", "place"], + "labels": [["Person"], ["Place"]], + } + ] + assert "PathMatchExec" in engine.explain(query) + # Equal endpoint labels do not make Place 1 interchangeable with Person 1. + rows = engine.execute( + f"MATCH p = {search}((a:Person {{id: 1}})" + f"-[:LIVES*{minimum}..]-(b:Person {{id: 2}})) " + "RETURN length(p) AS hops, nodes(p) AS ids, [n IN nodes(p) | n.name] AS names" + ).to_pylist() + assert rows == [{"hops": 2, "ids": [1, 1, 2], "names": ["one", "place", "two"]}]