diff --git a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs index 9b4ad83..4f4f943 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs @@ -358,6 +358,19 @@ impl DataFusionPlanner { .zip(cat.get(rel_type, false)) .map(|(out, inc)| Neighbors::undirected(out.clone(), inc.clone())), }); + let reverse_adjacency = self + .adjacency + .as_ref() + .filter(|_| adjacency.is_some()) + .and_then(|cat| match search { + SearchDirection::Outgoing => { + cat.get(rel_type, false).cloned().map(Neighbors::directed) + } + SearchDirection::Incoming => { + cat.get(rel_type, true).cloned().map(Neighbors::directed) + } + SearchDirection::Undirected => None, // Forward neighbours already visit both ends. + }); let source = self .get_catalog()? .relationship_source(&rel_map.relationship_type) @@ -400,7 +413,8 @@ impl DataFusionPlanner { edge_source: schema.field(src).name().clone(), edge_target: schema.field(dst).name().clone(), }; - let node = ShortestPathNode::try_new(input_plan, edges, adjacency, spec)?; + let node = ShortestPathNode::try_new(input_plan, edges, adjacency, spec)? + .with_reverse_adjacency(reverse_adjacency); Ok(LogicalPlan::Extension(Extension { node: Arc::new(node), })) diff --git a/crates/lance-csr/src/shortest_path.rs b/crates/lance-csr/src/shortest_path.rs index 2a47450..9c37dec 100644 --- a/crates/lance-csr/src/shortest_path.rs +++ b/crates/lance-csr/src/shortest_path.rs @@ -12,9 +12,9 @@ //! a second input, a scan of the relationship's two endpoint columns, and builds an adjacency //! from it once per query, shared by all partitions. //! -//! One search runs per distinct source id in a batch, level by level, until every target in -//! that batch is reached at its shortest depth (or the depth limit). Predecessors are kept -//! per node, so all shortest paths can be enumerated back from each target. +//! A source with one target searches from both ends when reverse adjacency is available, +//! expanding the smaller complete frontier. Sources with several targets share one +//! forward search. Predecessors retain every shortest path, including parallel edges. use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -78,6 +78,7 @@ pub struct ShortestPathNode { /// Scan of the relationship's endpoints; `None` when `adjacency` is given. edges: Option, adjacency: Option, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: DFSchemaRef, } @@ -111,10 +112,16 @@ impl ShortestPathNode { input, edges, adjacency, + reverse_adjacency: None, spec, schema: Arc::new(schema), }) } + + pub(crate) fn with_reverse_adjacency(mut self, adjacency: Option) -> Self { + self.reverse_adjacency = adjacency; + self + } } impl PartialEq for ShortestPathNode { @@ -122,6 +129,11 @@ impl PartialEq for ShortestPathNode { self.input == other.input && self.edges == other.edges && self.spec == other.spec + && match (&self.reverse_adjacency, &other.reverse_adjacency) { + (Some(a), Some(b)) => a.ptr_eq(b), + (None, None) => true, + _ => false, + } && match (&self.adjacency, &other.adjacency) { (Some(a), Some(b)) => a.ptr_eq(b), (None, None) => true, @@ -188,6 +200,7 @@ impl UserDefinedLogicalNodeCore for ShortestPathNode { let edges = (inputs.len() > 1).then(|| inputs.pop().expect("two inputs")); let input = inputs.pop().expect("an input"); Self::try_new(input, edges, self.adjacency.clone(), self.spec.clone()) + .map(|node| node.with_reverse_adjacency(self.reverse_adjacency.clone())) .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string())) } } @@ -200,6 +213,7 @@ pub struct ShortestPathExec { edges: Option>, /// The CSR, or the adjacency built from `edges` by the first batch that needs it. adjacency: Arc>>, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: SchemaRef, properties: Arc, @@ -210,6 +224,7 @@ impl ShortestPathExec { input: Arc, edges: Option>, adjacency: Option, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: SchemaRef, ) -> Self { @@ -223,6 +238,7 @@ impl ShortestPathExec { input, edges, adjacency: Arc::new(Mutex::new(adjacency)), + reverse_adjacency, spec, schema, properties, @@ -277,7 +293,14 @@ impl ExecutionPlan for ShortestPathExec { ) -> datafusion::common::Result> { let edges = (children.len() > 1).then(|| children.pop().expect("two children")); let input = children.pop().expect("a child"); - let mut exec = Self::new(input, edges, None, self.spec.clone(), self.schema.clone()); + let mut exec = Self::new( + input, + edges, + None, + self.reverse_adjacency.clone(), + self.spec.clone(), + self.schema.clone(), + ); exec.adjacency = self.adjacency.clone(); Ok(Arc::new(exec)) } @@ -297,7 +320,9 @@ impl ExecutionPlan for ShortestPathExec { self.schema.clone(), ); let out_schema = schema.clone(); + let reverse_adjacency = self.reverse_adjacency.clone(); let stream = input.then(move |batch| { + let reverse_adjacency = reverse_adjacency.clone(); let (cell, edges, spec, schema, context) = ( cell.clone(), edges.clone(), @@ -318,6 +343,9 @@ impl ExecutionPlan for ShortestPathExec { source_index, target_index, &adjacency, + reverse_adjacency.as_ref().or_else(|| { + (spec.direction == SearchDirection::Undirected).then_some(&adjacency) + }), &spec, &schema, ) @@ -372,22 +400,7 @@ fn search( let mut level = 0; while !frontier.is_empty() && !remaining.is_empty() && level < max_hops { level += 1; - let mut next = Vec::new(); - for &node in &frontier { - adjacency.for_each_neighbor(node, |n| { - let n = n as i64; - match depth.get(&n) { - // Another shortest way into a node first reached at this level. - Some(&d) if d == level => parents.entry(n).or_default().push(node), - Some(_) => {} - None => { - depth.insert(n, level); - parents.insert(n, vec![node]); - next.push(n); - } - } - }); - } + let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents); // Finish the level before stopping, so every shortest path is recorded. for n in &next { remaining.remove(n); @@ -397,6 +410,93 @@ fn search( (depth, parents) } +/// Finish the whole level so parallel edges and alternative shortest paths survive. +fn expand_level( + adjacency: &Neighbors, + frontier: &[i64], + level: u32, + depth: &mut FastMap, + parents: &mut FastMap>, +) -> Vec { + let mut next = Vec::new(); + for &node in frontier { + adjacency.for_each_neighbor(node, |n| { + let n = n as i64; + match depth.get(&n) { + Some(&d) if d == level => parents.entry(n).or_default().push(node), + Some(_) => {} + None => { + depth.insert(n, level); + parents.insert(n, vec![node]); + 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( + forward: &Neighbors, + reverse: &Neighbors, + source: i64, + target: i64, + max_hops: u32, +) -> (FastMap, FastMap>) { + let mut depth = FastMap::default(); + let mut parents = FastMap::default(); + depth.insert(source, 0); + if source == target { + return (depth, parents); + } + let mut reverse_depth = FastMap::default(); + let mut reverse_parents: FastMap> = FastMap::default(); + reverse_depth.insert(target, 0); + let (mut frontier, mut reverse_frontier) = (vec![source], vec![target]); + let (mut level, mut reverse_level) = (0, 0); + 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); + } else { + reverse_level += 1; + reverse_frontier = expand_level( + reverse, + &reverse_frontier, + reverse_level, + &mut reverse_depth, + &mut reverse_parents, + ); + } + let mut pending: Vec<_> = frontier + .iter() + .copied() + .filter(|n| reverse_depth.get(n) == Some(&reverse_level)) + .collect(); + if pending.is_empty() { + continue; + } + // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths. + let mut visited = HashSet::new(); + while let Some(node) = pending.pop() { + if !visited.insert(node) { + continue; + } + if let Some(next) = reverse_parents.get(&node) { + for &n in next { + parents.entry(n).or_default().push(node); + pending.push(n); + } + } + } + depth.insert(target, level + reverse_level); + break; + } + (depth, parents) +} + /// 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> { @@ -431,6 +531,7 @@ fn search_batch( source_index: usize, target_index: usize, adjacency: &Neighbors, + reverse_adjacency: Option<&Neighbors>, spec: &ShortestPathSpec, schema: &SchemaRef, ) -> datafusion::common::Result { @@ -450,7 +551,19 @@ fn search_batch( } let searches: FastMap = by_source .iter() - .map(|(&s, ts)| (s, search(adjacency, s, ts, spec.max_hops))) + .map(|(&s, ts)| { + let result = match reverse_adjacency.filter(|_| ts.len() == 1) { + Some(reverse) => search_bidirectional( + adjacency, + reverse, + s, + *ts.iter().next().expect("one target"), + spec.max_hops, + ), + None => search(adjacency, s, ts, spec.max_hops), + }; + (s, result) + }) .collect(); let mut rows = Vec::new(); let mut lengths = Vec::new(); @@ -498,6 +611,7 @@ pub(crate) fn plan_shortest_path( physical_inputs[0].clone(), edges, node.adjacency.clone(), + node.reverse_adjacency.clone(), node.spec.clone(), Arc::new(node.schema.as_arrow().clone()), )) @@ -514,7 +628,7 @@ fn plan_error(message: &str) -> GraphError { mod tests { use super::*; - fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + 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), @@ -547,7 +661,12 @@ mod tests { let inc = rt .block_on(Adjacency::build(stream(), "dst", "src")) .unwrap(); - Neighbors::undirected(Arc::new(out), Arc::new(inc)) + (Arc::new(out), Arc::new(inc)) + } + + fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + let (out, inc) = adjacencies(edges); + Neighbors::undirected(out, inc) } #[test] @@ -564,4 +683,105 @@ mod tests { let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1); assert!(!depth.contains_key(&4)); } + + #[test] + 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(&[ + (1, 2), + (1, 2), + (1, 3), + (1, 9), + (2, 4), + (2, 4), + (3, 4), + (3, 4), + (4, 5), + (5, 5), + (5, 1), + ]); + let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); + let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3); + assert_eq!(depth[&5], 3); + let mut all = paths_to(5, &parents, true); + all.sort(); + assert_eq!( + all, + vec![vec![1, 2, 4, 5]; 4] + .into_iter() + .chain(vec![vec![1, 3, 4, 5]; 2]) + .collect::>() + ); + let one = paths_to(5, &parents, false); + assert_eq!(one.len(), 1); + assert!(all.contains(&one[0])); + assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2) + .0 + .contains_key(&5)); + } + + #[test] + 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()) { + let selected: Vec<_> = edges + .iter() + .enumerate() + .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge)) + .collect(); + let (out, inc) = adjacencies(&selected); + let directed = Neighbors::directed(out.clone()); + let reverse = Neighbors::directed(inc.clone()); + let undirected = Neighbors::undirected(out, inc); + for (forward, reverse) in [(&directed, &reverse), (&undirected, &undirected)] { + for source in 0..3 { + for target in 0..3 { + for max_hops in 0..=3 { + let (expected_depth, expected_parents) = + search(forward, source, &HashSet::from([target]), max_hops); + let (depth, parents) = + search_bidirectional(forward, reverse, source, target, max_hops); + 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); + expected.sort(); + actual.sort(); + assert_eq!( + actual, expected, + "graph {mask}, {source}->{target}, max {max_hops}" + ); + } + } + } + } + } + } + } + + #[test] + fn bidirectional_avoids_expanding_the_wide_middle_layer() { + let mut edges = Vec::new(); + for branch in 1..=100 { + edges.push((0, branch)); + for leaf in 0..100 { + edges.push((branch, 101 + (branch - 1) * 100 + leaf)); + } + } + let (out, inc) = adjacencies(&edges); + 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]]); + assert_eq!(depth.len(), 102); // Root, first layer and stitched target. + assert_eq!( + search(&forward, 0, &HashSet::from([101]), 2).0.len(), + 10_101 + ); + } } diff --git a/docs/csr.md b/docs/csr.md index f819ec6..0d3ae5a 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -131,6 +131,12 @@ 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. +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 +one forward search. Undirected searches can use both ends without a persistent CSR too; directed searches +without a CSR keep their forward search over the adjacency built for that query. + !!! 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 04b1472..0645739 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -219,7 +219,11 @@ RETURN length(p), nodes(p) - **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. - With no path, `MATCH` drops the row and `OPTIONAL MATCH` returns `null`. + With no path, `MATCH` drops the row and `OPTIONAL MATCH` returns `null`. `shortestPath` chooses one path + when several have the same minimum length; `allShortestPaths` returns every shortest path, including + 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`. ### Expressions and predicates diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 1206e40..922a369 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -95,3 +95,45 @@ def test_search_without_csr_reads_every_edge_partition(): "MATCH p = shortestPath((a)-[:KNOWS*1..30]->(b)) RETURN length(p) AS l" ) assert got.to_pylist() == [{"l": 5}] + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("direction", ["outgoing", "incoming", "undirected"]) +def test_all_shortest_paths_keeps_every_meeting_node(use_csr, direction): + people = pa.table({"id": list(range(1, 10))}) + knows = pa.table( + { + "src": [1, 1, 1, 1, 2, 2, 3, 3, 4, 5, 5], + "dst": [2, 2, 3, 9, 4, 4, 4, 4, 5, 5, 6], + } + ) + engine = CypherEngine(CONFIG, {"Person": people, "KNOWS": knows}, use_csr=use_csr) + pattern = { + "outgoing": "(a)-[:KNOWS*1..3]->(b)", + "incoming": "(b)<-[:KNOWS*1..3]-(a)", + "undirected": "(a)-[:KNOWS*1..3]-(b)", + }[direction] + got = engine.execute( + "MATCH (a:Person {id: 1}), (b:Person {id: 5}) " + f"MATCH p = allShortestPaths({pattern}) RETURN length(p) AS l, nodes(p) AS ns" + ).to_pylist() + paths = [[1, 2, 4, 5]] * 4 + [[1, 3, 4, 5]] * 2 + if direction == "incoming": + paths = [p[::-1] for p in paths] + assert sorted((r["l"], r["ns"]) for r in got) == sorted((3, p) for p in paths) + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_shortest_paths_shares_a_source_across_targets(use_csr): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) + got = engine.execute( + "MATCH (a:Person {id: 5}), (b:Person) " + "MATCH p = allShortestPaths((a)-[:KNOWS*1..3]->(b)) " + "RETURN b.id AS target, length(p) AS l, nodes(p) AS ns" + ).to_pylist() + assert sorted((r["target"], r["l"], r["ns"]) for r in got) == [ + (1, 1, [5, 1]), + (2, 2, [5, 1, 2]), + (2, 2, [5, 1, 2]), + (3, 2, [5, 1, 3]), + ]