fix(planner): match a reused relationship variable as the same edge

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<dyn crate::GraphSourceCatalog>,
+        left_plan: LogicalPlan,
+        bound: &BoundRelationship,
+        source_variable: &str,
+        target_variable: &str,
+        target_label: &str,
+        direction: &RelationshipDirection,
+        relationship_properties: &HashMap<String, crate::ast::PropertyValue>,
+        target_properties: &HashMap<String, crate::ast::PropertyValue>,
+    ) -> Result<LogicalPlan> {
+        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.
     ///

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<String, String>,
+    /// The types each relationship variable was first written with.
+    relationship_types: HashMap<String, Vec<String>>,
     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<Vec<GraphPattern>> {
+        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<String, bool> = 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.

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::<Int64Array>()
+            .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();

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

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")