From ec16c39e5321d8a736258df416783ec6611e9d51 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Fri, 30 Oct 2026 21:35:10 -0400 Subject: [PATCH] perf(planner): count two-hop paths from per-node degrees MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MATCH (a)-[r1]->(b)-[r2]->(c) [WHERE ...] RETURN count(*)` built every path before counting it. The paths through each middle node b number in(b) × out(b), so plan each hop on its own, group it down to one row per b, join the two and sum the products (eager aggregation). The work follows the edges, not the paths. It applies when every WHERE conjunct reads only one hop's variables, and when the engine's mean fan-out at b says the paths outnumber the edges grouped there at least twice; with one edge per node on a side (q5 of graph-benchmark, one city per person) grouping only adds work. When both hops share a type, the planner's uniqueness filter compares endpoints, so the rewrite subtracts s1(b) × s2(b) self-loop pairs, matching the join plan's counts exactly, k² for k parallel self-loops. Same-type hops that share an end keep the join plan. graph-benchmark (100K Person, 2.4M FOLLOWS), release builds, M5: q8 120 -> 20 ms, q9 92 -> 19 ms, other queries unchanged. Both LDBC suites, memory and lazy: geomean branch/main 1.014, 0.995, 1.005, 1.001, all within noise. Closes #96. --- .../datafusion_planner/builder/basic_ops.rs | 3 + .../datafusion_planner/builder/count_ops.rs | 554 ++++++++++++++++++ .../src/datafusion_planner/builder/mod.rs | 2 + .../lance-csr/src/datafusion_planner/mod.rs | 4 + crates/lance-csr/src/query.rs | 1 + python/python/tests/test_two_hop_count.py | 80 +++ 6 files changed, 644 insertions(+) create mode 100644 crates/lance-csr/src/datafusion_planner/builder/count_ops.rs create mode 100644 python/python/tests/test_two_hop_count.py diff --git a/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs index 0408c43..0f08502 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs @@ -213,6 +213,9 @@ impl DataFusionPlanner { let input_plan = self.build_operator(ctx, input)?; return self.build_simple_project(input_plan, projections); } + if let Some(plan) = self.try_build_two_hop_count(ctx, input, projections)? { + return Ok(plan); + } // An ORDER BY right before an aggregation only orders what collect() gathers, as in // Neo4j (LDBC IC7 takes `head(collect(...))` of sorted rows): it becomes the // aggregate's own ORDER BY instead of a sort of its input. diff --git a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs new file mode 100644 index 0000000..3fdaa50 --- /dev/null +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -0,0 +1,554 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! `count(*)` over a two-hop pattern, computed from per-node degrees. +//! +//! The paths through a middle node `m` number `in(m) × out(m)`, so the count is +//! `Σ_m n1(m) · n2(m)`, where `n1` and `n2` count each hop's edges at `m` that pass that +//! hop's filters (eager aggregation, Yan and Larson). Each hop is planned as a one-hop +//! pattern of its own and grouped down to one row per `m` before the join, so the work +//! follows the edges, not the paths. + +use std::collections::HashSet; + +use crate::ast::{BooleanExpression, ComparisonOperator, RelationshipDirection, ValueExpression}; +use crate::case_insensitive::qualify_column; +use crate::datafusion_planner::analysis::PlanningContext; +use crate::datafusion_planner::DataFusionPlanner; +use crate::error::Result; +use crate::logical_plan::*; +use datafusion::common::Column; +use datafusion::functions::core::expr_fn::coalesce; +use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::logical_expr::{col, lit, when, Expr, JoinType, LogicalPlan, LogicalPlanBuilder}; + +/// One `Expand` of the pattern, with its edge's endpoints in stored direction. +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), +} + +impl<'a> Hop<'a> { + fn of(op: &'a LogicalOperator) -> Option<(Self, &'a LogicalOperator)> { + let LogicalOperator::Expand { + input, + source_variable, + target_variable, + relationship_types, + direction, + relationship_variable, + .. + } = op + else { + 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, + }, + input.as_ref(), + )) + } + + fn variables(&self) -> HashSet<&'a str> { + [self.source, self.target] + .into_iter() + .chain(self.relationship) + .collect() + } + + /// Mean edges per `middle` node, from the engine's key statistics. + fn fan_out(&self, fan_out: &FanOut, middle: &str) -> Option { + let outgoing = self.edge.0 == middle; + fan_out + .get(&(self.rel_type.to_lowercase(), outgoing)) + .copied() + } + + /// The endpoint that is not `middle`. + fn other(&self, middle: &str) -> &'a str { + if self.source == middle { + self.target + } else { + self.source + } + } +} + +/// Names of the variables `expr` reads; `None` for an `EXISTS` pattern, whose variables +/// this does not track. +fn referenced(expr: &BooleanExpression, out: &mut HashSet) -> Option<()> { + match expr { + BooleanExpression::And(a, b) | BooleanExpression::Or(a, b) => { + referenced(a, out)?; + referenced(b, out) + } + BooleanExpression::Not(inner) => referenced(inner, out), + BooleanExpression::Exists(p) => { + out.insert(p.variable.clone()); + Some(()) + } + BooleanExpression::PatternExists { .. } => None, + leaf => { + crate::ast::transform_boolean_values(&mut leaf.clone(), &mut |value| match value { + ValueExpression::Variable(v) => { + out.insert(v.clone()); + } + ValueExpression::Property(p) => { + out.insert(p.variable.clone()); + } + ValueExpression::Literal(crate::ast::PropertyValue::Property(p)) => { + out.insert(p.variable.clone()); + } + _ => {} + }); + Some(()) + } + } +} + +fn conjuncts<'a>(expr: &'a BooleanExpression, out: &mut Vec<&'a BooleanExpression>) { + match expr { + BooleanExpression::And(a, b) => { + conjuncts(a, out); + conjuncts(b, out); + } + other => out.push(other), + } +} + +impl DataFusionPlanner { + /// The node id column of `variable` (`b` -> `b__id`). + fn id_column(&self, ctx: &PlanningContext, variable: &str) -> Option { + let label = ctx.analysis.var_to_label.get(variable)?; + let mapping = self.config.get_node_mapping(label)?; + 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( + &self, + ctx: &PlanningContext, + expr: &BooleanExpression, + pairs: [(&str, &str); 2], + ) -> bool { + let BooleanExpression::Not(inner) = expr else { + return false; + }; + let BooleanExpression::And(first, second) = inner.as_ref() else { + return false; + }; + 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 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)) + }; + equal(first, pairs[0]) && equal(second, pairs[1]) + } + + /// `RETURN count(*)` (or `count(v)` of a pattern variable, never null) over + /// `(x)-[r1]-(m)-[r2]-(y)`, as `Σ_m n1(m) · n2(m) - s1(m) · s2(m)` without building + /// any path. `None` (plan as usual) for any other query. + /// + /// Applies when every `WHERE` conjunct reads the variables of only one hop, so it can + /// 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. + // 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. + pub(crate) fn try_build_two_hop_count( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + projections: &[ProjectionItem], + ) -> Result> { + let [item] = projections else { + return Ok(None); + }; + let ValueExpression::AggregateFunction { + name, + args, + distinct: false, + } = &item.expression + else { + return Ok(None); + }; + if !name.eq_ignore_ascii_case("count") { + return Ok(None); + } + // `count(b)` of a node arrives as `count(b.id)`; matched node ids are never null. + let counted = match args.as_slice() { + [ValueExpression::Variable(v)] => v.as_str(), + [ValueExpression::Property(p)] + if self.id_column(ctx, &p.variable) + == Some(qualify_column(&p.variable, &p.property)) => + { + p.variable.as_str() + } + _ => return Ok(None), + }; + + let mut predicates = Vec::new(); + let mut op = input; + while let LogicalOperator::Filter { input, predicate } = op { + conjuncts(predicate, &mut predicates); + op = input; + } + let Some((second, op)) = Hop::of(op) else { + return Ok(None); + }; + let Some((first, op)) = Hop::of(op) else { + return Ok(None); + }; + let LogicalOperator::ScanByLabel { + variable: root, + label: root_label, + properties: root_properties, + } = op + else { + return Ok(None); + }; + if first.source != root.as_str() { + return Ok(None); + } + // The second hop starts from the root (`m` re-rooted) or the first hop's target. + let middle = second.source; + let (middle_label, middle_properties) = if middle == root.as_str() { + (root_label, root_properties) + } else if middle == first.target { + let LogicalOperator::Expand { + target_label, + target_properties, + .. + } = first.op + else { + unreachable!("Hop::of matched an Expand") + }; + (target_label, target_properties) + } else { + return Ok(None); + }; + let (first_vars, second_vars) = (first.variables(), second.variables()); + // A closing edge (cycle) or a relationship bound twice is not a free second hop. + if first_vars.contains(second.target) + || second.relationship.is_some_and(|r| first_vars.contains(r)) + || first.relationship.is_some_and(|r| second_vars.contains(r)) + { + return Ok(None); + } + if counted != "*" && !first_vars.contains(counted) && !second_vars.contains(counted) { + return Ok(None); + } + // Grouping pays off only when paths through `m` outnumber the edges grouped there: + // at least twice, by mean fan-out. Without statistics, keep the join plan. + let fan_out = |hop: &Hop| self.fan_out.as_ref().and_then(|f| hop.fan_out(f, middle)); + let (Some(d1), Some(d2)) = (fan_out(&first), fan_out(&second)) else { + return Ok(None); + }; + if d1 * d2 < 2.0 * (d1 + d2) { + return Ok(None); + } + let ids = [middle, first.other(middle), second.other(middle)] + .map(|variable| self.id_column(ctx, variable)); + let [Some(middle_id), Some(first_other), Some(second_other)] = ids else { + 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 (mut first_filter, mut second_filter, mut self_loops) = (vec![], vec![], false); + for predicate in predicates { + if self.is_distinct_edges(ctx, predicate, pairs) { + if !chain { + return Ok(None); + } + self_loops = true; + continue; + } + let mut names = HashSet::new(); + if referenced(predicate, &mut names).is_none() { + return Ok(None); + } + if names.iter().all(|n| first_vars.contains(n.as_str())) { + first_filter.push(predicate.clone()); + } else if names.iter().all(|n| second_vars.contains(n.as_str())) { + second_filter.push(predicate.clone()); + } else { + return Ok(None); + } + } + + let filtered = |op: LogicalOperator, predicates: Vec| match predicates + .into_iter() + .reduce(|a, b| BooleanExpression::And(Box::new(a), Box::new(b))) + { + Some(predicate) => LogicalOperator::Filter { + input: Box::new(op), + predicate, + }, + None => op, + }; + let LogicalOperator::Expand { + target_variable, + target_label, + relationship_types, + direction, + relationship_variable, + properties, + target_properties, + .. + } = second.op + else { + unreachable!("Hop::of matched an Expand") + }; + let second_op = LogicalOperator::Expand { + input: Box::new(LogicalOperator::ScanByLabel { + variable: middle.to_string(), + label: middle_label.clone(), + properties: middle_properties.clone(), + }), + source_variable: middle.to_string(), + target_variable: target_variable.clone(), + target_label: target_label.clone(), + relationship_types: relationship_types.clone(), + direction: direction.clone(), + relationship_variable: relationship_variable.clone(), + properties: properties.clone(), + target_properties: target_properties.clone(), + }; + // Built in the original order, so each hop takes the same relationship instance. + let first_plan = self.build_operator(ctx, &filtered(first.op.clone(), first_filter))?; + let second_plan = self.build_operator(ctx, &filtered(second_op, second_filter))?; + + let column = |name: &str| col(Column::new_unqualified(name)); + let per_middle = |plan: LogicalPlan, other: &str, n: usize| -> Result { + let mut aggregates = vec![count(lit(1)).alias(format!("__n{n}"))]; + if self_loops { + let is_loop = column(other).eq(column(&middle_id)); + aggregates.push( + sum(when(is_loop, lit(1_i64)).otherwise(lit(0_i64))?).alias(format!("__s{n}")), + ); + } + LogicalPlanBuilder::from(plan) + .aggregate( + vec![column(&middle_id).alias(format!("__m{n}"))], + aggregates, + )? + .build() + .map_err(|e| self.plan_error("Failed to group a hop by its middle node", e)) + }; + 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"); + 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, + JoinType::Inner, + (vec!["__m1"], vec!["__m2"]), + None, + )? + .aggregate(Vec::::new(), vec![sum(paths).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))?; + Ok(Some(plan)) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + + use crate::{CypherEngine, EngineSource, GraphConfig}; + + fn batch(columns: &[(&str, Vec)]) -> RecordBatch { + let fields = columns + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int64, false)) + .collect::>(); + let arrays = columns + .iter() + .map(|(_, values)| Arc::new(Int64Array::from(values.clone())) as _) + .collect(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap() + } + + fn edges(pairs: &[(i64, i64)]) -> RecordBatch { + batch(&[ + ("src", pairs.iter().map(|p| p.0).collect()), + ("dst", pairs.iter().map(|p| p.1).collect()), + ]) + } + + /// An engine over Persons 1–4 (ages 10, 30, 60, 40) and the given edges. Each type + /// also gets a dense block of 50 edges between other nodes, which adds no two-hop + /// path but gives the mean fan-out at which the rewrite pays off; `LIVES_IN` has one + /// edge per person, so a hop over it never does. + async fn engine(follows: &[(i64, i64)], likes: &[(i64, i64)], use_csr: bool) -> CypherEngine { + let config = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("FOLLOWS", "Person", "Person", "src", "dst") + .rel_table("LIKES", "Person", "Person", "src", "dst") + .rel_table("LIVES_IN", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let mut ids = vec![1, 2, 3, 4]; + let mut ages = vec![10, 30, 60, 40]; + ids.extend(10..50); + ages.extend(10..50); + let dense = |from: i64, to: i64| (0..50).map(move |i| (from + i % 5, to + i / 10)); + let follows: Vec<_> = follows.iter().copied().chain(dense(10, 20)).collect(); + let likes: Vec<_> = likes.iter().copied().chain(dense(30, 40)).collect(); + let lives_in: Vec<_> = (1..50).map(|i| (i, 1)).collect(); + let sources = HashMap::from([ + ( + "Person".to_string(), + EngineSource::Arrow(batch(&[("id", ids), ("age", ages)])), + ), + ("FOLLOWS".to_string(), EngineSource::Arrow(edges(&follows))), + ("LIKES".to_string(), EngineSource::Arrow(edges(&likes))), + ( + "LIVES_IN".to_string(), + EngineSource::Arrow(edges(&lives_in)), + ), + ]); + CypherEngine::new(config, sources, use_csr).await.unwrap() + } + + /// The count from `query`, checking it uses the rewrite, equals the join plan's (forced + /// by `+ 0`, which the rewrite does not match), and is the same with the CSR. + async fn count(follows: &[(i64, i64)], likes: &[(i64, i64)], query: &str) -> i64 { + let close = query.find("count(").unwrap() + + query[query.find("count(").unwrap()..].find(')').unwrap() + + 1; + let join = format!("{} + 0{}", &query[..close], &query[close..]); + let mut counts = Vec::new(); + for use_csr in [false, true] { + let engine = engine(follows, likes, use_csr).await; + for (q, rewritten) in [(query, true), (join.as_str(), false)] { + let plan = engine.explain(q, HashMap::new()).await.unwrap(); + assert_eq!(plan.contains("__paths"), rewritten, "{plan}"); + let result = engine.execute(q, HashMap::new()).await.unwrap(); + let column = result.column(0).as_any().downcast_ref::(); + counts.push(column.unwrap().value(0)); + } + } + assert!(counts.iter().all(|c| *c == counts[0]), "{counts:?}"); + counts[0] + } + + const CHAIN: &str = + "MATCH (a:Person)-[r1:FOLLOWS]->(b:Person)-[r2:FOLLOWS]->(c:Person) RETURN count(*) AS n"; + + #[tokio::test] + async fn self_loops_at_the_middle_node() { + // in(2) × out(2) = 4 includes the self-loop paired with itself. + assert_eq!(count(&[(1, 2), (2, 3), (2, 2)], &[], CHAIN).await, 3); + // Endpoints are compared, so two parallel self-loops exclude all four of their pairs. + assert_eq!(count(&[(2, 2), (2, 2)], &[], CHAIN).await, 0); + assert_eq!( + count(&[(1, 2), (2, 3), (2, 2), (2, 2)], &[], CHAIN).await, + 5 + ); + // Parallel edges between different nodes count like any others. + assert_eq!( + count(&[(1, 2), (1, 2), (2, 3), (2, 3)], &[], CHAIN).await, + 4 + ); + // No paths at all is 0, not null. + assert_eq!(count(&[(1, 2)], &[], CHAIN).await, 0); + } + + #[tokio::test] + async fn filters_decide_which_self_loops_are_subtracted() { + let follows = [(1, 2), (2, 3), (2, 2), (3, 3), (3, 4)]; + // The self-loop at 2 fails `c.age > 50` (2 is 30), so nothing is subtracted for it. + let query = CHAIN.replace("RETURN", "WHERE c.age > 50 RETURN"); + assert_eq!(count(&follows, &[], &query).await, 3); + // q9's filters, which re-root the pattern at b. + let query = CHAIN.replace("RETURN", "WHERE b.age < 50 AND c.age > 25 RETURN"); + assert_eq!(count(&follows, &[], &query).await, 3); + let query = CHAIN.replace("RETURN", "WHERE a.age < b.age AND r2.src <> 4 RETURN"); + assert_eq!(count(&follows, &[], &query).await, 4); + } + + #[tokio::test] + async fn other_two_hop_shapes() { + // Different types never bind the same edge: the self-loops pair up. + let mixed = "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:LIKES]->(c:Person) RETURN count(*)"; + assert_eq!(count(&[(2, 2)], &[(2, 2), (2, 3)], mixed).await, 2); + let shared = "MATCH (a:Person)-[:FOLLOWS]->(b:Person)<-[:LIKES]-(c:Person) RETURN count(b)"; + assert_eq!(count(&[(1, 2), (3, 2)], &[(1, 2)], shared).await, 2); + let reversed = + "MATCH (c:Person)<-[:FOLLOWS]-(b:Person)<-[:FOLLOWS]-(a:Person) RETURN count(*)"; + assert_eq!(count(&[(1, 2), (2, 3), (2, 2)], &[], reversed).await, 3); + } + + #[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(*)", + // 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) \ + RETURN count(DISTINCT c)", + "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(c:Person) \ + RETURN b.id, count(*)", + "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(c:Person) \ + RETURN count(c.age)", + "MATCH (a:Person)-[:FOLLOWS*2]->(c:Person) RETURN count(*)", + // One edge per person on one side: no more paths than edges. + "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:LIVES_IN]->(c:Person) RETURN count(*)", + ] { + let plan = engine.explain(query, HashMap::new()).await.unwrap(); + assert!(!plan.contains("__paths"), "{query}: {plan}"); + } + } +} diff --git a/crates/lance-csr/src/datafusion_planner/builder/mod.rs b/crates/lance-csr/src/datafusion_planner/builder/mod.rs index bf4f379..b559d84 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/mod.rs @@ -9,6 +9,7 @@ //! - `basic_ops`: Basic operations (filter, project, sort, limit, offset, distinct) //! - `expand_ops`: Graph traversal operations (expand, variable-length expand) //! - `aggregate_ops`: Aggregation and grouping operations +//! - `count_ops`: `count(*)` of two-hop paths from per-node degrees //! - `join_builder`: Join inference and building //! - `helpers`: Utility functions @@ -16,6 +17,7 @@ mod aggregate_ops; mod basic_ops; mod call_ops; mod comprehension_ops; +mod count_ops; mod expand_ops; mod helpers; mod join_builder; diff --git a/crates/lance-csr/src/datafusion_planner/mod.rs b/crates/lance-csr/src/datafusion_planner/mod.rs index cb910d1..7864ab7 100644 --- a/crates/lance-csr/src/datafusion_planner/mod.rs +++ b/crates/lance-csr/src/datafusion_planner/mod.rs @@ -51,6 +51,8 @@ pub struct DataFusionPlanner { pub(crate) config: GraphConfig, pub(crate) catalog: Option>, pub(crate) adjacency: Option, + /// Mean fan-out per relationship type and direction, for cost-based rewrites. + pub(crate) fan_out: Option, pub(crate) parameter_types: std::collections::HashMap, } @@ -60,6 +62,7 @@ impl DataFusionPlanner { config, catalog: None, adjacency: None, + fan_out: None, parameter_types: Default::default(), } } @@ -69,6 +72,7 @@ impl DataFusionPlanner { config, catalog: Some(catalog), adjacency: None, + fan_out: None, parameter_types: Default::default(), } } diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index 71702c3..662ec06 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -1150,6 +1150,7 @@ impl CypherQuery { // Phase 3: DataFusion Logical Plan let mut df_planner = DataFusionPlanner::with_catalog(config.clone(), catalog); df_planner.adjacency = self.adjacency.clone(); + df_planner.fan_out = self.fan_out.clone(); if parameterized { df_planner.parameter_types = self .scalar_parameters() diff --git a/python/python/tests/test_two_hop_count.py b/python/python/tests/test_two_hop_count.py new file mode 100644 index 0000000..ca5b4db --- /dev/null +++ b/python/python/tests/test_two_hop_count.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""`count(*)` of two-hop paths computed from per-node degrees. + +Each count must equal the join plan's, which `count(*) + 0` forces, on random graphs +with self-loops, parallel edges, null endpoints and ids, dangling edges, and nulls in +filtered columns. +""" + +import random + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + +CONFIG = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("FOLLOWS", "Person", "Person", "src", "dst") + .rel_table("LIKES", "Person", "Person", "src", "dst") + .build() +) + +CHAIN = "MATCH (a:Person)-[r1:FOLLOWS]->(b:Person)-[r2:FOLLOWS]->(c:Person)" +QUERIES = [ + f"{CHAIN} RETURN count(*) AS n", + f"{CHAIN} WHERE b.age < 50 AND c.age > 25 RETURN count(*) AS n", + f"{CHAIN} WHERE a.age > 30 AND r1.w <> 2 AND r2.w IS NOT NULL RETURN count(*) AS n", + f"{CHAIN} WHERE a.age < b.age OR b.age IS NULL RETURN count(c) AS n", + "MATCH (a:Person)-[:FOLLOWS]->(:Person {age: 40})-[:FOLLOWS]->(c:Person) " + "RETURN count(*) AS n", + "MATCH (c:Person)<-[:FOLLOWS]-(b:Person)<-[:FOLLOWS]-(a:Person) " + "WHERE b.age >= 20 RETURN count(*) AS n", + "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:LIKES]->(c:Person) RETURN count(*) AS n", + "MATCH (a:Person)-[:FOLLOWS]->(b:Person)<-[:LIKES]-(c:Person) " + "WHERE c.age > 30 RETURN count(b) AS n", +] + + +def graph(seed): + rng = random.Random(seed) + n = 30 + ids = list(range(n)) + [None] + people = pa.table( + { + "id": pa.array(ids, pa.int64()), + "age": [rng.choice([None, *range(0, 80, 5)]) for _ in ids], + } + ) + + def edges(count): + # Few ids, so self-loops and parallel edges are common; 99 dangles. + ends = [None, 99, *range(n)] + pairs = [(rng.choice(ends), rng.choice(ends)) for _ in range(count)] + loops = rng.sample(range(n), 8) + pairs += [(v, v) for v in loops for _ in range(rng.randint(1, 3))] + pairs += pairs[: count // 4] + return pa.table( + { + "src": pa.array([p[0] for p in pairs], pa.int64()), + "dst": pa.array([p[1] for p in pairs], pa.int64()), + "w": [rng.choice([None, 1, 2, 3]) for _ in pairs], + } + ) + + return {"Person": people, "FOLLOWS": edges(200), "LIKES": edges(100)} + + +@pytest.mark.parametrize("seed", range(5)) +@pytest.mark.parametrize("use_csr", [False, True]) +def test_two_hop_count_matches_join_plan(seed, use_csr): + engine = CypherEngine(CONFIG, graph(seed), use_csr=use_csr) + for query in QUERIES: + assert "__paths" in engine.explain(query), query + join = query.replace(" AS n", " + 0 AS n") + assert "__paths" not in engine.explain(join), join + got, want = engine.execute(query), engine.execute(join) + assert got.to_pylist() == want.to_pylist(), query + assert got.schema == want.schema, query