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 f86a90c..8f23bad 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -273,8 +273,48 @@ impl DataFusionPlanner { conjuncts(predicate, &mut predicates); op = input; } - let Some((second, op)) = Hop::of(op) else { - return Ok(None); + // `OPTIONAL MATCH (m)-[r2]-(y)` after the first hop: the second hop, planned from + // the distinct `m`, and its own `WHERE`, which applies before rows are null-extended. + let mut optional_filters = None; + let (second, op) = match op { + LogicalOperator::OptionalMatch { + input, + optional, + keys, + } if predicates.is_empty() => { + let mut own = Vec::new(); + let mut op = optional.as_ref(); + while let LogicalOperator::Filter { input, predicate } = op { + conjuncts(predicate, &mut own); + op = input; + } + let Some((second, seed)) = Hop::of(op) else { + return Ok(None); + }; + let LogicalOperator::Distinct { input: seed } = seed else { + return Ok(None); + }; + let distinct_middle = matches!(seed.as_ref(), LogicalOperator::Project { projections, .. } + if matches!(projections.as_slice(), [ProjectionItem { + expression: ValueExpression::Variable(v), .. + }] if v == second.source)); + if !distinct_middle || keys.len() != 1 { + return Ok(None); + } + optional_filters = Some(own); + let mut op = input.as_ref(); + while let LogicalOperator::Filter { input, predicate } = op { + conjuncts(predicate, &mut predicates); + op = input; + } + (second, op) + } + op => { + let Some((second, op)) = Hop::of(op) else { + return Ok(None); + }; + (second, op) + } }; let Some((first, op)) = Hop::of(op) else { return Ok(None); @@ -363,6 +403,18 @@ impl DataFusionPlanner { } read.extend(names); } + // The optional hop's own conditions filter its edges, `m` included: a middle node + // they reject keeps its rows, null-extended. + for predicate in optional_filters.iter().flatten() { + let mut names = HashSet::new(); + if referenced(predicate, &mut names).is_none() + || !names.iter().all(|n| second_vars.contains(n.as_str())) + { + return Ok(None); + } + second_filter.push((*predicate).clone()); + read.extend(names); + } // The distinct-edges predicate became the self-loop correction above, so edge // identity is needed only by a filter on the relationship; an unread one is left // unbound, and the CSR does not fetch its edge rows. @@ -474,17 +526,22 @@ impl DataFusionPlanner { let first_plan = per_middle(first_plan, &first_other, 1)?; let second_plan = per_middle(second_plan, &second_other, 2)?; - let mut paths = column("__n1") * column("__n2"); + // A middle node without an optional match still has one row per first-hop edge, + // with nulls that `count` of an optional variable skips. + let (join, second_count) = match optional_filters { + None => (JoinType::Inner, column("__n2")), + Some(_) => { + let skips_nulls = counted != middle && second_vars.contains(counted); + let unmatched = lit(if skips_nulls { 0_i64 } else { 1_i64 }); + (JoinType::Left, coalesce(vec![column("__n2"), unmatched])) + } + }; + let mut paths = column("__n1") * second_count; if self_loops { paths = paths - column("__s1") * column("__s2"); } let plan = LogicalPlanBuilder::from(first_plan) - .join( - second_plan, - JoinType::Inner, - (vec!["__m1"], vec!["__m2"]), - None, - )? + .join(second_plan, join, (vec!["__m1"], vec!["__m2"]), None)? .aggregate(Vec::::new(), vec![sum(paths).alias("__paths")])? .project(vec![ coalesce(vec![column("__paths"), lit(0_i64)]).alias(output) @@ -1092,6 +1149,41 @@ mod tests { .contains(&"2,4".to_string())); } + #[tokio::test] + async fn optional_second_hop() { + // Counts beyond the engine's dense block, which `count` checks against the join plan. + async fn extra(follows: &[(i64, i64)], likes: &[(i64, i64)], query: &str) -> i64 { + count(follows, likes, query).await - count(&[], &[], query).await + } + let base = + "MATCH (a:Person)-[:FOLLOWS]->(b:Person) OPTIONAL MATCH (b)-[:LIKES]->(c:Person)"; + // 1 and 3 follow 2, which likes 3 and 1; 1 follows 4, which likes no one. + let (follows, likes) = ([(1, 2), (3, 2), (1, 4)], [(2, 3), (2, 1)]); + for (rest, n) in [ + // 2 · 2 at 2, and 4's one row with nulls. + ("RETURN count(*)", 5), + ("RETURN count(a)", 5), + // Counting the optional variable skips the null row. + ("RETURN count(c)", 4), + // The optional hop's own WHERE filters its edges; 2 keeps one, 4 its null row. + ("WHERE c.age > 50 RETURN count(*)", 3), + // Rejecting 2 itself leaves it a null row per follower. + ("WHERE b.age > 35 RETURN count(*)", 3), + ("WHERE b.age > 35 RETURN count(c)", 0), + ] { + let query = format!("{base} {rest}"); + assert_eq!(extra(&follows, &likes, &query).await, n, "{query}"); + } + // The first clause's WHERE drops 4 before the optional match. + let query = base.replace(" OPTIONAL", " WHERE b.age < 35 OPTIONAL") + " RETURN count(*)"; + assert_eq!(extra(&follows, &likes, &query).await, 4); + // A separate clause may walk the first hop's edge again: 2's self-loop pairs with + // itself. + let query = "MATCH (a:Person)-[:FOLLOWS]->(b:Person) \ + OPTIONAL MATCH (b)-[:FOLLOWS]->(c:Person) RETURN count(*)"; + assert_eq!(extra(&[(1, 2), (2, 2)], &[], query).await, 2); + } + #[tokio::test] async fn other_queries_keep_the_join_plan() { let engine = engine(&[(1, 2), (2, 3)], &[], false).await; diff --git a/docs/csr.md b/docs/csr.md index e96511b..53bbace 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -489,7 +489,8 @@ node's neighbours without emitting a row per edge, and the planner uses that whe 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)`, 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. + them takes grouping by both ends. The second hop can be an `OPTIONAL MATCH` from `b`, where a `b` with no + match counts once per incoming edge. - `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`. diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 7e4bf37..3e11c26 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -146,6 +146,13 @@ edges between `b` and `x` exclude `k²` pairs, as on the join path. These hops a 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 second hop can also be an `OPTIONAL MATCH (b)-[:S]->(c)` from the first hop's `b`. It is grouped per `b` the +same way and left-joined, and a `b` without a match contributes `n1(b)` rows, one per first-hop edge, with `c` null: +the count is `Σ n1(b) × coalesce(n2(b), 1)`, or `coalesce(n2(b), 0)` for `count(c)`, which skips nulls. The +optional clause's own `WHERE` filters the second hop's edges, including conditions on `b` alone, since a `b` it +rejects still keeps its null row. The two clauses don't share the rule that a path uses each edge once, so no pairs +are subtracted. + 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.