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