From fa2389fd23aa236f0b4e355b21e5796014352eec Mon Sep 17 00:00:00 2001 From: prrao87 Date: Fri, 6 Nov 2026 20:19:45 -0500 Subject: [PATCH] perf(planner): count same-type two-hop paths that share an end Two hops of one type that share an end, such as (a)-[:R]->(b)<-[:R]-(c), kept the join plan in the two-hop count rewrite because the paths whose far ends are the same node would use one edge twice. Group both hops by (b, far end) and then by b, and subtract the pairs with the same far end, so the count still follows the edges rather than the paths. --- .../datafusion_planner/builder/count_ops.rs | 135 ++++++++++++++++-- docs/csr.md | 4 +- docs/internals/optimizer.md | 7 +- 3 files changed, 131 insertions(+), 15 deletions(-) diff --git a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs index b9e3485..dfd7306 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -186,8 +186,11 @@ impl DataFusionPlanner { /// (`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. - // shortcut: two same-type hops that share an end (`(a)-->(m)<--(c)`) keep the join - // plan, as their correction needs grouping by both ends; add it if that shape matters. + /// + /// 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. pub(crate) fn try_build_two_hop_count( &self, ctx: &mut PlanningContext, @@ -289,15 +292,18 @@ impl DataFusionPlanner { 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 (mut middle_filter, mut first_filter, mut second_filter) = (vec![], vec![], vec![]); - let mut self_loops = false; + // The correction the distinct-edges predicate needs: self-loops for a chain, pairs + // with the same far end for hops that share an end. + 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 { - return Ok(None); + if chain { + self_loops = true; + } else { + shared_end = true; } - self_loops = true; continue; } let mut names = HashSet::new(); @@ -381,6 +387,22 @@ impl DataFusionPlanner { // Built in the original order, so each hop takes the same relationship instance. let first_plan = self.build_operator(ctx, &from_middle(&first, first_filter))?; let second_plan = self.build_operator(ctx, &from_middle(&second, second_filter))?; + let output = item.alias.as_ref().map_or_else( + || super::super::expression::to_cypher_column_name(&item.expression), + |alias| alias.to_lowercase(), + ); + if shared_end { + return self + .shared_end_count( + [ + (first_plan, first_other.as_str()), + (second_plan, second_other.as_str()), + ], + &middle_id, + &output, + ) + .map(Some); + } let column = |name: &str| col(Column::new_unqualified(name)); let per_middle = |plan: LogicalPlan, other: &str, n: usize| -> Result { @@ -414,10 +436,6 @@ impl DataFusionPlanner { if self_loops { paths = paths - column("__s1") * column("__s2"); } - let output = item.alias.as_ref().map_or_else( - || super::super::expression::to_cypher_column_name(&item.expression), - |alias| alias.to_lowercase(), - ); let plan = LogicalPlanBuilder::from(first_plan) .join( second_plan, @@ -433,6 +451,59 @@ impl DataFusionPlanner { .map_err(|e| self.plan_error("Failed to build two-hop count", e))?; Ok(Some(plan)) } + + /// The two-hop count for same-type hops that share an end `m`, from each hop's rows + /// (`m`'s id column `middle_id`, the far end's id column): `Σ_m n1(m) · n2(m) - c(m)`, + /// where `c(m) = Σ_x k1(m, x) · k2(m, x)` counts the excluded pairs, whose two edges + /// join `m` to the same `x`. Both hops' rows go through one grouping by `(m, x)` and + /// one by `m`, so the work still follows the edges. + fn shared_end_count( + &self, + hops: [(LogicalPlan, &str); 2], + middle_id: &str, + output: &str, + ) -> Result { + let column = |name: &str| col(Column::new_unqualified(name)); + // One row per edge, marked with the hop it belongs to. + let side = |plan: LogicalPlan, other: &str, marks: [i64; 2]| { + LogicalPlanBuilder::from(plan) + .project(vec![ + column(middle_id).alias("__m"), + column(other).alias("__x"), + lit(marks[0]).alias("__k1"), + lit(marks[1]).alias("__k2"), + ])? + .build() + }; + let [(first, first_other), (second, second_other)] = hops; + let pairs = (column("__k1") * column("__k2")).alias("__c"); + LogicalPlanBuilder::from(side(first, first_other, [1, 0])?) + .union(side(second, second_other, [0, 1])?)? + .aggregate( + vec![column("__m"), column("__x")], + vec![ + sum(column("__k1")).alias("__k1"), + sum(column("__k2")).alias("__k2"), + ], + )? + .aggregate( + vec![column("__m")], + vec![ + sum(column("__k1")).alias("__n1"), + sum(column("__k2")).alias("__n2"), + sum(pairs).alias("__c"), + ], + )? + .aggregate( + Vec::::new(), + vec![sum(column("__n1") * column("__n2") - column("__c")).alias("__paths")], + )? + .project(vec![ + coalesce(vec![column("__paths"), lit(0_i64)]).alias(output) + ])? + .build() + .map_err(|e| self.plan_error("Failed to build two-hop count", e)) + } } impl DataFusionPlanner { @@ -902,6 +973,48 @@ mod tests { assert_eq!(count(&[(1, 2), (2, 3), (2, 2)], &[], reversed).await, 3); } + #[tokio::test] + async fn same_type_hops_sharing_an_end() { + let into = + "MATCH (a:Person)-[r1:FOLLOWS]->(b:Person)<-[r2:FOLLOWS]-(c:Person) RETURN count(*)"; + let out_of = + "MATCH (a:Person)<-[:FOLLOWS]-(b:Person)-[:FOLLOWS]->(c:Person) RETURN count(c)"; + // The engine's dense block adds 400 paths to each shape: each of its five middle + // nodes has two parallel edges to five far ends, 10² - 5 · 2² pairs. + const DENSE: i64 = 400; + // 1 and 3 follow 2: each pairs with the other, never with itself. + assert_eq!(count(&[(1, 2), (3, 2)], &[], into).await, DENSE + 2); + // Endpoints are compared: 1's two parallel edges to 2 exclude all four of their + // pairs, as on the join path, leaving 1 and 3 paired both ways, twice each. + assert_eq!(count(&[(1, 2), (1, 2), (3, 2)], &[], into).await, DENSE + 4); + // A self-loop at 2 is one more far end, 2 itself. + assert_eq!(count(&[(1, 2), (2, 2)], &[], into).await, DENSE + 2); + assert_eq!( + count(&[(2, 1), (2, 3), (2, 2)], &[], out_of).await, + DENSE + 6 + ); + // The shared end as the pattern's root. + let rooted = + "MATCH (b:Person)-[:FOLLOWS]->(a:Person), (b)-[:FOLLOWS]->(c:Person) RETURN count(*)"; + assert_eq!( + count(&[(2, 1), (2, 3), (2, 2)], &[], rooted).await, + DENSE + 6 + ); + // A far end that passes only one hop's filter is never excluded; the counts are + // checked against the join plan. + let follows = [(1, 2), (3, 2), (4, 2), (4, 2), (2, 2), (4, 3)]; + for filter in [ + "WHERE a.age > 15", + "WHERE a.age > 15 AND c.age < 50", + "WHERE b.age < 50 AND c.age > 25", + "WHERE r1.src <> 3", + ] { + let query = into.replace("RETURN", &format!("{filter} RETURN")); + let n = count(&follows, &[], &query).await; + assert!(n > 0, "{query}"); + } + } + #[tokio::test] async fn other_queries_keep_the_join_plan() { let engine = engine(&[(1, 2), (2, 3)], &[], false).await; @@ -909,8 +1022,6 @@ mod tests { // A filter across the two hops. "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(c:Person) \ WHERE a.age < c.age RETURN count(*)", - // Same-type hops sharing an end: the correction needs both ends. - "MATCH (a:Person)-[:FOLLOWS]->(b:Person)<-[:FOLLOWS]-(c:Person) 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) \ diff --git a/docs/csr.md b/docs/csr.md index 1efa3dc..9b1617c 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -486,7 +486,9 @@ cheap there, and a CSR hop that emits one row per edge is no faster. What the CS node's neighbours without emitting a row per edge, and the planner uses that wherever a query only counts: - `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. + 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 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 diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 439f52d..e451178 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -137,8 +137,11 @@ 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. It -applies when every `WHERE` condition reads one hop's variables, and the hops' mean fan-outs at `b` (from the +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. A relationship variable no condition reads is left unbound, so the CSR doesn't fetch its edge rows. Each side's