From 1bf6d42141b9cc6ee1824d000f49cdc0ebd9d8fe Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:25:46 -0700 Subject: [PATCH 1/2] perf(csr): traverse variable-length paths iteratively Resolve #13 with bounded typed traversal over CSR snapshots, retaining mixed relationship types without unrolling every length or type sequence. Implement zero-hop identity in indexed and join paths while preserving the 20-hop limit. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/csr_expand.rs | 103 +++++- crates/lance-graph/src/csr_traversal.rs | 328 ++++++++++++++++++ .../src/datafusion_planner/analysis.rs | 12 +- .../datafusion_planner/builder/expand_ops.rs | 217 +++++++++++- .../src/datafusion_planner/builder/helpers.rs | 7 +- .../src/datafusion_planner/join_ops.rs | 2 +- .../src/datafusion_planner/scan_ops.rs | 1 + crates/lance-graph/src/lib.rs | 1 + crates/lance-graph/src/logical_plan.rs | 78 ++++- .../lance-graph/src/pattern_alternatives.rs | 132 ++++--- crates/lance-graph/src/query.rs | 4 +- docs/csr.md | 20 +- docs/cypher.md | 9 +- .../python/tests/test_engine_alternation.py | 45 +++ python/python/tests/test_engine_csr.py | 32 +- 15 files changed, 912 insertions(+), 79 deletions(-) create mode 100644 crates/lance-graph/src/csr_traversal.rs diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 201a3cc..9bb4c20 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -554,7 +554,7 @@ impl NodeRows { self.positions.len() } - fn position(&self, id: i64) -> Option { + pub(crate) fn position(&self, id: i64) -> Option { self.positions.get(&id).copied() } @@ -773,7 +773,7 @@ pub struct Neighbors { /// Occurrence offsets are sufficient to pause even a high-degree source row. #[derive(Default)] -struct NeighborPosition { +pub(crate) struct NeighborPosition { forward: usize, reverse: usize, } @@ -798,7 +798,7 @@ impl Neighbors { .for_each(f); } - fn neighbors_from<'a>( + pub(crate) fn neighbors_from<'a>( &'a self, key: i64, position: &'a mut NeighborPosition, @@ -859,6 +859,7 @@ pub struct CsrExpandNode { /// Closing edge: input column holding the already-bound target id. Only neighbours /// equal to it are kept, replacing a filter after a full expansion. closing_column: Option, + traversal: Option>, schema: DFSchemaRef, } @@ -888,6 +889,7 @@ impl CsrExpandNode { target, target_filter: None, closing_column: None, + traversal: None, schema: Arc::new(schema), }) } @@ -898,6 +900,11 @@ impl CsrExpandNode { self } + pub(crate) fn with_traversal(mut self, traversal: crate::csr_traversal::Traversal) -> Self { + self.traversal = Some(Arc::new(traversal)); + self + } + /// The same node over `input` with `target`, keeping the moved filter and closing column. fn rebuilt(&self, input: LogicalPlan, target: Option) -> Result { let mut node = Self::try_new( @@ -910,6 +917,7 @@ impl CsrExpandNode { )?; node.target_filter = self.target_filter.clone(); node.closing_column = self.closing_column.clone(); + node.traversal = self.traversal.clone(); Ok(node) } @@ -965,6 +973,11 @@ impl PartialEq for CsrExpandNode { } && self.target_filter == other.target_filter && self.closing_column == other.closing_column + && match (&self.traversal, &other.traversal) { + (Some(a), Some(b)) => Arc::ptr_eq(a, b), + (None, None) => true, + _ => false, + } } } impl Eq for CsrExpandNode {} @@ -1067,6 +1080,9 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { if let Some(column) = &self.closing_column { write!(f, ", match={}", column)?; } + if let Some(walk) = &self.traversal { + write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; + } Ok(()) } @@ -1091,6 +1107,7 @@ pub struct CsrExpandExec { /// The moved target filter, if any. target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, properties: Arc, } @@ -1104,6 +1121,7 @@ impl CsrExpandExec { target: Option, target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, ) -> Self { let properties = Arc::new(PlanProperties::new( @@ -1119,6 +1137,7 @@ impl CsrExpandExec { target, target_filter, closing_column, + traversal, schema, properties, } @@ -1132,6 +1151,7 @@ struct BatchExpansion { bound: Option, row: usize, position: NeighborPosition, + walk: Option, } impl BatchExpansion { @@ -1151,6 +1171,7 @@ impl BatchExpansion { bound, row: 0, position: NeighborPosition::default(), + walk: None, }) } @@ -1192,6 +1213,49 @@ impl BatchExpansion { } chunk } + + async fn next_path_chunk( + &mut self, + graph: &crate::csr_traversal::Traversal, + target: Option<&TargetLookup>, + batch_size: usize, + ) -> ExpandedRows { + let mut chunk = ExpandedRows::default(); + while self.row < self.keys.len() { + let i = self.row; + if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { + let key = self.keys.value(i); + let walk = self + .walk + .get_or_insert_with(|| crate::csr_traversal::TraversalCursor::new(graph, key)); + let mut candidates = 0usize; + while let Some(neighbor) = walk.next(graph).await { + candidates += 1; + if candidates.is_multiple_of(1024) { + cooperative_yield().await; + } + if self.bound.as_ref().is_some_and(|b| b.value(i) != neighbor) { + continue; + } + if let Some(target) = target { + let Some(position) = target.rows.position(neighbor) else { + continue; + }; + chunk.target_positions.push(position); + } + chunk.rows.push(i as u32); + chunk.key_values.push(key); + chunk.neighbor_values.push(neighbor); + if chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + self.row += 1; + self.walk = None; + } + chunk + } } /// Expanded rows of one input batch: input row index, both endpoints, and the @@ -1285,6 +1349,20 @@ fn retain_by_mask(values: &mut Vec, keep: &[bool]) { values.retain(|_| *keep.next().expect("one mask entry per value")); } +/// Cooperate with the stream's executor without requiring a particular async runtime. +pub(crate) async fn cooperative_yield() { + let mut yielded = false; + std::future::poll_fn(move |cx| { + if std::mem::replace(&mut yielded, true) { + std::task::Poll::Ready(()) + } else { + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await +} + impl DisplayAs for CsrExpandExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!( @@ -1305,6 +1383,9 @@ impl DisplayAs for CsrExpandExec { if let Some(column) = &self.closing_column { write!(f, ", match={}", column)?; } + if let Some(walk) = &self.traversal { + write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; + } Ok(()) } } @@ -1323,6 +1404,9 @@ impl ExecutionPlan for CsrExpandExec { use datafusion::common::stats::Precision; let input = self.input.partition_statistics(partition)?; let mut stats = datafusion::common::Statistics::new_unknown(&self.schema); + if self.traversal.is_some() { + return Ok(Arc::new(stats)); + } let degree = self.adjacency.mean_degree(); stats.num_rows = match input.num_rows { Precision::Exact(n) | Precision::Inexact(n) => { @@ -1356,6 +1440,7 @@ impl ExecutionPlan for CsrExpandExec { self.target.clone(), self.target_filter.clone(), self.closing_column.clone(), + self.traversal.clone(), self.schema.clone(), ))) } @@ -1379,6 +1464,7 @@ impl ExecutionPlan for CsrExpandExec { let schema = self.schema.clone(); let out_schema = schema.clone(); let batch_size = context.session_config().batch_size(); + let traversal = self.traversal.clone(); // Keep one input and one materialization in flight per partition. A hub's // remaining neighbours stay in the immutable adjacency until the next poll. let stream = futures::stream::try_unfold((input, None), move |(mut input, mut pending)| { @@ -1388,6 +1474,7 @@ impl ExecutionPlan for CsrExpandExec { target_filter.clone(), schema.clone(), ); + let traversal = traversal.clone(); async move { loop { if pending.is_none() { @@ -1397,7 +1484,14 @@ impl ExecutionPlan for CsrExpandExec { pending = Some(BatchExpansion::try_new(batch?, key_index, closing_index)?); } let expansion = pending.as_mut().expect("input batch just set"); - let chunk = expansion.next_chunk(&adjacency, target.as_deref(), batch_size); + let chunk = match &traversal { + Some(walk) => { + expansion + .next_path_chunk(walk, target.as_deref(), batch_size) + .await + } + None => expansion.next_chunk(&adjacency, target.as_deref(), batch_size), + }; if chunk.rows.is_empty() { pending = None; continue; @@ -1455,6 +1549,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { node.target.clone(), target_filter, node.closing_column.clone(), + node.traversal.clone(), Arc::new(node.schema.as_arrow().clone()), )))) } diff --git a/crates/lance-graph/src/csr_traversal.rs b/crates/lance-graph/src/csr_traversal.rs new file mode 100644 index 0000000..a19eb36 --- /dev/null +++ b/crates/lance-graph/src/csr_traversal.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Resumable variable-length walks over immutable, typed CSR snapshots. + +use std::sync::Arc; + +use crate::ast::RelationshipDirection; +use crate::config::GraphConfig; +use crate::csr_expand::{ + cooperative_yield, AdjacencyCatalog, NeighborPosition, Neighbors, NodeRows, +}; + +#[derive(Debug)] +struct Step { + from: usize, + to: usize, + adjacency: Neighbors, +} + +#[derive(Debug)] +pub(crate) struct Traversal { + steps: Vec, + nodes: Vec>, + start: usize, + target: usize, + pub min_hops: u32, + pub max_hops: u32, + undirected: bool, +} + +impl Traversal { + /// Every intermediate label needs a node index: joins drop dangling edges at every hop. + #[allow(clippy::too_many_arguments)] + pub fn new( + config: &GraphConfig, + catalog: &AdjacencyCatalog, + types: &[String], + direction: &RelationshipDirection, + source_label: &str, + target_label: &str, + min_hops: u32, + max_hops: u32, + ) -> Option { + let mut labels: Vec = Vec::new(); + let mut nodes = Vec::new(); + let mut label = |name: &str| -> Option { + if let Some(index) = labels.iter().position(|l| l.eq_ignore_ascii_case(name)) { + return Some(index); + } + nodes.push(catalog.nodes(name)?.clone()); + labels.push(name.to_owned()); + Some(labels.len() - 1) + }; + let (start, target) = (label(source_label)?, label(target_label)?); + let mut steps = Vec::new(); + for name in types { + let mapping = config.get_relationship_mapping(name)?; + let (from, to) = ( + label(mapping.source_label.as_deref()?)?, + label(mapping.target_label.as_deref()?)?, + ); + match direction { + RelationshipDirection::Outgoing => steps.push(Step { + from, + to, + adjacency: Neighbors::directed(catalog.get(name, true)?.clone()), + }), + RelationshipDirection::Incoming => steps.push(Step { + from: to, + to: from, + adjacency: Neighbors::directed(catalog.get(name, false)?.clone()), + }), + RelationshipDirection::Undirected if from == to => steps.push(Step { + from, + to, + adjacency: Neighbors::undirected( + catalog.get(name, true)?.clone(), + catalog.get(name, false)?.clone(), + ), + }), + RelationshipDirection::Undirected => { + steps.push(Step { + from, + to, + adjacency: Neighbors::directed(catalog.get(name, true)?.clone()), + }); + steps.push(Step { + from: to, + to: from, + adjacency: Neighbors::directed(catalog.get(name, false)?.clone()), + }); + } + } + } + (!steps.is_empty()).then_some(Self { + steps, + nodes, + start, + target, + min_hops, + max_hops, + undirected: *direction == RelationshipDirection::Undirected, + }) + } +} + +struct Frame { + label: usize, + key: i64, + step: usize, + position: NeighborPosition, + emit: bool, +} + +impl Frame { + fn new(label: usize, key: i64) -> Self { + Self { + label, + key, + step: 0, + position: NeighborPosition::default(), + emit: true, + } + } +} + +pub(crate) struct TraversalCursor { + stack: Vec, +} + +impl TraversalCursor { + pub fn new(graph: &Traversal, key: i64) -> Self { + Self { + stack: vec![Frame::new(graph.start, key)], + } + } + + /// Keep one frame per hop, and cooperate even when candidate edges produce no rows. + pub async fn next(&mut self, graph: &Traversal) -> Option { + let mut work = 0; + while !self.stack.is_empty() { + if work == 1024 { + cooperative_yield().await; + work = 0; + } + work += 1; + let depth = self.stack.len() as u32 - 1; + let frame = self.stack.last_mut().expect("nonempty stack"); + if frame.emit { + frame.emit = false; + if depth >= graph.min_hops && frame.label == graph.target { + return Some(frame.key); + } + } + if depth == graph.max_hops || frame.step == graph.steps.len() { + self.stack.pop(); + continue; + } + let step = &graph.steps[frame.step]; + let next = (step.from == frame.label) + .then(|| { + step.adjacency + .neighbors_from(frame.key, &mut frame.position) + .next() + }) + .flatten(); + let Some(key) = next.map(|n| n as i64) else { + frame.step += 1; + frame.position = NeighborPosition::default(); + continue; + }; + if graph.nodes[step.to].position(key).is_none() { + continue; + } + // Preserve the existing undirected immediate-backtracking rule, using typed ids. + // shortcut: full relationship uniqueness needs edge row ids, as on the join path. + if graph.undirected + && self.stack.len() >= 2 + && self.stack[self.stack.len() - 2].label == step.to + && self.stack[self.stack.len() - 2].key == key + { + continue; + } + self.stack.push(Frame::new(step.to, key)); + } + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::csr_expand::Adjacency; + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + + async fn graph() -> (GraphConfig, AdjacencyCatalog) { + let config = GraphConfig::builder() + .node_table("Tag", "id") + .node_table("Class", "id") + .rel_table("hasType", "Tag", "Class", "src", "dst") + .rel_table("isSubclass", "Class", "Class", "src", "dst") + .build() + .unwrap(); + let mut catalog = AdjacencyCatalog::default(); + for (name, ids) in [("Tag", vec![1, 2]), ("Class", vec![1, 2, 3])] { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids))]).unwrap(); + catalog.insert_nodes( + name, + Arc::new(NodeRows::from_batch(batch, "id").unwrap().unwrap()), + ); + } + for (name, src, dst) in [ + ("hasType", vec![1, 1, 2, 1], vec![1, 1, 2, 99]), + ("isSubclass", vec![1, 2, 3], vec![2, 3, 1]), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::Int64, false), + Field::new("dst", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(src)), + Arc::new(Int64Array::from(dst)), + ], + ) + .unwrap(); + for (outgoing, from, to) in [(true, "src", "dst"), (false, "dst", "src")] { + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::iter([Ok(batch.clone())]), + )); + catalog.insert( + name, + outgoing, + Arc::new(Adjacency::build(stream, from, to).await.unwrap()), + ); + } + } + (config, catalog) + } + + #[tokio::test] + async fn typed_walk_keeps_duplicates_and_skips_missing_intermediate_nodes() { + let (config, catalog) = graph().await; + let walk = Traversal::new( + &config, + &catalog, + &["hasType".into(), "isSubclass".into()], + &RelationshipDirection::Outgoing, + "Tag", + "Class", + 0, + 3, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut out = Vec::new(); + while let Some(key) = cursor.next(&walk).await { + out.push(key); + assert!(cursor.stack.len() <= 4); + } + out.sort(); + // Tag 1 and Class 1 are different nodes; no extra zero-hop Class 1 result. + assert_eq!(out, vec![1, 1, 2, 2, 3, 3]); + assert!(cursor.stack.is_empty()); + } + + #[tokio::test] + async fn zero_hop_is_identity_and_depth_bound_survives_cycles() { + let (config, catalog) = graph().await; + let walk = Traversal::new( + &config, + &catalog, + &["isSubclass".into()], + &RelationshipDirection::Outgoing, + "Class", + "Class", + 0, + 20, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + for depth in 0..=20 { + assert_eq!(cursor.next(&walk).await, Some(1 + depth % 3)); + assert_eq!(cursor.stack.len(), depth as usize + 1); + } + assert_eq!(cursor.next(&walk).await, None); + let zero = Traversal { + max_hops: 0, + ..walk + }; + let mut cursor = TraversalCursor::new(&zero, 2); + assert_eq!(cursor.next(&zero).await, Some(2)); + assert_eq!(cursor.next(&zero).await, None); + } + + #[tokio::test] + async fn walk_yields_while_no_candidate_can_match_the_target_label() { + let (config, mut catalog) = graph().await; + let adjacency = catalog + .get("isSubclass", true) + .unwrap() + .with_changes(&vec![(1, 1); 4096], &[]); + catalog.insert("isSubclass", true, Arc::new(adjacency)); + let walk = Traversal::new( + &config, + &catalog, + &["isSubclass".into()], + &RelationshipDirection::Outgoing, + "Class", + "Tag", + 1, + 2, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut next = Box::pin(cursor.next(&walk)); + assert!(futures::poll!(&mut next).is_pending()); + drop(next); + assert!(cursor.stack.len() <= 3); + } +} diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 89abee0..11d1e52 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -144,6 +144,7 @@ fn analyze_operator( input, source_variable, target_variable, + target_label, relationship_types, direction, relationship_variable, @@ -154,13 +155,9 @@ fn analyze_operator( // Recursively analyze input first analyze_operator(input, analysis, rel_counter)?; - // Infer target variable's label from source variable - // For (a:Person)-[:KNOWS]->(b), b also gets label Person - if let Some(source_label) = analysis.var_to_label.get(source_variable).cloned() { - analysis - .var_to_label - .insert(target_variable.clone(), source_label); - } + analysis + .var_to_label + .insert(target_variable.clone(), target_label.clone()); // For variable-length paths, register multiple instances (one per hop) // We need to register instances for all possible hop counts @@ -371,6 +368,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".to_string(), target_variable: "b".to_string(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".to_string()], direction: RelationshipDirection::Outgoing, relationship_variable: None, diff --git a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs index 9b4ad83..26777d3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs @@ -406,11 +406,7 @@ impl DataFusionPlanner { })) } - /// Build variable-length path expansion using unrolling + UNION strategy - /// - /// For a query like: (a)-[:KNOWS*1..3]->(b) - /// This generates: - /// 1-hop plan UNION 2-hop plan UNION 3-hop plan + /// Use one iterative CSR expansion when indexed; otherwise union the fixed-length plans. #[allow(clippy::too_many_arguments)] pub(crate) fn build_variable_length_expand( &self, @@ -424,7 +420,7 @@ impl DataFusionPlanner { max_length: Option, target_properties: &HashMap, ) -> Result { - let min_hops = min_length.unwrap_or(1).max(1); + let min_hops = min_length.unwrap_or(1); let max_hops = max_length.unwrap_or(crate::MAX_VARIABLE_LENGTH_HOPS); // Validate range @@ -465,6 +461,30 @@ impl DataFusionPlanner { .map(|f| f.name().clone()), ); + if let Some(plan) = self.try_build_csr_traversal( + ctx, + input_plan.clone(), + source_variable, + target_variable, + relationship_types, + direction, + min_hops, + max_hops, + target_properties, + )? { + let projection: Vec<_> = plan + .schema() + .fields() + .iter() + .filter(|f| expected_columns.contains(f.name().as_str())) + .map(|f| col(f.name())) + .collect(); + return LogicalPlanBuilder::from(plan) + .project(projection) + .and_then(|p| p.build()) + .map_err(|e| self.plan_error("Failed to project variable-length CSR", e)); + } + // Generate a plan for each hop count and UNION them let mut plans = Vec::new(); @@ -544,6 +564,59 @@ impl DataFusionPlanner { hop_count: u32, target_properties: &HashMap, ) -> Result { + if hop_count == 0 { + let (source_label, source_map) = self.get_target_node_mapping(ctx, source_variable)?; + let (target_label, target_map) = self.get_target_node_mapping(ctx, target_variable)?; + let source = col(qualify_column(source_variable, &source_map.id_field)); + let target = qualify_column(target_variable, &target_map.id_field); + let same_label = source_label.eq_ignore_ascii_case(&target_label); + let mut builder = LogicalPlanBuilder::from(input_plan); + if builder + .schema() + .field_with_unqualified_name(&target) + .is_ok() + { + if !target_properties.is_empty() { + return Err(crate::error::GraphError::PlanError { + message: format!("Property filters on reused variable '{target_variable}' are not yet supported"), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + builder = builder + .filter( + source + .eq(col(&target)) + .and(datafusion::logical_expr::lit(same_label)), + ) + .map_err(|e| self.plan_error("Failed to filter zero-hop identity", e))?; + } else { + let target_source = + self.get_catalog()? + .node_source(&target_label) + .ok_or_else(|| crate::error::GraphError::PlanError { + message: format!("No node source for {target_label}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let scan = self.build_target_scan( + target_source, + &target_label, + target_variable, + target_properties, + )?; + builder = builder + .join_on( + scan, + datafusion::logical_expr::JoinType::Inner, + [source + .eq(col(&target)) + .and(datafusion::logical_expr::lit(same_label))], + ) + .map_err(|e| self.plan_error("Failed to join zero-hop identity", e))?; + } + return builder + .build() + .map_err(|e| self.plan_error("Failed to build zero-hop identity", e)); + } let mut current_plan = input_plan; let mut current_source = source_variable.to_string(); let mut path_nodes = vec![current_source.clone()]; @@ -604,6 +677,130 @@ impl DataFusionPlanner { .map_err(|e| self.plan_error("Failed to build undirected path", e)) } + #[allow(clippy::too_many_arguments)] + fn try_build_csr_traversal( + &self, + ctx: &mut PlanningContext, + input: LogicalPlan, + source_variable: &str, + target_variable: &str, + types: &[String], + direction: &RelationshipDirection, + min_hops: u32, + max_hops: u32, + target_properties: &HashMap, + ) -> Result> { + use crate::csr_expand::{CsrExpandNode, Neighbors}; + use datafusion::logical_expr::Extension; + if max_hops == 0 { + return Ok(None); + } + let Some(catalog) = &self.adjacency else { + return Ok(None); + }; + let (source_label, source_map) = self.get_target_node_mapping(ctx, source_variable)?; + let (target_label, target_map) = self.get_target_node_mapping(ctx, target_variable)?; + let Some(walk) = crate::csr_traversal::Traversal::new( + &self.config, + catalog, + types, + direction, + &source_label, + &target_label, + min_hops, + max_hops, + ) else { + return Ok(None); + }; + let rel_type = &types[0]; + let rel_map = self.get_relationship_mapping(rel_type)?; + let source_key = qualify_column(source_variable, &source_map.id_field); + if !input + .schema() + .field_with_unqualified_name(&source_key) + .is_ok_and(|f| crate::csr_expand::is_supported_id_type(f.data_type())) + { + return Ok(None); + } + let adjacency = match direction { + RelationshipDirection::Outgoing => Neighbors::directed( + catalog + .get(rel_type, true) + .expect("validated traversal") + .clone(), + ), + RelationshipDirection::Incoming => Neighbors::directed( + catalog + .get(rel_type, false) + .expect("validated traversal") + .clone(), + ), + RelationshipDirection::Undirected => Neighbors::undirected( + catalog + .get(rel_type, true) + .expect("validated traversal") + .clone(), + catalog + .get(rel_type, false) + .expect("validated traversal") + .clone(), + ), + }; + let cat = self.get_catalog()?; + let target = target_properties + .is_empty() + .then(|| { + self.csr_target_lookup(cat, &input, target_variable, &target_label, target_map) + }) + .flatten(); + let fused = target.is_some(); + let bound_target = qualify_column(target_variable, &target_map.id_field); + let closing = input + .schema() + .field_with_unqualified_name(&bound_target) + .is_ok(); + let alias = ctx.next_relationship_instance(rel_type)?.alias.clone(); + let output = |name| { + arrow_schema::Field::new( + qualify_column(&alias, name), + arrow_schema::DataType::Int64, + true, + ) + }; + let mut node = CsrExpandNode::try_new( + input, + source_key, + output(Self::get_source_join_key(direction, rel_map)), + output(Self::get_target_join_key(direction, rel_map)), + adjacency, + target, + )? + .with_traversal(walk); + if closing && target_properties.is_empty() { + node = node.with_closing_column(bound_target); + } + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(node), + }); + if fused || (closing && target_properties.is_empty()) { + return Ok(Some(plan)); + } + self.join_relationship_to_target( + LogicalPlanBuilder::from(plan), + cat, + &TargetJoinParams { + target_variable, + target_label: &target_label, + rel_qualifier: &alias, + node_map: target_map, + rel_map, + direction, + target_properties, + }, + ) + .map(Some) + } + /// Build a single-hop expansion on top of an existing plan (one hop of a /// variable-length path). Uses the same planning as a single hop, so a hop gets /// the CSR expansion whenever a plain `(a)-[:R]->(b)` would. @@ -712,6 +909,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: Some("r".into()), @@ -747,6 +945,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -782,6 +981,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -831,6 +1031,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -879,6 +1080,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -928,6 +1130,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -977,6 +1180,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -1026,6 +1230,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index f2b8333..32a438b 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -138,14 +138,12 @@ impl DataFusionPlanner { } LogicalOperator::VariableLengthExpand { input, - source_variable, target_variable, + target_label, .. } => { Self::collect_labels(input, labels); - if let Some(label) = labels.get(source_variable).cloned() { - labels.insert(target_variable.clone(), label); - } + labels.insert(target_variable.clone(), target_label.clone()); } LogicalOperator::Filter { input, .. } | LogicalOperator::Project { input, .. } @@ -303,6 +301,7 @@ mod tests { input: Box::new(scan), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: RelationshipDirection::Outgoing, relationship_variable: None, diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index da9ed1c..a67b0d4 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -43,7 +43,7 @@ pub(crate) struct TargetJoinParams<'a> { impl DataFusionPlanner { /// Build a target node scan with property filters and qualified column names - fn build_target_scan( + pub(crate) fn build_target_scan( &self, target_source: Arc, target_label: &str, diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 33434f8..5330648 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -358,6 +358,7 @@ mod tests { input: Box::new(scan), source_variable: "foo_bar".to_string(), // Will generate _temp_foo_bar_1 target_variable: "baz".to_string(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".to_string()], direction: crate::ast::RelationshipDirection::Outgoing, min_length: Some(2), diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index 37083e7..6d89ced 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -40,6 +40,7 @@ pub mod case_insensitive; pub mod config; pub mod csr_expand; pub mod csr_index; +mod csr_traversal; pub mod datafusion_planner; pub mod error; pub mod key_statistics; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index ed89e08..d2ede95 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -72,7 +72,7 @@ pub enum LogicalOperator { /// Variable-length path expansion (*1..2 syntax) /// /// Represents multi-hop relationship traversals: (source)-[rel*min..max]->(target) - /// This is implemented by unrolling into multiple fixed-length paths and unioning them + /// Uses iterative CSR traversal when indexed, otherwise fixed-length branches and a union. VariableLengthExpand { /// The input operator (typically a node scan) input: Box, @@ -80,6 +80,8 @@ pub enum LogicalOperator { source_variable: String, /// Variable name for the target node (reachable in min..max hops) target_variable: String, + /// The label of the final node, which may differ from intermediate nodes. + target_label: String, /// Types of relationships to traverse in each hop relationship_types: Vec, /// Direction of traversal for each hop @@ -210,6 +212,7 @@ pub struct LogicalPlanner<'a> { variables: HashMap, // variable -> label config: &'a GraphConfig, fan_out: Option<&'a FanOut>, + adjacency: Option<&'a crate::csr_expand::AdjacencyCatalog>, /// Mark columns created for EXISTS subqueries so far (names `__exists_`). exists_marks: usize, /// Aliases holding lists of node ids (`collect(node) AS friends`), with the label. @@ -230,6 +233,7 @@ impl<'a> LogicalPlanner<'a> { variables: HashMap::new(), config, fan_out: None, + adjacency: None, exists_marks: 0, node_lists: HashMap::new(), comprehensions: 0, @@ -244,6 +248,73 @@ impl<'a> LogicalPlanner<'a> { self } + pub(crate) fn with_adjacency( + mut self, + adjacency: Option<&'a crate::csr_expand::AdjacencyCatalog>, + ) -> Self { + self.adjacency = adjacency; + self + } + + fn retain_variable_alternation(&self, patterns: &[GraphPattern]) -> bool { + let Some(catalog) = self.adjacency else { + return false; + }; + let labels = |n: &NodePattern| -> Vec { + n.variable + .as_ref() + .and_then(|v| self.variables.get(v)) + .map(|l| vec![l.clone()]) + .unwrap_or_else(|| n.labels.clone()) + }; + patterns.iter().all(|p| { + let GraphPattern::Path(path) = p else { + return true; + }; + let mut previous = &path.start_node; + path.segments.iter().all(|segment| { + let start = previous; + previous = &segment.end_node; + let rel = &segment.relationship; + let Some(range) = &rel.length else { + return true; + }; + if rel.types.len() <= 1 { + return true; + } + if rel.variable.is_some() + || !rel.properties.is_empty() + || rel.direction == RelationshipDirection::Undirected + || (range.min == Some(1) && range.max == Some(1)) + { + return false; + } + let (starts, ends) = (labels(start), labels(&segment.end_node)); + let (Some(start), Some(end)) = (starts.first(), ends.first()) else { + return false; + }; + if starts + .iter() + .chain(&ends) + .any(|l| catalog.nodes(l).is_none()) + { + return false; + } + crate::csr_traversal::Traversal::new( + self.config, + catalog, + &rel.types, + &rel.direction, + start, + end, + 0, + 0, + ) + .is_some() + }) + }) + } + /// For a path that continues earlier rows, the path read from the other end when that /// is the better start: always when only the far end is bound (expansion must start /// from a bound node), and, with both ends bound, when the far end's first hop has the @@ -575,11 +646,13 @@ impl<'a> LogicalPlanner<'a> { ..match_clause.clone() }; let match_clause = &inferred; - if crate::pattern_alternatives::needs_expansion(&match_clause.patterns) { + let retain_variable = self.retain_variable_alternation(&match_clause.patterns); + if crate::pattern_alternatives::needs_expansion(&match_clause.patterns, retain_variable) { let branches = crate::pattern_alternatives::expand( &match_clause.patterns, self.config, &self.variables, + retain_variable, )?; return self.plan_match_branches(base, branches, filter); } @@ -1235,6 +1308,7 @@ impl<'a> LogicalPlanner<'a> { input: Box::new(plan), source_variable: current_src.clone(), target_variable: target_variable.clone(), + target_label: target_label.clone(), relationship_types: segment.relationship.types.clone(), direction: segment.relationship.direction.clone(), relationship_variable: segment.relationship.variable.clone(), diff --git a/crates/lance-graph/src/pattern_alternatives.rs b/crates/lance-graph/src/pattern_alternatives.rs index f2d1356..64b37df 100644 --- a/crates/lance-graph/src/pattern_alternatives.rs +++ b/crates/lance-graph/src/pattern_alternatives.rs @@ -31,6 +31,12 @@ use crate::error::{GraphError, Result}; /// Branch limit, so a pattern full of alternatives can't explode the plan. const MAX_BRANCHES: usize = 64; +fn zero_hop(rel: &RelationshipPattern) -> bool { + rel.length + .as_ref() + .is_some_and(|r| r.min == Some(0) && r.max == Some(0)) +} + /// `patterns` with labels filled in for unlabeled nodes that the configured endpoint labels /// pin down: `(friend)<-[:postHasCreator]-(post)` makes `post` a Post when /// `postHasCreator` goes from Post to Person. Labels flow along paths and between @@ -149,16 +155,17 @@ pub fn infer_labels( } /// True when any relationship has several types or any node several labels. -pub fn needs_expansion(patterns: &[GraphPattern]) -> bool { +pub fn needs_expansion(patterns: &[GraphPattern], retain_variable: bool) -> bool { let nodes_multi = |n: &NodePattern| n.labels.len() > 1; patterns.iter().any(|p| match p { GraphPattern::Node(n) => nodes_multi(n), GraphPattern::Path(path) => { nodes_multi(&path.start_node) - || path - .segments - .iter() - .any(|s| s.relationship.types.len() > 1 || nodes_multi(&s.end_node)) + || path.segments.iter().any(|s| { + (s.relationship.types.len() > 1 + && !(retain_variable && s.relationship.length.is_some())) + || nodes_multi(&s.end_node) + }) } }) } @@ -178,11 +185,16 @@ pub fn expand( patterns: &[GraphPattern], config: &GraphConfig, bound: &HashMap, + retain_variable: bool, ) -> Result>> { - if let Some(variants) = unroll_variable_length_alternation(patterns, config, bound)? { + if let Some(variants) = if retain_variable { + None + } else { + unroll_variable_length_alternation(patterns, config, bound)? + } { let mut branches = Vec::new(); for variant in variants { - branches.extend(expand(&variant, config, bound)?); + branches.extend(expand(&variant, config, bound, retain_variable)?); if branches.len() > MAX_BRANCHES { return Err(too_many()); } @@ -227,12 +239,14 @@ pub fn expand( for segment in &path.segments { let rel = &segment.relationship; let right = slot(&segment.end_node, &mut allowed); - rels.push(Rel { - left, - right, - direction: rel.direction.clone(), - types: rel.types.clone(), - }); + if !(zero_hop(rel) || retain_variable && rel.length.is_some()) { + rels.push(Rel { + left, + right, + direction: rel.direction.clone(), + types: rel.types.clone(), + }); + } left = right; } } @@ -338,7 +352,13 @@ pub fn expand( } }); if consistent { - branches.push(rewrite(patterns, &occurrence_slots, &labels, &types)); + branches.push(rewrite( + patterns, + &occurrence_slots, + &labels, + &types, + retain_variable, + )); } } if branches.is_empty() { @@ -414,7 +434,7 @@ fn unroll_variable_length_alternation( labels.is_empty() || labels.iter().any(|l| l.eq_ignore_ascii_case(label)) }; let range = rel.length.as_ref().expect("a variable-length segment"); - let min = range.min.unwrap_or(1).max(1); + let min = range.min.unwrap_or(1); let max = range.max.unwrap_or(crate::MAX_VARIABLE_LENGTH_HOPS); if max > crate::MAX_VARIABLE_LENGTH_HOPS { return Err(unsupported(&format!( @@ -449,43 +469,52 @@ fn unroll_variable_length_alternation( } } } - if sequences.is_empty() { + if sequences.is_empty() && min != 0 { return Err(plan_error( "no sequence of relationship types in the variable-length pattern chains its \ endpoint labels" .into(), )); } - let variants = sequences - .into_iter() - .map(|sequence| { - let last = sequence.len() - 1; - let hops = sequence - .into_iter() - .enumerate() - .map(|(h, (t, label))| PathSegment { - relationship: RelationshipPattern { - types: vec![t], - length: None, - ..rel.clone() - }, - end_node: if h == last { - path.segments[i].end_node.clone() - } else { - NodePattern { - variable: None, - labels: vec![label], - properties: HashMap::new(), - } - }, - }); - let mut unrolled = path.clone(); - unrolled.segments.splice(i..=i, hops); - let mut variant = patterns.to_vec(); - variant[p] = GraphPattern::Path(unrolled); - variant - }) - .collect(); + let mut variants = Vec::new(); + if min == 0 { + let mut identity = path.clone(); + identity.segments[i].relationship.types = vec![rel.types[0].clone()]; + identity.segments[i].relationship.length = Some(crate::ast::LengthRange { + min: Some(0), + max: Some(0), + }); + let mut variant = patterns.to_vec(); + variant[p] = GraphPattern::Path(identity); + variants.push(variant); + } + variants.extend(sequences.into_iter().map(|sequence| { + let last = sequence.len() - 1; + let hops = sequence + .into_iter() + .enumerate() + .map(|(h, (t, label))| PathSegment { + relationship: RelationshipPattern { + types: vec![t], + length: None, + ..rel.clone() + }, + end_node: if h == last { + path.segments[i].end_node.clone() + } else { + NodePattern { + variable: None, + labels: vec![label], + properties: HashMap::new(), + } + }, + }); + let mut unrolled = path.clone(); + unrolled.segments.splice(i..=i, hops); + let mut variant = patterns.to_vec(); + variant[p] = GraphPattern::Path(unrolled); + variant + })); Ok(Some(variants)) } @@ -515,6 +544,7 @@ fn rewrite( occurrence_slots: &[usize], labels: &[Option<&str>], types: &[&str], + retain_variable: bool, ) -> Vec { let mut occurrences = occurrence_slots.iter(); let mut types = types.iter(); @@ -532,8 +562,12 @@ fn rewrite( let mut path = path.clone(); path.start_node = relabel(&path.start_node); for segment in &mut path.segments { - segment.relationship.types = - vec![types.next().expect("one type per relationship").to_string()]; + if !(zero_hop(&segment.relationship) + || retain_variable && segment.relationship.length.is_some()) + { + segment.relationship.types = + vec![types.next().expect("one type per relationship").to_string()]; + } segment.end_node = relabel(&segment.end_node); } GraphPattern::Path(path) @@ -588,7 +622,7 @@ mod tests { let ReadingClause::Match(m) = &ast.reading_clauses[0] else { panic!("not a MATCH") }; - let out = expand(&m.patterns, &ldbc_config(), &HashMap::new())?; + let out = expand(&m.patterns, &ldbc_config(), &HashMap::new(), false)?; Ok(out.iter().map(|b| describe(b)).collect()) } diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 3b43c20..2223482 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -890,7 +890,9 @@ impl CypherQuery { } // Phase 2: Graph Logical Plan - let mut logical_planner = LogicalPlanner::new(config).with_fan_out(self.fan_out.as_ref()); + let mut logical_planner = LogicalPlanner::new(config) + .with_fan_out(self.fan_out.as_ref()) + .with_adjacency(self.adjacency.as_ref()); // Use the transformed AST (with parameters substituted) let logical_plan = logical_planner.plan(&semantic.ast)?; diff --git a/docs/csr.md b/docs/csr.md index f819ec6..e8d746c 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -121,7 +121,8 @@ When those hold, it covers: | --- | --- | | `(a)-[:KNOWS]->(b)`, `(a)<-[:KNOWS]-(b)` | :material-check: outgoing or incoming adjacency | | `(a)-[:KNOWS]-(b)` | :material-check: both adjacencies | -| `(a)-[:KNOWS*1..3]->(b)` | :material-check: every unrolled hop | +| `(a)-[:KNOWS*1..3]->(b)` | :material-check: one iterative expansion for the whole range | +| (t:Tag)-[:hasType|isSubclassOf*1..3]->(c:TagClass) | :material-check: typed iterative expansion when every type and label is indexed | | (a)-[:LIKES_POST|LIKES_COMMENT]->(m) | :material-check: one expansion per branch | | `OPTIONAL MATCH (a)-[:KNOWS]->(b)` | :material-check: | | `shortestPath((a)-[:KNOWS*]-(b))` | :material-check: the search walks the CSR adjacency | @@ -131,6 +132,23 @@ 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. +### Variable-length traversal + +An eligible variable-length path keeps one cursor per hop and yields bounded output batches. Its plan +contains one `CsrExpandExec` with `iterative hops=...`, rather than a branch for every possible length or +sequence of relationship types. It preserves duplicate matches and skips missing nodes at intermediate +hops. A node's label and key travel together, so the same integer key in two labels identifies two nodes. + +The iterative path needs CSR indexes for every traversed type and direction, node indexes for every +intermediate label, and resolved labels for the endpoints. Unsupported index shapes retain the join +fallback. The 20-hop maximum still applies. `*0..` includes a zero-hop identity match when both labels and +keys agree. Undirected paths keep the existing rule that rejects immediate backtracking; neither execution +path enforces full relationship uniqueness on longer cycles. + +Traversal state grows with the hop limit, not the number of pending paths. Resident graph indexes, input +batches, fetched target properties and downstream operators still contribute to query memory. Branching +graphs can produce exponentially many matches even though traversal resumes between output batches. + !!! tip "Leave relationship variables out unless you need them" `MATCH (a)-[k:KNOWS]->(b) RETURN b.name` falls back to joins only because it names `k`. If you don't read diff --git a/docs/cypher.md b/docs/cypher.md index d292977..4c8d0d5 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -209,8 +209,13 @@ MATCH p = shortestPath((a:Person {id: 1})-[:KNOWS*]-(b:Person {id: 4})) RETURN length(p), nodes(p) ``` -- **Variable length** defaults to a minimum of 1 and a maximum of 20 hops; a maximum above 20 is rejected, - and zero-length paths (`*0..`) are not supported. Each length becomes one branch of a union. +- **Variable length** defaults to a minimum of 1 and a maximum of 20 hops; a maximum above 20 is rejected. + `*0..` also matches the source itself, once per input row; `*0..0` matches only that identity. The start + and end must have the same label and key for a zero-hop match. Eligible CSR paths use one iterative + 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. - **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. diff --git a/python/python/tests/test_engine_alternation.py b/python/python/tests/test_engine_alternation.py index 2d8d9fa..09e632a 100644 --- a/python/python/tests/test_engine_alternation.py +++ b/python/python/tests/test_engine_alternation.py @@ -128,6 +128,51 @@ def test_variable_length_alternation_follows_endpoint_labels(datasets, use_csr): assert rows(got) == [(20, 10), (21, 10)] +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("minimum", [0, 1]) +def test_variable_length_mixed_labels_keep_identity_separate(use_csr, minimum): + builder = GraphConfig.builder().node_table("Tag", "id").node_table("TagClass", "id") + cfg = ( + builder.rel_table("hasType", "Tag", "TagClass", "src", "dst") + .rel_table("isSubclassOf", "TagClass", "TagClass", "src", "dst") + .build() + ) + tables = { + "Tag": pa.table({"id": [1, 2]}), + "TagClass": pa.table({"id": [1, 2, 3], "name": ["one", "two", "three"]}), + "hasType": pa.table({"src": [1, 1, 2, 1], "dst": [1, 1, 2, 99]}), + "isSubclassOf": pa.table({"src": [1, 2], "dst": [2, 3]}), + } + engine = CypherEngine(cfg, tables, use_csr=use_csr) + query = ( + f"MATCH (t:Tag {{id: 1}})-[:hasType|isSubclassOf*{minimum}..3]->(c:TagClass) " + "RETURN c.id, c.name" + ) + assert rows(engine.execute(query)) == [ + (1, "one"), + (1, "one"), + (2, "two"), + (2, "two"), + (3, "three"), + (3, "three"), + ] + if use_csr: + plan = engine.explain(query) + physical = plan[plan.index("physical_plan") :] + assert physical.count("CsrExpandExec") == 1 + assert "UnionExec" not in physical + # Equal integer keys in different labels never form a zero-hop path. + identity = ( + "MATCH (t:Tag {id: 1})-[:hasType|isSubclassOf*0..0]->(c:TagClass) RETURN c.id" + ) + assert rows(engine.execute(identity)) == [] + same_label = ( + "MATCH (t:TagClass {id: 1})-[:hasType|isSubclassOf*0..2]->(c:TagClass) " + "RETURN c.id" + ) + assert rows(engine.execute(same_label)) == [(1,), (2,), (3,)] + + def test_collect_distinct(datasets): engine = CypherEngine(config(), datasets) got = engine.execute( diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index b5a4dd7..b96a8b9 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -119,11 +119,39 @@ def test_variable_length_hops_use_csr(datasets): plan = csr.explain("MATCH (a:Person)-[:KNOWS*1..2]-(b:Person) RETURN b.id") lines = plan[plan.index("physical_plan") :].splitlines() expands = [line for line in lines if "CsrExpandExec" in line] - # One expansion for the 1-hop branch and two for the 2-hop branch. - assert len(expands) == 3 and all("undirected" in line for line in expands) + assert len(expands) == 1 and "iterative hops=1..2" in expands[0] + assert "undirected" in expands[0] + assert "UnionExec" not in "\n".join(lines) assert "HashJoinExec" not in "\n".join(lines) +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("hops", ["0..0", "0..2"]) +def test_variable_length_zero_hop_is_identity(datasets, use_csr, hops): + engine = CypherEngine(CONFIG, datasets, use_csr=use_csr) + got = engine.execute( + f"MATCH (a:Person {{id: 1}})-[:KNOWS*{hops}]->(b:Person) " + "RETURN b.id AS id, b.name AS name" + ).to_pylist() + expected = [(1, "a")] + if hops == "0..2": + expected += [(2, "b")] * 2 + [(3, "c")] * 4 + assert sorted((r["id"], r["name"]) for r in got) == sorted(expected) + + +def test_variable_length_plan_does_not_grow_with_maximum(datasets): + engine = CypherEngine(CONFIG, datasets, use_csr=True) + for maximum in [2, 5, 20]: + plan = engine.explain( + f"MATCH (a:Person)-[:KNOWS*1..{maximum}]->(b:Person) RETURN b.id" + ) + physical = plan[plan.index("physical_plan") :] + assert physical.count("CsrExpandExec") == 1 + assert "UnionExec" not in physical + with pytest.raises(NotImplementedError, match="20"): + engine.execute("MATCH (a:Person)-[:KNOWS*1..21]->(b:Person) RETURN b.id") + + def test_source_filter_pushed_below_csr(datasets): csr = CypherEngine(CONFIG, datasets, use_csr=True) plan = csr.explain( From c6bc5c150a2e198e2346e1870509db0f4ce69887 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 23:13:49 -0700 Subject: [PATCH 2/2] fix(csr): bound raw and empty work in iterative traversal --- crates/lance-graph/src/csr_expand.rs | 145 +++++++++++++++++++++--- crates/lance-graph/src/csr_traversal.rs | 58 +++++++++- docs/csr.md | 3 + 3 files changed, 185 insertions(+), 21 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 9bb4c20..a3fa5bb 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -215,6 +215,7 @@ impl Adjacency { self.neighbors_from(key, &mut 0).for_each(f); } + #[inline(always)] fn base_neighbors(&self, key: i64) -> &[u64] { self.base .dense @@ -223,21 +224,34 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, offset: &'a mut usize, ) -> impl Iterator + 'a { + self.candidates_from(key, offset).flatten() + } + + /// Keep tombstones visible to cancellation without repeating lookups for every edge. + #[inline(always)] + fn candidates_from<'a>( + &'a self, + key: i64, + offset: &'a mut usize, + ) -> impl Iterator> + 'a { let base = self.base_neighbors(key); let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; - base.iter() - .chain(added) - .enumerate() - .skip(*offset) - .filter_map(move |(position, &neighbor)| { - *offset = position + 1; + let added_offset = offset.saturating_sub(base.len()); + base.get(*offset..) + .unwrap_or_default() + .iter() + .chain(added.get(added_offset..).unwrap_or_default()) + .map(move |&neighbor| { + let position = *offset; + *offset += 1; if removed.first() == Some(&position) { removed = &removed[1..]; None @@ -772,7 +786,7 @@ pub struct Neighbors { } /// Occurrence offsets are sufficient to pause even a high-degree source row. -#[derive(Default)] +#[derive(Clone, Copy, Default)] pub(crate) struct NeighborPosition { forward: usize, reverse: usize, @@ -798,21 +812,32 @@ impl Neighbors { .for_each(f); } + #[inline(always)] pub(crate) fn neighbors_from<'a>( &'a self, key: i64, position: &'a mut NeighborPosition, ) -> impl Iterator + 'a { - let forward = self.forward.neighbors_from(key, &mut position.forward); + self.candidates_from(key, position).flatten() + } + + /// Borrow each adjacency once while exposing tombstones and suppressed self-loops as work. + #[inline(always)] + pub(crate) fn candidates_from<'a>( + &'a self, + key: i64, + position: &'a mut NeighborPosition, + ) -> impl Iterator> + 'a { + let forward = self.forward.candidates_from(key, &mut position.forward); let reverse = self .reverse .as_ref() - .map(|r| r.neighbors_from(key, &mut position.reverse)); + .map(|r| r.candidates_from(key, &mut position.reverse)); forward.chain( reverse .into_iter() .flatten() - .filter(move |&n| n as i64 != key), + .map(move |candidate| candidate.filter(|&n| n as i64 != key)), ) } @@ -1183,12 +1208,14 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); - while self.row < self.keys.len() { - let i = self.row; + let mut position = self.position; + let mut row = self.row; + 'rows: while row < self.keys.len() { + let i = row; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); let required = self.bound.as_ref().map(|b| b.value(i)); - for neighbor in adjacency.neighbors_from(key, &mut self.position) { + for neighbor in adjacency.neighbors_from(key, &mut position) { let neighbor = neighbor as i64; if required.is_some_and(|r| r != neighbor) { continue; @@ -1204,13 +1231,15 @@ impl BatchExpansion { chunk.key_values.push(key); chunk.neighbor_values.push(neighbor); if chunk.rows.len() == batch_size.max(1) { - return chunk; + break 'rows; } } } - self.row += 1; - self.position = NeighborPosition::default(); + row += 1; + position = NeighborPosition::default(); } + self.row = row; + self.position = position; chunk } @@ -1221,7 +1250,12 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + let mut source_work = 0_usize; while self.row < self.keys.len() { + source_work += 1; + if source_work.is_multiple_of(1024) { + cooperative_yield().await; + } let i = self.row; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); @@ -1847,6 +1881,33 @@ mod tests { assert!(Arc::ptr_eq(&base.base, &later.base)); } + #[tokio::test] + async fn resumed_neighbors_cross_the_base_and_delta_boundary() { + let base = Adjacency::build( + edge_stream(vec![Some(1), Some(1)], vec![Some(10), Some(11)]), + "src", + "dst", + ) + .await + .unwrap(); + let next = base.with_changes(&[(1, 12), (1, 12), (1, 13)], &[(1, 11), (1, 12)]); + for (mut offset, expected) in [ + (0, vec![10, 12, 13]), + (1, vec![12, 13]), + (2, vec![12, 13]), + (3, vec![12, 13]), + (4, vec![13]), + (5, vec![]), + (6, vec![]), + ] { + assert_eq!( + next.neighbors_from(1, &mut offset).collect::>(), + expected + ); + } + assert_eq!(neighbors(&base, 1), vec![10, 11]); + } + #[tokio::test] async fn debug_output_does_not_grow_with_edges() { let n = 10_000; @@ -1949,6 +2010,58 @@ mod tests { assert!(expansion.next_chunk(&adj, None, 1).rows.is_empty()); } + #[tokio::test] + async fn cancellation_without_edges_in_variable_paths_yields_between_source_rows() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let nodes = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + let mut catalog = AdjacencyCatalog::default(); + catalog.insert_nodes( + "Person", + Arc::new(NodeRows::from_batch(nodes, "id").unwrap().unwrap()), + ); + catalog.insert( + "KNOWS", + true, + Arc::new( + Adjacency::build(edge_stream(vec![], vec![]), "src", "dst") + .await + .unwrap(), + ), + ); + let walk = crate::csr_traversal::Traversal::new( + &config, + &catalog, + &["KNOWS".into()], + &crate::ast::RelationshipDirection::Outgoing, + "Person", + "Person", + 1, + 1, + ) + .unwrap(); + let input = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![1; 100_000])) as ArrayRef, + )]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_path_chunk(&walk, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.row > 0 && expansion.row <= 1024); + assert!(expansion + .next_path_chunk(&walk, None, 1) + .await + .rows + .is_empty()); + } + #[tokio::test] async fn high_degree_expansion_visits_only_the_first_chunk() { let n = 100_000; diff --git a/crates/lance-graph/src/csr_traversal.rs b/crates/lance-graph/src/csr_traversal.rs index a19eb36..cdf11cb 100644 --- a/crates/lance-graph/src/csr_traversal.rs +++ b/crates/lance-graph/src/csr_traversal.rs @@ -161,14 +161,18 @@ impl TraversalCursor { let next = (step.from == frame.label) .then(|| { step.adjacency - .neighbors_from(frame.key, &mut frame.position) + .candidates_from(frame.key, &mut frame.position) .next() }) .flatten(); - let Some(key) = next.map(|n| n as i64) else { - frame.step += 1; - frame.position = NeighborPosition::default(); - continue; + let key = match next { + Some(Some(key)) => key as i64, + Some(None) => continue, + None => { + frame.step += 1; + frame.position = NeighborPosition::default(); + continue; + } }; if graph.nodes[step.to].position(key).is_none() { continue; @@ -271,6 +275,50 @@ mod tests { assert!(cursor.stack.is_empty()); } + #[tokio::test] + async fn cancellation_visits_raw_tombstones_before_returning_the_live_path() { + let (config, mut catalog) = graph().await; + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::Int64, false), + Field::new("dst", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1; 5000])), + Arc::new(Int64Array::from(vec![1; 5000])), + ], + ) + .unwrap(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(vec![Ok(batch)]), + )); + let base = Adjacency::build(stream, "src", "dst").await.unwrap(); + catalog.insert( + "hasType", + true, + Arc::new(base.with_changes(&[], &vec![(1, 1); 4999])), + ); + let walk = Traversal::new( + &config, + &catalog, + &["hasType".into()], + &RelationshipDirection::Outgoing, + "Tag", + "Class", + 1, + 1, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut next = Box::pin(cursor.next(&walk)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert_eq!(cursor.next(&walk).await, Some(1)); + assert_eq!(cursor.next(&walk).await, None); + } + #[tokio::test] async fn zero_hop_is_identity_and_depth_bound_survives_cycles() { let (config, catalog) = graph().await; diff --git a/docs/csr.md b/docs/csr.md index e8d746c..3b530a6 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -99,6 +99,9 @@ Node index size; wide properties and multiple active partitions still increase working memory. This does not bound the resident graph, the input data, or buffering in joins, sorts and aggregations elsewhere in the plan. + Iteration checkpoints count raw edge occurrences, including tombstones, and empty source rows. + Traversal therefore remains cooperative even when no path survives. + !!! warning "Experimental" The CSR is marked experimental. Results are identical with it on or off (the test suite runs each