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 ed30c57..ce1a721 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs @@ -10,10 +10,19 @@ use crate::datafusion_planner::join_ops::{SourceJoinParams, TargetJoinParams}; use crate::datafusion_planner::DataFusionPlanner; use crate::error::Result; use crate::logical_plan::*; -use datafusion::logical_expr::{col, LogicalPlan, LogicalPlanBuilder}; +use datafusion::logical_expr::{col, lit, Expr, JoinType, LogicalPlan, LogicalPlanBuilder}; use std::collections::HashMap; use std::sync::Arc; +/// A relationship variable that an earlier hop bound, met again in a pattern. +struct BoundRelationship<'a> { + variable: &'a str, + rel_map: &'a crate::config::RelationshipMapping, + rel_schema: arrow_schema::SchemaRef, + /// Whether this hop's type matches the bound one; a mismatch matches nothing. + same_type: bool, +} + impl DataFusionPlanner { /// Build a relationship expansion (graph traversal) as a series of joins #[allow(clippy::too_many_arguments)] @@ -87,6 +96,40 @@ impl DataFusionPlanner { return Ok(left_plan); }; + if let Some(var) = relationship_variable { + let bound_source = qualify_column(var, &rel_map.source_id_field); + if left_plan + .schema() + .field_with_unqualified_name(&bound_source) + .is_ok() + { + let bound_type = ctx + .analysis + .relationship_instances + .iter() + .find(|r| r.alias == var) + .map(|r| &r.rel_type); + let bound = BoundRelationship { + variable: var, + rel_map, + rel_schema: rel_source.schema(), + same_type: bound_type.is_some_and(|t| t.eq_ignore_ascii_case(rel_type)), + }; + return self.join_bound_relationship( + ctx, + cat, + left_plan, + &bound, + source_variable, + target_variable, + target_label, + direction, + relationship_properties, + target_properties, + ); + } + } + // A bound relationship variable or relationship property filter needs the // relationship's other columns, which the CSR does not carry. let csr_eligible = relationship_variable.is_none() @@ -177,6 +220,96 @@ impl DataFusionPlanner { self.join_relationship_to_target(builder, cat, &target_params) } + /// A hop over a relationship variable the input already binds: the same edge again, so + /// its endpoint columns must equal the two nodes' ids, and no relationship is scanned. + // shortcut: an undirected binding holds its edge as walked, so a later directed reuse + // compares against that orientation; upgrade by keeping the stored endpoints too. + #[allow(clippy::too_many_arguments)] + fn join_bound_relationship( + &self, + ctx: &PlanningContext, + cat: &Arc, + left_plan: LogicalPlan, + bound: &BoundRelationship, + source_variable: &str, + target_variable: &str, + target_label: &str, + direction: &RelationshipDirection, + relationship_properties: &HashMap, + target_properties: &HashMap, + ) -> Result { + let (_, source_map) = self.get_target_node_mapping(ctx, source_variable)?; + let target_map = self.config.get_node_mapping(target_label).ok_or_else(|| { + crate::error::GraphError::ConfigError { + message: format!("No mapping found for target label: {}", target_label), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + let rel_column = |name: &str| { + let physical = bound + .rel_schema + .fields() + .iter() + .find(|f| f.name().eq_ignore_ascii_case(name)) + .map_or(name, |f| f.name()); + col(qualify_column(bound.variable, physical)) + }; + let source = col(qualify_column(source_variable, &source_map.id_field)); + let target_key = qualify_column(target_variable, &target_map.id_field); + let target = col(&target_key); + let ends = |from: Expr, to: Expr| source.clone().eq(from).and(target.clone().eq(to)); + let (src, dst) = ( + rel_column(&bound.rel_map.source_id_field), + rel_column(&bound.rel_map.target_id_field), + ); + let mut condition = match direction { + RelationshipDirection::Outgoing => ends(src, dst), + RelationshipDirection::Incoming => ends(dst, src), + RelationshipDirection::Undirected => ends(src.clone(), dst.clone()).or(ends(dst, src)), + } + .and(lit(bound.same_type)); + for (k, v) in relationship_properties { + let value = + self.unresolved_value_expr(&crate::ast::ValueExpression::Literal(v.clone())); + condition = condition.and(rel_column(k).eq(value)); + } + + let builder = LogicalPlanBuilder::from(left_plan); + let target_bound = builder + .schema() + .field_with_unqualified_name(&target_key) + .is_ok(); + let builder = if target_bound { + if !target_properties.is_empty() { + return Err(crate::error::GraphError::PlanError { + message: format!( + "Property filters on reused variable '{}' are not yet supported", + target_variable + ), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + builder.filter(condition) + } else { + let target_source = cat.node_source(target_label).ok_or_else(|| { + crate::error::GraphError::ConfigError { + message: format!("No node source for target label: {}", target_label), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + let scan = self.build_target_scan( + target_source, + target_label, + target_variable, + target_properties, + )?; + builder.join_on(scan, JoinType::Inner, [condition]) + }; + builder + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to plan bound relationship", e)) + } + /// Replace "relationship scan + source join" with a CSR lookup, if an adjacency /// exists for this relationship and direction and every id column is an integer. /// diff --git a/crates/lance-csr/src/logical_plan.rs b/crates/lance-csr/src/logical_plan.rs index 7346c8c..30eb9d3 100644 --- a/crates/lance-csr/src/logical_plan.rs +++ b/crates/lance-csr/src/logical_plan.rs @@ -250,6 +250,8 @@ pub struct LogicalPlanner<'a> { calls: usize, /// The node label of each path variable's nodes, for `nodes(path)`. path_labels: HashMap, + /// The types each relationship variable was first written with. + relationship_types: 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 @@ -269,6 +271,7 @@ impl<'a> LogicalPlanner<'a> { comprehensions: 0, calls: 0, path_labels: HashMap::new(), + relationship_types: HashMap::new(), rich_paths: false, seed_source: None, } @@ -853,9 +856,10 @@ impl<'a> LogicalPlanner<'a> { location: snafu::Location::new(file!(), line!(), column!()), }); } + let typed = self.reuse_relationship_types(&match_clause.patterns)?; let inferred = MatchClause { patterns: crate::pattern_alternatives::infer_labels( - &match_clause.patterns, + &typed, self.config, &self.variables, ), @@ -951,6 +955,64 @@ impl<'a> LogicalPlanner<'a> { } } + /// A relationship variable written again without a type (`(c)-[r]->(d)`) takes the types + /// it was first written with, so its endpoints get labels and it plans as the same edge. + // shortcut: types are remembered by name for the whole query, even once WITH drops the + // variable; scope them like `variables` if a query reuses a name for a new relationship. + fn reuse_relationship_types(&mut self, patterns: &[GraphPattern]) -> Result> { + fn relationships(patterns: &mut [GraphPattern]) -> Vec<&mut RelationshipPattern> { + patterns + .iter_mut() + .filter_map(|p| match p { + GraphPattern::Path(path) => Some(path), + GraphPattern::Node(_) => None, + }) + .flat_map(|path| path.segments.iter_mut()) + .map(|s| &mut s.relationship) + .filter(|r| r.variable.is_some()) + .collect() + } + let mut patterns = patterns.to_vec(); + // Whether each variable seen in this clause was variable-length. + let mut seen: HashMap = HashMap::new(); + for rel in relationships(&mut patterns) { + let var = rel + .variable + .as_ref() + .expect("filtered to named relationships"); + let variable_length = rel.length.is_some(); + if let Some(earlier_variable_length) = seen.insert(var.clone(), variable_length) { + if variable_length || earlier_variable_length { + return Err(GraphError::UnsupportedFeature { + feature: format!( + "reusing relationship variable '{var}' in a variable-length segment" + ), + location: Location::new(file!(), line!(), column!()), + }); + } + } + } + for rel in relationships(&mut patterns) { + if !rel.types.is_empty() { + let var = rel + .variable + .clone() + .expect("filtered to named relationships"); + self.relationship_types + .entry(var) + .or_insert_with(|| rel.types.clone()); + } + } + for rel in relationships(&mut patterns) { + if rel.types.is_empty() { + if let Some(types) = self.relationship_types.get(rel.variable.as_ref().unwrap()) { + rel.types = types.clone(); + } + } + } + Ok(patterns) + } + /// openCypher binds each relationship at most once per MATCH. For single-hop relationships /// that could bind the same edge, returns the patterns with their anonymous nodes named /// and a predicate that the edges differ, compared by their endpoint node ids. diff --git a/crates/lance-csr/tests/test_datafusion_pipeline.rs b/crates/lance-csr/tests/test_datafusion_pipeline.rs index 15b3a2c..7c7c310 100644 --- a/crates/lance-csr/tests/test_datafusion_pipeline.rs +++ b/crates/lance-csr/tests/test_datafusion_pipeline.rs @@ -5025,6 +5025,37 @@ async fn test_datafusion_variable_reuse_with_count_distinct() { assert_eq!(counts.value(0), 1); } +#[tokio::test] +async fn test_datafusion_relationship_variable_reuse_binds_the_same_edge() { + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("KNOWS".to_string(), create_knows_dataset()); + + // The second part has no type and its endpoints no labels: both come from `r`. + let query = CypherQuery::new( + "MATCH (a:Person)-[r:KNOWS]->(b:Person), (c)-[r]->(d) \ + RETURN COUNT(*) AS total, \ + SUM(CASE WHEN a.id = c.id AND b.id = d.id THEN 1 ELSE 0 END) AS same", + ) + .unwrap() + .with_config(create_graph_config()); + let result = query + .execute(datasets, Some(ExecutionStrategy::DataFusion)) + .await + .unwrap(); + + let value = |i: usize| { + result + .column(i) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + // One row per KNOWS edge, each with c = a and d = b. + assert_eq!((value(0), value(1)), (5, 5)); +} + #[tokio::test] async fn test_datafusion_variable_reuse_triangle_pattern() { let config = create_graph_config(); diff --git a/docs/cypher.md b/docs/cypher.md index c8a1bbd..9faf5db 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -319,6 +319,13 @@ RETURN length(p), nodes(p) `(a)-[:KNOWS]->(b)<-[:KNOWS]-(c)` therefore never matches `c = a`. Separate `MATCH` clauses may bind the same edge again. Edges are compared by their endpoints, so two parallel edges between the same nodes count as one. Variable-length and untyped relationships are not checked, except inside a named path. +- **Reused relationship variables**: writing a bound relationship variable again, as in + `MATCH (a)-[r:KNOWS]->(b), (c)-[r]->(d)`, matches that same relationship, so `c = a` and `d = b`. This + also works in a later `MATCH` clause or after `WITH r`. The type can be left out; a different type + matches nothing. An undirected reuse matches both orientations. When the first binding is undirected, + a later directed reuse follows the direction the first one walked the edge, not its stored direction. + Reusing a relationship variable in a variable-length segment is not supported; within one `MATCH` it + raises an error. - **`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`. `shortestPath` chooses one path diff --git a/python/python/tests/test_relationship_uniqueness.py b/python/python/tests/test_relationship_uniqueness.py index 66c218b..618eca6 100644 --- a/python/python/tests/test_relationship_uniqueness.py +++ b/python/python/tests/test_relationship_uniqueness.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The lance-csr Authors -"""One MATCH binds each relationship at most once (relationship isomorphism).""" +"""Relationship identity: one MATCH binds each relationship at most once (relationship +isomorphism), and a relationship variable written again is that same relationship.""" import pyarrow as pa import pytest @@ -42,3 +43,45 @@ def test_one_match_binds_each_relationship_once(use_csr, query, expected): engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) assert engine.execute(query + RETURN).to_pylist() == expected + + +# A->B (w 10) and B->C (w 20). +CHAIN = pa.table({"src": [1, 2], "dst": [2, 3], "w": [10, 20]}) +SAME = "MATCH (a:Person)-[r:KNOWS]->(b:Person), " +ENDS = " RETURN a.name AS a, b.name AS b, c.name AS c, d.name AS d ORDER BY a, c" +AB, BC = ("A", "B"), ("B", "C") + + +def ends(*rows): + return [{"a": a, "b": b, "c": c, "d": d} for (a, b), (c, d) in rows] + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize( + ("query", "expected"), + [ + (SAME + "(c:Person)-[r]->(d:Person)" + ENDS, ends((AB, AB), (BC, BC))), + (SAME + "(c)-[r:KNOWS]->(d)" + ENDS, ends((AB, AB), (BC, BC))), + (SAME + "(d:Person)<-[r]-(c:Person)" + ENDS, ends((AB, AB), (BC, BC))), + ( + SAME + "(c:Person)-[r]-(d:Person)" + ENDS, + ends((AB, AB), (AB, AB[::-1]), (BC, BC), (BC, BC[::-1])), + ), + ( + "MATCH (a:Person)-[r:KNOWS]->(b:Person) MATCH (c)-[r]->(d)" + ENDS, + ends((AB, AB), (BC, BC)), + ), + (SAME + "(b)-[r]->(a) RETURN a.name AS a", []), + (SAME + "(c)-[r {w: 10}]->(d) RETURN a.name AS a", [{"a": "A"}]), + (SAME + "(c)-[r]->(d) RETURN r.w AS w ORDER BY w", [{"w": 10}, {"w": 20}]), + ], +) +def test_reused_relationship_variable_is_the_same_edge(use_csr, query, expected): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": CHAIN}, use_csr=use_csr) + assert engine.execute(query).to_pylist() == expected + + +def test_reused_relationship_variable_rejects_variable_length(): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": CHAIN}) + with pytest.raises(Exception, match="variable-length segment"): + engine.execute(SAME + "(c)-[r*1..2]->(d) RETURN c")