perf(planner): count two-hop paths over undirected hops

crates/lance-csr/src/datafusion_planner/builder/count_ops.rs

@@ -26,15 +26,25 @@ use datafusion::logical_expr::{
     col, lit, when, Expr, Extension, JoinType, LogicalPlan, LogicalPlanBuilder,
 };
 
-/// One `Expand` of the pattern, with its edge's endpoints in stored direction.
+/// The pairs of a two-hop pattern's edges that the distinct-edges predicate excludes, as
+/// the paths that would use one edge twice. Ordered so that the larger set wins.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum Excluded {
+    /// Both edges are self-loops at the middle node: a directed chain.
+    SelfLoops,
+    /// Both edges join the middle node to the same far end: hops that share an end, or
+    /// any undirected hop.
+    SameFarEnd,
+}
+
+/// One `Expand` of the pattern.
 struct Hop<'a> {
     op: &'a LogicalOperator,
     source: &'a str,
     target: &'a str,
     relationship: Option<&'a str>,
     rel_type: &'a str,
-    /// `(src, dst)` of the edge: the traversal reversed for an incoming hop.
-    edge: (&'a str, &'a str),
+    direction: &'a RelationshipDirection,
 }
 
 impl<'a> Hop<'a> {
@@ -52,19 +62,14 @@ impl<'a> Hop<'a> {
             return None;
         };
         let (source, target) = (source_variable.as_str(), target_variable.as_str());
-        let edge = match direction {
-            RelationshipDirection::Outgoing => (source, target),
-            RelationshipDirection::Incoming => (target, source),
-            RelationshipDirection::Undirected => return None,
-        };
         (relationship_types.len() == 1 && source != target).then_some((
             Self {
                 op,
                 source,
                 target,
                 relationship: relationship_variable.as_deref(),
                 rel_type: &relationship_types[0],
-                edge,
+                direction,
             },
             input.as_ref(),
         ))
@@ -77,12 +82,19 @@ impl<'a> Hop<'a> {
             .collect()
     }
 
-    /// Mean edges per `middle` node, from the engine's key statistics.
+    /// Mean edges per `middle` node, from the engine's key statistics: both directions'
+    /// for an undirected hop.
     fn fan_out(&self, fan_out: &FanOut, middle: &str) -> Option<f64> {
-        let outgoing = self.edge.0 == middle;
-        fan_out
-            .get(&(self.rel_type.to_lowercase(), outgoing))
-            .copied()
+        let mean = |outgoing| {
+            fan_out
+                .get(&(self.rel_type.to_lowercase(), outgoing))
+                .copied()
+        };
+        match self.direction {
+            RelationshipDirection::Outgoing => mean(self.source == middle),
+            RelationshipDirection::Incoming => mean(self.target == middle),
+            RelationshipDirection::Undirected => Some(mean(true)? + mean(false)?),
+        }
     }
 
     /// The endpoint that is not `middle`.
@@ -145,33 +157,64 @@ impl DataFusionPlanner {
         Some(qualify_column(variable, &mapping.id_field))
     }
 
-    /// Whether `expr` is `NOT (x1 = y1 AND x2 = y2)` over node ids, with `(x1, y1)` and
-    /// `(x2, y2)` being `pairs`: the predicate the logical planner adds so that two hops of
-    /// one type never bind the same edge (`distinct_relationships`).
-    fn is_distinct_edges(
+    /// Which pairs of edges `expr` excludes, if it is the predicate the logical planner adds
+    /// so that two hops of one type never bind the same edge (`distinct_relationships`):
+    /// `NOT (x1 = y1 AND x2 = y2)` over node ids, or, when a hop is undirected, an `OR` of
+    /// that and the swapped endpoints. Every endpoint is the middle node `m` or a far end,
+    /// `a` of the first hop or `c` of the second; `None` for any other predicate.
+    fn excluded_pairs(
         &self,
         ctx: &PlanningContext,
         expr: &BooleanExpression,
-        pairs: [(&str, &str); 2],
-    ) -> bool {
+        [m, a, c]: [&str; 3],
+    ) -> Option<Excluded> {
         let BooleanExpression::Not(inner) = expr else {
-            return false;
-        };
-        let BooleanExpression::And(first, second) = inner.as_ref() else {
-            return false;
+            return None;
         };
-        let is_id = |value: &ValueExpression, variable: &str| {
-            matches!(value, ValueExpression::Property(p)
-                if p.variable == variable
-                    && self.id_column(ctx, variable)
-                        == Some(qualify_column(variable, &p.property)))
+        let variable = |value: &ValueExpression| match value {
+            ValueExpression::Property(p)
+                if self.id_column(ctx, &p.variable)
+                    == Some(qualify_column(&p.variable, &p.property)) =>
+            {
+                [m, a, c].into_iter().find(|v| *v == p.variable)
+            }
+            _ => None,
         };
-        let equal = |expr: &BooleanExpression, (x, y): (&str, &str)| {
-            matches!(expr, BooleanExpression::Comparison {
-                left, operator: ComparisonOperator::Equal, right
-            } if is_id(left, x) && is_id(right, y))
+        // The nodes one `x1 = y1 AND x2 = y2` makes equal: all three, or `a` and `c`.
+        let same = |expr: &BooleanExpression| {
+            let BooleanExpression::And(first, second) = expr else {
+                return None;
+            };
+            let mut equal = Vec::new();
+            for expr in [first, second] {
+                let BooleanExpression::Comparison {
+                    left,
+                    operator: ComparisonOperator::Equal,
+                    right,
+                } = expr.as_ref()
+                else {
+                    return None;
+                };
+                let pair = [variable(left)?, variable(right)?];
+                if pair[0] != pair[1] {
+                    equal.push(pair);
+                }
+            }
+            let joins = |x: &str, y: &str| equal.iter().any(|p| p.contains(&x) && p.contains(&y));
+            match equal.len() {
+                1 if joins(a, c) => Some(Excluded::SameFarEnd),
+                2 if joins(a, m) && joins(m, c) => Some(Excluded::SelfLoops),
+                _ => None,
+            }
         };
-        equal(first, pairs[0]) && equal(second, pairs[1])
+        match inner.as_ref() {
+            BooleanExpression::Or(first, second) => {
+                let (first, second) = (same(first)?, same(second)?);
+                // Self-loops at `m` already have the same far end.
+                Some(first.max(second))
+            }
+            inner => same(inner),
+        }
     }
 
     /// `RETURN count(*)` (or `count(v)` of a pattern variable, never null) over
@@ -182,15 +225,16 @@ impl DataFusionPlanner {
     /// filter that hop's edges before grouping, and the hops' mean fan-outs at `m` say
     /// there are many more paths than edges (not so for `(i)<--(p)-->(city)`, one city
     /// per person, where grouping only adds work). When both hops have the same type the
-    /// logical planner also requires that they not have the same endpoints; for a chain
-    /// (`m` the target of one edge and the source of the other) that only excludes pairs
-    /// of self-loops at `m`, which `s1` and `s2` count. Endpoints are compared, not edge
-    /// ids, so k parallel self-loops exclude k² pairs, as on the join path.
+    /// logical planner also requires that they not have the same endpoints; for a directed
+    /// chain (`m` the target of one edge and the source of the other) that only excludes
+    /// pairs of self-loops at `m`, which `s1` and `s2` count. Endpoints are compared, not
+    /// edge ids, so k parallel self-loops exclude k² pairs, as on the join path.
     ///
-    /// Two same-type hops that share an end (`(a)-->(m)<--(c)`) exclude every pair whose
-    /// far ends are the same node `x`: `Σ_x k1(m, x) · k2(m, x)`, where `k1` and `k2` count
-    /// each hop's edges between `m` and `x`. Those hops are grouped by both ends instead,
-    /// which [`Self::shared_end_count`] sums to the count.
+    /// Two same-type hops that share an end (`(a)-->(m)<--(c)`), or where either hop is
+    /// undirected, exclude every pair whose far ends are the same node `x`:
+    /// `Σ_x k1(m, x) · k2(m, x)`, where `k1` and `k2` count each hop's edges between `m`
+    /// and `x`. Those hops are grouped by both ends instead, which
+    /// [`Self::shared_end_count`] sums to the count.
     pub(crate) fn try_build_two_hop_count(
         &self,
         ctx: &mut PlanningContext,
@@ -289,20 +333,18 @@ impl DataFusionPlanner {
             return Ok(None);
         };
 
-        let pairs = [(first.edge.0, second.edge.0), (first.edge.1, second.edge.1)];
-        let chain = first.edge.1 == second.edge.0 || first.edge.0 == second.edge.1;
+        let ends = [middle, first.other(middle), second.other(middle)];
         let (mut middle_filter, mut first_filter, mut second_filter) = (vec![], vec![], vec![]);
-        // The correction the distinct-edges predicate needs: self-loops for a chain, pairs
-        // with the same far end for hops that share an end.
+        // The correction the distinct-edges predicate needs: self-loops for a directed
+        // chain, pairs with the same far end for hops that share an end or are undirected.
         let (mut self_loops, mut shared_end) = (false, false);
         // Variables the kept filters read: a relationship no filter reads needs no edge.
         let mut read = HashSet::new();
         for predicate in predicates {
-            if self.is_distinct_edges(ctx, predicate, pairs) {
-                if chain {
-                    self_loops = true;
-                } else {
-                    shared_end = true;
+            if let Some(excluded) = self.excluded_pairs(ctx, predicate, ends) {
+                match excluded {
+                    Excluded::SelfLoops => self_loops = true,
+                    Excluded::SameFarEnd => shared_end = true,
                 }
                 continue;
             }
@@ -646,7 +688,8 @@ impl DataFusionPlanner {
         };
         let reversed = match direction {
             RelationshipDirection::Outgoing => RelationshipDirection::Incoming,
-            _ => RelationshipDirection::Outgoing,
+            RelationshipDirection::Incoming => RelationshipDirection::Outgoing,
+            RelationshipDirection::Undirected => RelationshipDirection::Undirected,
         };
         let hop = LogicalOperator::Expand {
             input: Box::new(scan.clone()),
@@ -1015,14 +1058,47 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn undirected_hops() {
+        // Counts beyond the engine's dense block, which `count` checks against the join plan.
+        let extra = |edges: &'static [(i64, i64)],
+                     likes: &'static [(i64, i64)],
+                     q: &'static str| async move {
+            count(edges, likes, q).await - count(&[], &[], q).await
+        };
+        let mixed = "MATCH (a:Person)-[:FOLLOWS]-(b:Person)-[:LIKES]->(c:Person) RETURN count(*)";
+        // 2's undirected FOLLOWS neighbours 1 and 3, times its one LIKES edge.
+        assert_eq!(extra(&[(1, 2), (2, 3)], &[(2, 4)], mixed).await, 2);
+        // An undirected self-loop is one neighbour, and different types never share an edge.
+        assert_eq!(extra(&[(2, 2)], &[(2, 2)], mixed).await, 1);
+        // Same type: a pair with the same far end would walk one edge twice.
+        let both = "MATCH (a:Person)-[:FOLLOWS]-(b:Person)-[:FOLLOWS]-(c:Person) RETURN count(*)";
+        assert_eq!(extra(&[(1, 2), (2, 3)], &[], both).await, 2);
+        assert_eq!(extra(&[(1, 2), (2, 3), (2, 2)], &[], both).await, 6);
+        // Endpoints are compared either way round, so 1→2 and 2→1 never pair up.
+        assert_eq!(extra(&[(1, 2), (2, 1), (2, 3)], &[], both).await, 4);
+        // Only 1-2-3: 3-2-3 would walk 2→3, or 3→2 against it, twice.
+        let half = "MATCH (a:Person)-[:FOLLOWS]-(b:Person)-[:FOLLOWS]->(c:Person) RETURN count(c)";
+        assert_eq!(extra(&[(1, 2), (2, 3), (3, 2)], &[], half).await, 1);
+        // At 2: three edges in from a.age > 15, four neighbours, three far ends shared.
+        let into = "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]-(c:Person) \
+                    WHERE a.age > 15 RETURN count(*)";
+        assert_eq!(extra(&[(1, 2), (3, 2), (4, 2), (2, 2)], &[], into).await, 9);
+        // The root hop, counted per node.
+        let follows = [(1, 2), (3, 2), (2, 2), (2, 4)];
+        let query = "MATCH (a:Person)-[:FOLLOWS]-(b:Person) RETURN b.id AS b, count(a) AS n";
+        assert!(grouped(&follows, query, true)
+            .await
+            .contains(&"2,4".to_string()));
+    }
+
     #[tokio::test]
     async fn other_queries_keep_the_join_plan() {
         let engine = engine(&[(1, 2), (2, 3)], &[], false).await;
         for query in [
             // A filter across the two hops.
             "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(c:Person) \
              WHERE a.age < c.age RETURN count(*)",
-            "MATCH (a:Person)-[:FOLLOWS]-(b:Person)-[:FOLLOWS]->(c:Person) RETURN count(*)",
             "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(a) RETURN count(*)",
             "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(c:Person) \
              RETURN count(DISTINCT c)",

docs/csr.md

@@ -487,9 +487,11 @@ node's neighbours without emitting a row per edge, and the planner uses that whe
 
 - `RETURN count(*)` over a two-hop chain `(a)-->(b)-->(c)` is the sum, over each middle node `b`, of its
   incoming times its outgoing neighbours, less the paths that would use one self-loop twice. Two hops of one
-  type that share an end, `(a)-->(b)<--(c)`, are counted from edges too, though not in place: the paths where
-  `a` and `c` are the same node would use one edge twice, and finding them takes grouping by both ends.
-- `RETURN m.name, count(x)` over `(x)-[:R]->(m)...`, when nothing else reads `x` or the relationship, counts
+  type that share an end, `(a)-->(b)<--(c)`, or that are undirected, `(a)--(b)--(c)`, are counted from edges
+  too, though not in place: the paths where `a` and `c` are the same node would use one edge twice, and finding
+  them takes grouping by both ends.
+- `RETURN m.name, count(x)` over `(x)-[:R]->(m)...` (in either direction, or undirected), when nothing else
+  reads `x` or the relationship, counts
   each `m`'s `x` neighbours once and plans the rest of the pattern from `m`.
 - Grouping by a node's id and its other properties groups by the id alone (see
   [grouping by a node](#grouping-by-a-node)).

docs/internals/optimizer.md

@@ -137,12 +137,17 @@ when the node table is in memory, rather than a hash join against the whole tabl
 `Σ n1(b) × n2(b)` (eager aggregation, after Yan and Larson), so the work follows the edges, not the paths. Each hop
 is planned from `b` on its own, after any filter on `b` alone, and grouped to one row per `b` before the two sides
 join on `b`'s id. When both hops have the same type, Cypher's rule that a path uses each edge once excludes the pairs
-that reuse a self-loop at `b`, so the planner subtracts `s1(b) × s2(b)`, the self-loops counted on each side. Two
-same-type hops that share an end, such as `(a)-[:R]->(b)<-[:R]-(c)`, instead exclude every pair whose far ends are
-the same node `x`, and their edges are compared by endpoints, so `k` parallel edges between `b` and `x` exclude `k²`
-pairs. Those hops are grouped by both ends, `b` and the far end, and then by `b`, which gives `n1(b)`, `n2(b)` and
-the excluded pairs `Σ k1(b, x) × k2(b, x)` in one pass over the edges; they don't count in place. It applies when every `WHERE` condition reads one hop's variables, and the hops' mean fan-outs at `b` (from the
-[join statistics](#join-statistics-for-graph-workloads)) say there are at least twice as many paths as edges.
+that reuse a self-loop at `b`, so the planner subtracts `s1(b) × s2(b)`, the self-loops counted on each side.
+
+Same-type hops that share an end, such as `(a)-[:R]->(b)<-[:R]-(c)`, or where either hop is undirected, exclude
+every pair whose far ends are the same node `x`: an undirected hop may walk its edge either way round, so the
+planner's predicate compares endpoints in both orders. Edges are compared by endpoints, not identity, so `k` parallel
+edges between `b` and `x` exclude `k²` pairs, as on the join path. These hops are grouped by both ends, `b` and the
+far end, and then by `b`, which gives `n1(b)`, `n2(b)` and the excluded pairs `Σ k1(b, x) × k2(b, x)` in one pass
+over the edges; they don't count in place. An undirected hop's mean fan-out is the sum of both directions'.
+
+The rewrite applies when every `WHERE` condition reads one hop's variables, and the hops' mean fan-outs at `b` (from
+the [join statistics](#join-statistics-for-graph-workloads)) say there are at least twice as many paths as edges.
 
 A relationship variable no condition reads is left unbound, so the CSR doesn't fetch its edge rows. Each side's
 grouping is a `HopDegreesNode` over the aggregate. After `PushCsrTargetFilters` has moved the side's far-end filter