From fa2bd31142e7be11535b2224143d2abb04287734 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 19:37:01 -0500 Subject: [PATCH 1/7] perf(planner): leave unread relationships unbound in the two-hop count The two-hop count rewrite copied each hop's relationship variable, so a CSR hop over `[r1:FOLLOWS]` fetched edge rows nothing read. The distinct-edges predicate is already the self-loop correction, so keep the variable only when one of the hop's own filters reads it. --- .../datafusion_planner/builder/count_ops.rs | 36 +++++++++++++++++-- 1 file changed, 34 insertions(+), 2 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 03f90fe..9d8c144 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -285,6 +285,8 @@ 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 first_filter, mut second_filter, mut self_loops) = (vec![], vec![], 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 { @@ -304,7 +306,16 @@ impl DataFusionPlanner { } else { return Ok(None); } + 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. + let bound = |relationship: Option<&str>| { + relationship + .filter(|r| read.contains(*r)) + .map(str::to_string) + }; let filtered = |op: LogicalOperator, predicates: Vec| match predicates .into_iter() @@ -340,12 +351,20 @@ impl DataFusionPlanner { target_label: target_label.clone(), relationship_types: relationship_types.clone(), direction: direction.clone(), - relationship_variable: relationship_variable.clone(), + relationship_variable: bound(relationship_variable.as_deref()), properties: properties.clone(), target_properties: target_properties.clone(), }; + let mut first_op = first.op.clone(); + if let LogicalOperator::Expand { + relationship_variable, + .. + } = &mut first_op + { + *relationship_variable = bound(relationship_variable.as_deref()); + } // 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 first_plan = self.build_operator(ctx, &filtered(first_op, first_filter))?; let second_plan = self.build_operator(ctx, &filtered(second_op, second_filter))?; let column = |name: &str| col(Column::new_unqualified(name)); @@ -514,6 +533,19 @@ mod tests { assert_eq!(count(&follows, &[], &query).await, 4); } + #[tokio::test] + async fn only_relationships_a_filter_reads_fetch_edges() { + let engine = engine(&[(1, 2), (2, 3)], &[], true).await; + let plan = engine.explain(CHAIN, HashMap::new()).await.unwrap(); + assert!( + plan.contains("__paths") && !plan.contains("edge="), + "{plan}" + ); + let query = CHAIN.replace("RETURN", "WHERE r2.src <> 4 RETURN"); + let plan = engine.explain(&query, HashMap::new()).await.unwrap(); + assert_eq!(plan.matches("edge=").count(), 2, "{plan}"); + } + #[tokio::test] async fn other_two_hop_shapes() { // Different types never bind the same edge: the self-loops pair up. From 4529718d9fdcdae5910bec76a448fe9f0e12b303 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 19:38:55 -0500 Subject: [PATCH 2/7] perf(csr): stop fetching late-lookup columns in the hop below PruneCsrTargetColumns collected used columns by name over the whole plan, so `b__name` fetched by a NodeLookup after the aggregate kept the CSR hop under it taking `b__name` for every edge. Prune each lookup's input against what that input itself outputs. --- crates/lance-csr/src/csr_expand.rs | 74 +++++++++++++++---- .../lance-csr/src/csr_expand/node_lookup.rs | 8 ++ docs/internals/optimizer.md | 4 +- python/python/tests/test_engine_with.py | 5 +- 4 files changed, 75 insertions(+), 16 deletions(-) diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 6fac611..5772d4c 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -2635,30 +2635,76 @@ impl OptimizerRule for PruneCsrTargetColumns { plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> datafusion::common::Result> { - let used = used_columns(&plan)?; - plan.transform_up(|node| { + prune_unused_columns(plan) + } +} + +fn prune_unused_columns(plan: LogicalPlan) -> datafusion::common::Result> { + let used = used_columns(&plan)?; + prune_within(plan, &used) +} + +/// `plan` with the CSR and lookup columns not in `used` dropped, visiting each node once. +fn prune_within( + plan: LogicalPlan, + used: &HashSet, +) -> datafusion::common::Result> { + if let LogicalPlan::Extension(ext) = &plan { + if let Some(lookup) = ext.node.as_any().downcast_ref::() { + // The columns a lookup fetches shadow any of the same name below it (a CSR + // hop's `take` of `b__name` under `NodeLookup: take=b[b__name]`), so its input + // needs only what the input itself outputs. Find that set only when some hop + // below fetches one of them, as it takes a walk of the input. + let shadows = lookup.input().exists(|node| { + Ok(matches!(node, LogicalPlan::Extension(ext) + if ext.node.as_any().downcast_ref::().is_some_and(|csr| { + csr.added_columns() + .any(|c| lookup.fields().iter().any(|f| f.name() == c)) + }))) + })?; + let input = if shadows { + prune_unused_columns(lookup.input().clone())? + } else { + prune_within(lookup.input().clone(), used)? + }; + let rebuilt; + let lookup = if input.transformed { + rebuilt = UserDefinedLogicalNodeCore::with_exprs_and_inputs( + lookup, + vec![], + vec![input.data], + )?; + &rebuilt + } else { + lookup + }; + let extension = |node: NodeLookupNode| { + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }) + }; + return Ok(match lookup.pruned(used) { + Some(Ok(pruned)) => Transformed::yes(extension(pruned)), + // Nothing it fetches is used: drop the lookup. + Some(Err(input)) => Transformed::yes(input), + None if input.transformed => Transformed::yes(extension(lookup.clone())), + None => Transformed::no(plan), + }); + } + } + plan.map_children(|child| prune_within(child, used))? + .transform_data(|node| { if let LogicalPlan::Extension(ext) = &node { if let Some(csr) = ext.node.as_any().downcast_ref::() { - if let Some(pruned) = csr.with_target_columns(&used) { + if let Some(pruned) = csr.with_target_columns(used) { return Ok(Transformed::yes(LogicalPlan::Extension(Extension { node: Arc::new(pruned), }))); } } - if let Some(lookup) = ext.node.as_any().downcast_ref::() { - return Ok(match lookup.pruned(&used) { - Some(Ok(pruned)) => Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(pruned), - })), - // Nothing it fetches is used: drop the lookup. - Some(Err(input)) => Transformed::yes(input), - None => Transformed::no(node), - }); - } } Ok(Transformed::no(node)) }) - } } fn used_columns(plan: &LogicalPlan) -> datafusion::common::Result> { diff --git a/crates/lance-csr/src/csr_expand/node_lookup.rs b/crates/lance-csr/src/csr_expand/node_lookup.rs index 7214dae..217422b 100644 --- a/crates/lance-csr/src/csr_expand/node_lookup.rs +++ b/crates/lance-csr/src/csr_expand/node_lookup.rs @@ -110,6 +110,14 @@ impl NodeLookupNode { .map(Ok) } + pub(super) fn input(&self) -> &LogicalPlan { + &self.input + } + + pub(super) fn fields(&self) -> &[Field] { + &self.fields + } + pub(super) fn key_column(&self) -> &str { &self.key_column } diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 5fe00af..50e5baa 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -31,7 +31,9 @@ expansion was already small, which re-rooting usually ensures. DataFusion's projection pushdown prunes an extension node's inputs but not its outputs, so without this rule a fused lookup would fetch every target column for every row. The rule collects the columns used anywhere in the plan and narrows each `CsrExpand`'s target and relationship columns, and each `NodeLookup`, to that set. A `NodeLookup` whose -columns are all unused is removed. +columns are all unused is removed. The columns a `NodeLookup` fetches shadow any of the same name below it, so the rule +prunes the plan under each lookup again, against only what that plan outputs: a hop under +`NodeLookup: take=b[b__name]` doesn't fetch `b__name` for every edge just because the lookup's `b__name` is used above. "Used" is conservative: expression references, the plan's output columns, each CSR node's key and closing column, and all inputs of nodes whose column use is implicit, such as `DISTINCT` and `UNION`. diff --git a/python/python/tests/test_engine_with.py b/python/python/tests/test_engine_with.py index cb93ab7..c337041 100644 --- a/python/python/tests/test_engine_with.py +++ b/python/python/tests/test_engine_with.py @@ -107,7 +107,10 @@ def test_grouped_node_columns_fetched_after_grouping(): datasets = {"Person": PEOPLE, "KNOWS": KNOWS} csr = CypherEngine(CONFIG, datasets, use_csr=True) join = CypherEngine(CONFIG, datasets) - assert "NodeLookupExec: key=b__id" in csr.explain(query) + plan = csr.explain(query) + assert "NodeLookupExec: key=b__id" in plan + # The hop under the aggregate takes only b's id; the lookup fetches the rest. + assert "take=b[b__id*]" in plan, plan assert "NodeLookup" not in join.explain(query) by_id = lambda t: sorted(t.to_pylist(), key=lambda r: r["id"]) # noqa: E731 assert by_id(csr.execute(query)) == by_id(join.execute(query)) From 20ff5a48692c3164ecfda5b38282f7bcd4079a7f Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 19:40:45 -0500 Subject: [PATCH 3/7] perf(planner): group by a node's id when its properties are grouped too `RETURN b.id, b.name, count(a)` grouped by both columns, so the CSR hop fetched `b.name` for every edge and the aggregate hashed it. When the group keys include the id of a node with an in-memory node index, group by the id alone and fetch the node's other grouped properties after the aggregate, as `WITH b` already does. Grouping by properties without the id keeps the plain plan. --- .../builder/aggregate_ops.rs | 96 ++++++++++++++----- docs/csr.md | 9 +- docs/internals/operators.md | 18 ++-- python/python/tests/test_engine_with.py | 39 ++++++++ 4 files changed, 129 insertions(+), 33 deletions(-) diff --git a/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs index c3f70e0..25c0cbb 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use crate::ast::ValueExpression; +use crate::case_insensitive::qualify_column; use crate::csr_expand::{NodeLookupNode, NodeRows}; use crate::datafusion_planner::analysis::PlanningContext; use crate::datafusion_planner::DataFusionPlanner; @@ -110,23 +111,24 @@ fn count_values_without_ids( } impl DataFusionPlanner { - /// For `WITH b, ...` grouped by a node: when `b`'s label has an in-memory node index - /// (unique integer ids), group by the id alone and look the other columns up - /// afterwards, once per node. Rows before the grouping can outnumber nodes many times - /// over (an expansion), and carrying each node's columns through the aggregate costs - /// more than the lookup. With unique ids, grouping by the id and by all the columns are + /// For `WITH b, ...` grouped by a node, or `RETURN b.id, b.name, ...` grouped by its + /// id and other properties: when `b`'s label has an in-memory node index (unique + /// integer ids), group by the id alone and look the other columns up afterwards, + /// once per node. Rows before the grouping can outnumber nodes many times over (an + /// expansion), and carrying each node's columns through the aggregate costs more + /// than the lookup. With unique ids, grouping by the id and by all the columns are /// the same. A Lance node table keeps the plain grouping: the input rows already hold - /// the node's columns, read by an earlier step, so grouping by them reads nothing more, - /// while a lookup afterwards would read them from disk a second time. - fn late_node( + /// the node's columns, read by an earlier step, so grouping by them reads nothing + /// more, while a lookup afterwards would read them from disk a second time. + /// + /// `names` are the grouped columns of `var` (`b__id`, `b__name`); `None` unless they + /// include the id and at least one other column of the node table. + fn late_node<'a>( &self, ctx: &PlanningContext, - expr: &ValueExpression, - columns: &[Expr], + var: &str, + names: impl IntoIterator, ) -> Option { - let ValueExpression::Variable(var) = expr else { - return None; - }; let label = ctx.analysis.var_to_label.get(var)?; let rows = self.adjacency.as_ref()?.nodes(label)?.clone(); if !rows.in_memory() { @@ -136,22 +138,21 @@ impl DataFusionPlanner { let prefix = format!("{}__", var.to_lowercase()); let table = rows.schema().clone(); let mut found_id = false; - let mut others = Vec::new(); - for column in columns { - let Expr::Alias(alias) = column else { - return None; - }; - if alias.name == id { + let mut others: Vec<(String, String)> = Vec::new(); + for name in names { + if name == id { found_id = true; continue; } - let suffix = alias.name.strip_prefix(&prefix)?; + let suffix = name.strip_prefix(&prefix)?; // Only the table's own columns can be fetched again. let field = table .fields() .iter() .find(|f| f.name().eq_ignore_ascii_case(suffix))?; - others.push((alias.name.clone(), field.name().clone())); + if !others.iter().any(|(n, _)| n == name) { + others.push((name.to_string(), field.name().clone())); + } } (found_id && !others.is_empty()).then(|| LateNode { variable: var.to_lowercase(), @@ -161,6 +162,35 @@ impl DataFusionPlanner { }) } + /// [`Self::late_node`] for each variable whose properties `projections` group by, + /// unless it is also grouped whole. + fn late_properties( + &self, + ctx: &PlanningContext, + projections: &[ProjectionItem], + ) -> Vec { + let mut grouped: Vec<(String, Vec)> = Vec::new(); + for p in projections { + if let ValueExpression::Property(prop) = &p.expression { + let name = qualify_column(&prop.variable, &prop.property); + match grouped.iter_mut().find(|(v, _)| *v == prop.variable) { + Some((_, names)) => names.push(name), + None => grouped.push((prop.variable.clone(), vec![name])), + } + } + } + let whole = |var: &str| { + projections + .iter() + .any(|p| matches!(&p.expression, ValueExpression::Variable(v) if v == var)) + }; + grouped + .into_iter() + .filter(|(var, _)| !whole(var)) + .filter_map(|(var, names)| self.late_node(ctx, &var, names.iter().map(String::as_str))) + .collect() + } + pub(crate) fn build_project_with_aggregates( &self, ctx: &PlanningContext, @@ -184,11 +214,22 @@ impl DataFusionPlanner { }) .collect(); let projections = bound.as_slice(); + let late_properties = self.late_properties(ctx, projections); let mut late = Vec::new(); for p in projections { // Grouping by a node groups by all its columns (`WITH b, count(*) AS n`). if let Some(columns) = Self::entity_columns(&schema, &p.expression, None) { - if let Some(node) = self.late_node(ctx, &p.expression, &columns) { + let ValueExpression::Variable(var) = &p.expression else { + unreachable!("entity_columns matched a variable") + }; + let names: Option> = columns + .iter() + .map(|c| match c { + Expr::Alias(alias) => Some(alias.name.as_str()), + _ => None, + }) + .collect(); + if let Some(node) = names.and_then(|names| self.late_node(ctx, var, names)) { group_exprs.push(col(Column::new_unqualified(&node.id))); late.push(node); continue; @@ -226,10 +267,21 @@ impl DataFusionPlanner { } outer.push(self.unresolved_value_expr(&rewritten).alias(alias)); } else { + // A property fetched after the aggregate is not grouped by. + if let ValueExpression::Property(prop) = &p.expression { + let name = qualify_column(&prop.variable, &prop.property); + if late_properties + .iter() + .any(|node| node.columns.iter().any(|(n, _)| *n == name)) + { + continue; + } + } // Group expressions: use raw expression for grouping, no alias group_exprs.push(self.value_expr(&p.expression, &schema)?); } } + late.extend(late_properties); // After aggregation, add a projection to apply aliases to group columns let mut final_projection = Vec::new(); diff --git a/docs/csr.md b/docs/csr.md index f37a788..948d4f5 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -219,9 +219,9 @@ cache across queries or dataset versions. ### Grouping by a node -`WITH b, count(*) AS n` groups by the node `b`. When `b`'s label has a node index and its table is in -memory, the grouping uses `b`'s id alone and looks up `b`'s other properties afterwards, once per node, by the -same row positions. Rows before a grouping can outnumber nodes many times over (an expansion), and carrying a +`WITH b, count(*) AS n` groups by the node `b`, and `RETURN b.id, b.name, count(*) AS n` by its id and name. +When `b`'s label has a node index and its table is in memory, either grouping uses `b`'s id alone and looks up +`b`'s other properties afterwards, once per node, by the same row positions. Rows before a grouping can outnumber nodes many times over (an expansion), and carrying a node's properties through the aggregate for each of them costs more than one lookup per node. With unique ids, the groups are the same. `explain()` shows the lookup as `NodeLookupExec`, with only the properties the query uses: @@ -234,7 +234,8 @@ print([line.strip(" |") for line in plan.splitlines() if "NodeLookupExec" in lin # ['NodeLookupExec: key=b__id, take=b[b__name]'] ``` -This applies only when the node table is held in memory. For a Lance dataset read lazily, grouping by a node +The grouping needs the id: `RETURN b.name, count(*)` groups by the name, because two people can share one. +It applies only when the node table is held in memory. For a Lance dataset read lazily, grouping by a node still groups by all of its properties, as before. By the time the query groups by `b`, its rows already hold `b`'s properties, read from the dataset by an earlier step, so grouping by them reads nothing more. Looking them up after the grouping would read them from disk a second time, which costs more than keeping them. diff --git a/docs/internals/operators.md b/docs/internals/operators.md index 3cc285d..da7972a 100644 --- a/docs/internals/operators.md +++ b/docs/internals/operators.md @@ -128,11 +128,12 @@ has the full list. Any other hop stays a join. ## `NodeLookup`: group by id, then look up -`WITH b, count(*) AS n` groups by the node `b`, which relationally means grouping by every one of `b`'s columns. Rows -before a grouping can outnumber nodes many times over, because they come out of an expansion, so carrying `b`'s -properties through the aggregate for each of them is expensive. When `b`'s label has a node index (unique ids), the -planner groups by `b`'s id alone and adds a `NodeLookup` above the aggregate that fetches `b`'s other properties once -per group: +`WITH b, count(*) AS n` groups by the node `b`, which relationally means grouping by every one of `b`'s columns, and +`RETURN b.id, b.name, count(*)` groups by two of them. Rows before a grouping can outnumber nodes many times over, +because they come out of an expansion, so carrying `b`'s properties through the aggregate for each of them is +expensive. When `b`'s label has a node index (unique ids) and the grouping includes `b`'s id, whole or as `b.id`, the +planner groups by the id alone and adds a `NodeLookup` above the aggregate that fetches `b`'s other grouped properties +once per group: ```text MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, count(*) AS n RETURN b.name, n @@ -144,11 +145,14 @@ ProjectionExec: expr=[b__name@2 as b.name, __agg_0@1 as n] AggregateExec: mode=Partial, gby=[b__id@0 as b__id], aggr=[count(1) as __agg_0] RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 ProjectionExec: expr=[b__id@3 as b__id] - CsrExpandExec: key=a__id, edges=4, take=b[b__id*, b__name] + CsrExpandExec: key=a__id, edges=4, take=b[b__id*] ... ``` -Unlike a join back to the node table, its cost follows the number of groups, not the size of the table. It applies to +The expansion below takes only `b`'s id: the lookup's `b__name` is a column of its own, so +[`PruneCsrTargetColumns`](optimizer.md#prunecsrtargetcolumns) doesn't count it as a use of the expansion's. Grouping by +`b.name` without `b.id` keeps the plain grouping, since two nodes may share a name. Unlike a join back to the node +table, the lookup's cost follows the number of groups, not the size of the table. It applies to in-memory node tables only; see [grouping by a node](../csr.md#grouping-by-a-node) for why lazily read Lance tables keep the old plan. A [list comprehension](optimizer.md#list-comprehensions-keep-each-row-once) over a list of nodes uses it too, to bind each element to its node. diff --git a/python/python/tests/test_engine_with.py b/python/python/tests/test_engine_with.py index c337041..231584d 100644 --- a/python/python/tests/test_engine_with.py +++ b/python/python/tests/test_engine_with.py @@ -116,6 +116,45 @@ def test_grouped_node_columns_fetched_after_grouping(): assert by_id(csr.execute(query)) == by_id(join.execute(query)) +@pytest.mark.parametrize( + ("query", "lookup"), + [ + # Grouping by b's id and other properties groups by the id alone. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "RETURN b.id AS id, b.name AS name, count(a.id) AS n", + True, + ), + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "RETURN b.name AS x, a.age AS age, b.id AS id, b.name AS y, count(*) AS n", + True, + ), + ( + "MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS]->(b:Person) " + "RETURN b.id AS id, b.age AS age, count(*) AS n", + True, + ), + # Without the id, two nodes may share the other properties: plain grouping. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "RETURN b.name AS name, count(*) AS n", + False, + ), + ("MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.id AS id, count(*) AS n", False), + ], +) +def test_grouped_node_properties_fetched_after_grouping(query, lookup): + datasets = {"Person": PEOPLE, "KNOWS": KNOWS} + csr = CypherEngine(CONFIG, datasets, use_csr=True) + plan = csr.explain(query) + assert ("NodeLookupExec: key=b__id" in plan) is lookup, plan + if lookup: + assert "take=b[b__id*]" in plan, plan + rows = lambda t: sorted(map(str, t.to_pylist())) # noqa: E731 + assert rows(csr.execute(query)) == rows(CypherEngine(CONFIG, datasets).execute(query)) + + @pytest.mark.parametrize("use_csr", [False, True]) def test_unmatched_optional_node_groups_with_null_columns(use_csr): engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) From 0de7e0b91233643618c6018073968f3f92ff3706 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 19:48:26 -0500 Subject: [PATCH 4/7] perf(csr): count two-hop paths from per-node degrees on the CSR The two-hop count grouped each hop's CSR expansion, one row per edge, by its middle node in a hash aggregate. Re-root both hops at the middle node and apply filters on it before expanding. Each hop's aggregate is wrapped in a HopDegreesNode, which a new optimizer rule, CountCsrDegrees, lowers after the hop's target filter has moved into its CSR node (so after query parameters are bound): when the hop is a CSR expansion with at most a far-end filter, the expansion counts each node's neighbours in place, one row per node with its count and self-loop count. Neighbours without a node row, or failing the far-end filter, are not counted, as on the join path. Other hops keep the aggregate. --- crates/lance-csr/src/csr_expand.rs | 393 +++++++++++++++--- .../datafusion_planner/builder/count_ops.rs | 145 +++++-- crates/lance-csr/src/engine.rs | 4 +- docs/internals/index.md | 11 +- docs/internals/operators.md | 1 + docs/internals/optimizer.md | 44 ++ 6 files changed, 486 insertions(+), 112 deletions(-) diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 5772d4c..770117a 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -1490,6 +1490,9 @@ pub struct CsrExpandNode { /// equal to it are kept, replacing a filter after a full expansion. closing_column: Option, traversal: Option>, + /// When set, one row per input row with its neighbours' count and the self-loops + /// among them (see [`CsrExpandNode::degrees`]) instead of one row per neighbour. + degrees: Option<[Field; 2]>, schema: DFSchemaRef, } @@ -1524,10 +1527,66 @@ impl CsrExpandNode { edge_filter: None, closing_column: None, traversal: None, + degrees: None, schema: Arc::new(schema), }) } + /// `plan`, a hop over a fused target with at most a filter on the target above it + /// (and column-only projections), as one row per input row: the input columns, then + /// `count` (neighbours with a target row that pass the filter) and `loops` (those + /// equal to the input key). The node's degree in `plan`'s direction, without a row + /// per edge for a hash aggregate to group again. `None` for any other plan. + fn degrees(plan: &LogicalPlan, count: &str, loops: &str) -> Option { + let mut plan = PushCsrTargetFilters + .rewrite( + plan.clone(), + &datafusion::optimizer::OptimizerContext::new(), + ) + .ok()? + .data; + while let LogicalPlan::Projection(p) = &plan { + if !p.expr.iter().all(|e| matches!(e, Expr::Column(_))) { + return None; + } + plan = p.input.as_ref().clone(); + } + let LogicalPlan::Extension(ext) = &plan else { + return None; + }; + let node = ext.node.as_any().downcast_ref::()?; + let simple = node.edge.is_none() + && node.closing_column.is_none() + && node.traversal.is_none() + && node.degrees.is_none(); + let target = node.target.as_ref().filter(|_| simple)?; + let target = TargetLookup { + columns: vec![], + fields: vec![], + ..target.clone() + }; + let mut counted = node.rebuilt(node.input.clone(), Some(target), None).ok()?; + counted + .count_degrees([count, loops].map(|n| Field::new(n, DataType::Int64, false))) + .ok()?; + Some(LogicalPlan::Extension(Extension { + node: Arc::new(counted), + })) + } + + fn count_degrees(&mut self, fields: [Field; 2]) -> Result<()> { + let added = DFSchema::try_from(Schema::new(fields.to_vec())) + .map_err(|e| plan_error(format!("CsrExpand schema: {}", e)))?; + self.schema = Arc::new( + self.input + .schema() + .join(&added) + .map_err(|e| plan_error(format!("CsrExpand schema: {}", e)))?, + ); + self.degrees = Some(fields); + Ok(()) + } + /// Keep only neighbours equal to the bound target id in `column` (a closing edge). pub fn with_closing_column(mut self, column: String) -> Self { self.closing_column = Some(column); @@ -1560,6 +1619,9 @@ impl CsrExpandNode { node.edge_filter = self.edge_filter.clone(); node.closing_column = self.closing_column.clone(); node.traversal = self.traversal.clone(); + if let Some(fields) = &self.degrees { + node.count_degrees(fields.clone())?; + } Ok(node) } @@ -1600,9 +1662,13 @@ impl CsrExpandNode { self.rebuilt(self.input.clone(), target, edge).ok() } - /// Names of the columns this node adds (endpoints, then any target and edge columns). - fn added_columns(&self) -> impl Iterator { - [self.key_output.name(), self.neighbor_output.name()] + /// Names of the columns this node adds (endpoints, then any target and edge columns, + /// or the degree columns). + fn added_columns(&self) -> Box + '_> { + if let Some(fields) = &self.degrees { + return Box::new(fields.iter().map(|f| f.name())); + } + let columns = [self.key_output.name(), self.neighbor_output.name()] .into_iter() .chain( self.target @@ -1613,7 +1679,8 @@ impl CsrExpandNode { self.edge .iter() .flat_map(|e| e.fields.iter().map(|f| f.name())), - ) + ); + Box::new(columns) } } @@ -1651,6 +1718,7 @@ impl PartialEq for CsrExpandNode { (None, None) => true, _ => false, } + && self.degrees == other.degrees } } impl Eq for CsrExpandNode {} @@ -1668,6 +1736,10 @@ impl Hash for CsrExpandNode { self.target_filter.hash(state); self.edge_filter.hash(state); self.closing_column.hash(state); + self.degrees + .as_ref() + .map(|f| [f[0].name(), f[1].name()]) + .hash(state); } } impl PartialOrd for CsrExpandNode { @@ -1775,6 +1847,9 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { if let Some(walk) = &self.traversal { write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; } + if let Some([count, loops]) = &self.degrees { + write!(f, ", degrees=[{}, {}]", count.name(), loops.name())?; + } Ok(()) } @@ -1806,6 +1881,8 @@ pub struct CsrExpandExec { target_state: Arc, closing_column: Option, traversal: Option>, + /// Count each input row's neighbours instead of emitting them. + degrees: bool, schema: SchemaRef, properties: Arc, } @@ -1822,6 +1899,7 @@ impl CsrExpandExec { edge_filter: Option, closing_column: Option, traversal: Option>, + degrees: bool, schema: SchemaRef, ) -> Self { let properties = Arc::new(PlanProperties::new( @@ -1842,6 +1920,7 @@ impl CsrExpandExec { target_state: Arc::default(), closing_column, traversal, + degrees, schema, properties, } @@ -1860,6 +1939,8 @@ struct BatchExpansion { edge_rows: bool, /// Passing edge rows of the relationship filter, once built: others are skipped. edge_mask: Option>, + /// Per input row, in degree mode: neighbours kept, and those equal to the key. + degrees: Option<(Vec, Vec)>, } impl BatchExpansion { @@ -1882,9 +1963,31 @@ impl BatchExpansion { walk: None, edge_rows: false, edge_mask: None, + degrees: None, }) } + /// In degree mode, the input rows with at least one neighbour, each with its two + /// counts; `None` otherwise or when no row has one. + fn degree_batch(self, schema: &SchemaRef) -> datafusion::common::Result> { + let Some((counts, loops)) = self.degrees else { + return Ok(None); + }; + let keep = arrow::array::BooleanArray::from_iter(counts.iter().map(|&c| Some(c > 0))); + if keep.true_count() == 0 { + return Ok(None); + } + let mut columns = self + .batch + .columns() + .iter() + .map(|c| filter_array(c.as_ref(), &keep)) + .collect::, _>>()?; + columns.push(filter_array(&Int64Array::from(counts), &keep)?); + columns.push(filter_array(&Int64Array::from(loops), &keep)?); + Ok(Some(RecordBatch::try_new(schema.clone(), columns)?)) + } + /// Collect only the next chunk; target properties are fetched after this returns. async fn next_chunk( &mut self, @@ -2072,63 +2175,19 @@ async fn materialize( .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?; expanded.retain(&filter.keep(columns)?); } + let reused = match (target, target_filter) { + (Some(target), Some(filter)) => { + filter_targets(target, scanned, filter, &mut expanded).await? + } + _ => vec![], + }; let ExpandedRows { - mut rows, - mut key_values, - mut neighbor_values, - mut edge_rows, - mut target_positions, + rows, + key_values, + neighbor_values, + edge_rows, + target_positions, } = expanded; - let mut reused = Vec::new(); - if let (Some(target), Some(filter)) = (target, target_filter) { - // Take only the filter's columns, evaluate, and keep passing rows before - // taking anything else. Null counts as false, as in a WHERE clause. - let fetch: Vec<_> = filter - .columns - .iter() - .copied() - .filter(|&c| c != target.rows.id_index) - .collect(); - let mut taken = take_target(target, scanned, &target_positions, &fetch) - .await? - .into_iter(); - let arrays = filter - .columns - .iter() - .enumerate() - .map(|(i, &c)| { - if c == target.rows.id_index { - cast( - &Int64Array::from(neighbor_values.clone()), - filter.schema.field(i).data_type(), - ) - } else { - Ok(taken - .next() - .expect("one fetched array per non-id filter column")) - } - }) - .collect::, _>>()?; - let candidates = RecordBatch::try_new(filter.schema.clone(), arrays)?; - let mask = filter - .predicate - .evaluate(&candidates)? - .into_array(candidates.num_rows())?; - let mask = mask.as_boolean(); - let keep: Vec = (0..mask.len()) - .map(|i| mask.is_valid(i) && mask.value(i)) - .collect(); - for (i, &column) in filter.columns.iter().enumerate() { - if column != target.rows.id_index && target.columns.contains(&column) { - reused.push((column, filter_array(candidates.column(i), mask)?)); - } - } - retain_by_mask(&mut rows, &keep); - retain_by_mask(&mut key_values, &keep); - retain_by_mask(&mut neighbor_values, &keep); - retain_by_mask(&mut edge_rows, &keep); - retain_by_mask(&mut target_positions, &keep); - } let rows = UInt32Array::from(rows); let mut columns: Vec = batch .columns() @@ -2175,6 +2234,60 @@ async fn materialize( Ok(RecordBatch::try_new(schema.clone(), columns)?) } +/// Keep the rows of `expanded` whose target passes the moved target filter. Takes only +/// the filter's columns; returns those the target also outputs, for the rows kept. +/// Null counts as false, as in a WHERE clause. +async fn filter_targets( + target: &TargetLookup, + scanned: Option<&TargetColumns>, + filter: &CompiledFilter, + expanded: &mut ExpandedRows, +) -> datafusion::common::Result> { + let fetch: Vec<_> = filter + .columns + .iter() + .copied() + .filter(|&c| c != target.rows.id_index) + .collect(); + let mut taken = take_target(target, scanned, &expanded.target_positions, &fetch) + .await? + .into_iter(); + let arrays = filter + .columns + .iter() + .enumerate() + .map(|(i, &c)| { + if c == target.rows.id_index { + cast( + &Int64Array::from(expanded.neighbor_values.clone()), + filter.schema.field(i).data_type(), + ) + } else { + Ok(taken + .next() + .expect("one fetched array per non-id filter column")) + } + }) + .collect::, _>>()?; + let candidates = RecordBatch::try_new(filter.schema.clone(), arrays)?; + let mask = filter + .predicate + .evaluate(&candidates)? + .into_array(candidates.num_rows())?; + let mask = mask.as_boolean(); + let keep: Vec = (0..mask.len()) + .map(|i| mask.is_valid(i) && mask.value(i)) + .collect(); + let mut reused = Vec::new(); + for (i, &column) in filter.columns.iter().enumerate() { + if column != target.rows.id_index && target.columns.contains(&column) { + reused.push((column, filter_array(candidates.column(i), mask)?)); + } + } + expanded.retain(&keep); + Ok(reused) +} + /// Keep the elements of `values` whose position in `keep` is true. fn retain_by_mask(values: &mut Vec, keep: &[bool]) { let mut keep = keep.iter(); @@ -2210,6 +2323,10 @@ impl DisplayAs for CsrExpandExec { if let Some(walk) = &self.traversal { write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; } + if self.degrees { + let names: Vec<_> = self.schema.fields().iter().rev().take(2).collect(); + write!(f, ", degrees=[{}, {}]", names[1].name(), names[0].name())?; + } Ok(()) } } @@ -2231,7 +2348,12 @@ impl ExecutionPlan for CsrExpandExec { if self.traversal.is_some() { return Ok(Arc::new(stats)); } - let degree = self.adjacency.mean_degree(); + // At most one row per input row when counting. + let degree = if self.degrees { + 1.0 + } else { + self.adjacency.mean_degree() + }; stats.num_rows = match input.num_rows { Precision::Exact(n) | Precision::Inexact(n) => { Precision::Inexact((n as f64 * degree).ceil() as usize) @@ -2267,6 +2389,7 @@ impl ExecutionPlan for CsrExpandExec { self.edge_filter.clone(), self.closing_column.clone(), self.traversal.clone(), + self.degrees, self.schema.clone(), ))) } @@ -2324,6 +2447,7 @@ impl ExecutionPlan for CsrExpandExec { let schema = self.schema.clone(); let out_schema = schema.clone(); let traversal = self.traversal.clone(); + let degrees = self.degrees; // Keep one input and one materialization in flight per partition. A hub's // remaining neighbours stay in the immutable adjacency until the next poll. let stream = futures::stream::try_unfold((input, None), move |(mut input, mut pending)| { @@ -2346,6 +2470,8 @@ impl ExecutionPlan for CsrExpandExec { let mut expansion = BatchExpansion::try_new(batch?, key_index, closing_index)?; expansion.edge_rows = edge.is_some(); + let rows = expansion.keys.len(); + expansion.degrees = degrees.then(|| (vec![0; rows], vec![0; rows])); pending = Some(expansion); } let expansion = pending.as_mut().expect("input batch just set"); @@ -2370,7 +2496,10 @@ impl ExecutionPlan for CsrExpandExec { } }; if chunk.rows.is_empty() { - pending = None; + let done = pending.take().expect("input batch just set"); + if let Some(batch) = done.degree_batch(&schema)? { + return Ok(Some((batch, (input, None)))); + } continue; } let mut chunk = chunk; @@ -2415,6 +2544,20 @@ impl ExecutionPlan for CsrExpandExec { } } let scanned = target_state.columns.get().cloned().flatten(); + if let Some((counts, loops)) = expansion.degrees.as_mut() { + if let (Some(target), Some(compiled)) = + (target.as_deref(), target_filter.as_deref()) + { + filter_targets(target, scanned.as_deref(), compiled, &mut chunk) + .await?; + } + let ends = chunk.key_values.iter().zip(&chunk.neighbor_values); + for (&row, (&key, &neighbor)) in chunk.rows.iter().zip(ends) { + counts[row as usize] += 1; + loops[row as usize] += i64::from(key == neighbor); + } + continue; + } let batch = materialize( &expansion.batch, chunk, @@ -2480,6 +2623,10 @@ impl ExtensionPlanner for CsrExtensionPlanner { if let Some(node) = node.as_any().downcast_ref::() { return Ok(Some(path_exists::plan_path_exists(node, physical_inputs))); } + if node.as_any().is::() { + // Not lowered by `CountCsrDegrees`: run the aggregate it wraps. + return Ok(Some(physical_inputs[0].clone())); + } let Some(node) = node.as_any().downcast_ref::() else { return Ok(None); }; @@ -2501,6 +2648,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { edge_filter, node.closing_column.clone(), node.traversal.clone(), + node.degrees.is_some(), Arc::new(node.schema.as_arrow().clone()), )))) } @@ -2614,6 +2762,129 @@ impl OptimizerRule for PushCsrTargetFilters { } } +/// A hop's per-node count: an aggregate grouping a hop by its key, `count(*)` as `count` +/// and, optionally, the rows whose other end is the key as `loops`. [`CountCsrDegrees`] +/// replaces it with [`CsrExpandNode::degrees`] when the hop qualifies, and otherwise +/// with the aggregate. It is lowered in the optimizer rather than when planning, as the +/// hop's target filter may hold query parameters, bound only after planning. +#[derive(Clone, PartialEq, Eq, Hash, PartialOrd)] +pub struct HopDegreesNode { + aggregate: LogicalPlan, + key: String, + count: String, + loops: String, +} + +impl HopDegreesNode { + pub fn new(aggregate: LogicalPlan, key: &str, count: &str, loops: &str) -> Self { + Self { + aggregate, + key: key.to_string(), + count: count.to_string(), + loops: loops.to_string(), + } + } + + /// The degree plan with the aggregate's output columns, if the hop qualifies. + fn lowered(&self) -> Option { + let LogicalPlan::Aggregate(aggregate) = &self.aggregate else { + return None; + }; + let degrees = CsrExpandNode::degrees(&aggregate.input, &self.count, &self.loops)?; + let column = |name: &str| Expr::Column(Column::new_unqualified(name)); + let fields = self.aggregate.schema().fields(); + let mut columns = vec![column(&self.key).alias(fields[0].name())]; + columns.extend(fields[1..].iter().map(|f| column(f.name()))); + LogicalPlanBuilder::from(degrees) + .project(columns) + .and_then(|b| b.build()) + .ok() + } +} + +impl std::fmt::Debug for HopDegreesNode { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + UserDefinedLogicalNodeCore::fmt_for_explain(self, f) + } +} + +impl UserDefinedLogicalNodeCore for HopDegreesNode { + fn name(&self) -> &str { + "HopDegrees" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.aggregate] + } + + fn schema(&self) -> &DFSchemaRef { + self.aggregate.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + /// Keeps filters above, so that the aggregate stays the node's input. + fn prevent_predicate_push_down_columns(&self) -> HashSet { + self.aggregate + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect() + } + + fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option>> { + Some(vec![output_columns.to_vec()]) + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "HopDegrees: key={}", self.key) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + inputs: Vec, + ) -> datafusion::common::Result { + let aggregate = inputs.into_iter().next().expect("HopDegrees has one input"); + Ok(Self { + aggregate, + ..self.clone() + }) + } +} + +/// Optimizer rule: lower each [`HopDegreesNode`], after [`PushCsrTargetFilters`] has +/// moved the hop's target filter into its CSR node. +#[derive(Debug, Default)] +pub struct CountCsrDegrees; + +impl OptimizerRule for CountCsrDegrees { + fn name(&self) -> &str { + "count_csr_degrees" + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> datafusion::common::Result> { + plan.transform_up(|node| { + let LogicalPlan::Extension(ext) = &node else { + return Ok(Transformed::no(node)); + }; + let Some(hop) = ext.node.as_any().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + Ok(Transformed::yes( + hop.lowered().unwrap_or_else(|| hop.aggregate.clone()), + )) + }) + } +} + /// Optimizer rule: drop fused target columns, and node lookup columns, that nothing /// in the plan uses. /// 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 9d8c144..d04b91f 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -6,13 +6,15 @@ //! 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. +//! pattern of its own from `m` and grouped down to one row per `m` before the join, so +//! the work follows the edges, not the paths. On the CSR the grouping is the expansion +//! itself: it counts each `m`'s neighbours in place (`HopDegreesNode`). use std::collections::HashSet; use crate::ast::{BooleanExpression, ComparisonOperator, RelationshipDirection, ValueExpression}; use crate::case_insensitive::qualify_column; +use crate::csr_expand::HopDegreesNode; use crate::datafusion_planner::analysis::PlanningContext; use crate::datafusion_planner::DataFusionPlanner; use crate::error::Result; @@ -20,7 +22,9 @@ 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}; +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. struct Hop<'a> { @@ -284,7 +288,8 @@ 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 first_filter, mut second_filter, mut self_loops) = (vec![], vec![], false); + let (mut middle_filter, mut first_filter, mut second_filter) = (vec![], vec![], vec![]); + let mut self_loops = false; // Variables the kept filters read: a relationship no filter reads needs no edge. let mut read = HashSet::new(); for predicate in predicates { @@ -299,7 +304,9 @@ impl DataFusionPlanner { if referenced(predicate, &mut names).is_none() { return Ok(None); } - if names.iter().all(|n| first_vars.contains(n.as_str())) { + if names.iter().all(|n| n == middle) { + middle_filter.push(predicate.clone()); + } else 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()); @@ -327,62 +334,78 @@ impl DataFusionPlanner { }, 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 { + // Each hop is re-rooted at `m`, after the filters on `m` alone, so that on the CSR + // it reads `m`'s neighbours in one direction and can count them in place. + let from_middle = |hop: &Hop, predicates: Vec| { + let LogicalOperator::Expand { + target_label, + relationship_types, + direction, + relationship_variable, + properties, + target_properties, + .. + } = hop.op + else { + unreachable!("Hop::of matched an Expand") + }; + let (other_label, other_properties, direction) = if hop.source == middle { + (target_label, target_properties, direction.clone()) + } else { + // Only the first hop ends at `m`, and its source is then the root. + let reversed = match direction { + RelationshipDirection::Outgoing => RelationshipDirection::Incoming, + RelationshipDirection::Incoming => RelationshipDirection::Outgoing, + RelationshipDirection::Undirected => RelationshipDirection::Undirected, + }; + (root_label, root_properties, reversed) + }; + let scan = 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: bound(relationship_variable.as_deref()), - properties: properties.clone(), - target_properties: target_properties.clone(), + }; + let expand = LogicalOperator::Expand { + input: Box::new(filtered(scan, middle_filter.clone())), + source_variable: middle.to_string(), + target_variable: hop.other(middle).to_string(), + target_label: other_label.clone(), + relationship_types: relationship_types.clone(), + direction, + relationship_variable: bound(relationship_variable.as_deref()), + properties: properties.clone(), + target_properties: other_properties.clone(), + }; + filtered(expand, predicates) }; - let mut first_op = first.op.clone(); - if let LogicalOperator::Expand { - relationship_variable, - .. - } = &mut first_op - { - *relationship_variable = bound(relationship_variable.as_deref()); - } // Built in the original order, so each hop takes the same relationship instance. - let first_plan = self.build_operator(ctx, &filtered(first_op, first_filter))?; - let second_plan = self.build_operator(ctx, &filtered(second_op, second_filter))?; + 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 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}"))]; + let (paths, loops) = (format!("__n{n}"), format!("__s{n}")); + let mut aggregates = vec![count(lit(1)).alias(&paths)]; 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}")), - ); + aggregates + .push(sum(when(is_loop, lit(1_i64)).otherwise(lit(0_i64))?).alias(&loops)); } - LogicalPlanBuilder::from(plan) + let aggregate = 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)) + .map_err(|e| self.plan_error("Failed to group a hop by its middle node", e))?; + if self.adjacency.is_none() { + return Ok(aggregate); + } + // On the CSR, the optimizer may count each `m`'s neighbours in place instead. + let hop = HopDegreesNode::new(aggregate, &middle_id, &paths, &loops); + Ok(LogicalPlan::Extension(Extension { + node: std::sync::Arc::new(hop), + })) }; let first_plan = per_middle(first_plan, &first_other, 1)?; let second_plan = per_middle(second_plan, &second_other, 2)?; @@ -546,6 +569,38 @@ mod tests { assert_eq!(plan.matches("edge=").count(), 2, "{plan}"); } + #[tokio::test] + async fn csr_hops_count_degrees_in_place() { + // Edges to and from 99, which has no Person row, match nothing on either path. + let dangling = [(1, 2), (2, 3), (2, 99), (99, 2), (2, 2)]; + assert_eq!(count(&dangling, &[], CHAIN).await, 3); + let engine = engine(&dangling, &[], true).await; + for (filter, degrees) in [ + ("", 2), + // Filters on the far ends and the middle node count in place. + ("WHERE b.age < 50 AND c.age > 25 AND a.age > 5", 2), + // A filter across a hop's ends, or on its edge, needs that hop's rows. + ("WHERE a.age < b.age", 1), + ("WHERE r2.src <> 4", 1), + ] { + let query = CHAIN.replace("RETURN", &format!("{filter} RETURN")); + let plan = engine.explain(&query, HashMap::new()).await.unwrap(); + assert_eq!(plan.matches("CsrExpandExec").count(), 2, "{plan}"); + // Lowered in the optimizer: only the physical plan counts degrees. + assert_eq!(plan.matches("degrees=").count(), degrees, "{query}: {plan}"); + } + // A far-end filter with a parameter, bound into a cached plan after planning. + let literal = CHAIN.replace("RETURN", "WHERE c.age > 25 RETURN"); + assert_eq!(count(&dangling, &[], &literal).await, 3); + let query = CHAIN.replace("RETURN", "WHERE c.age > $age RETURN"); + let params = HashMap::from([("age".to_string(), serde_json::Value::from(25))]); + for _ in 0..2 { + let result = engine.execute(&query, params.clone()).await.unwrap(); + let column = result.column(0).as_any().downcast_ref::(); + assert_eq!(column.unwrap().value(0), 3); + } + } + #[tokio::test] async fn other_two_hop_shapes() { // Different types never bind the same edge: the self-loops pair up. diff --git a/crates/lance-csr/src/engine.rs b/crates/lance-csr/src/engine.rs index de01402..d30d9ec 100644 --- a/crates/lance-csr/src/engine.rs +++ b/crates/lance-csr/src/engine.rs @@ -48,7 +48,8 @@ use serde_json::Value as JsonValue; use crate::config::GraphConfig; use crate::csr_expand::{ edge_changes, edge_changes_since, is_supported_id_type, Adjacency, AdjacencyCatalog, - CsrQueryPlanner, EdgeChanges, EdgeRows, NodeRows, PruneCsrTargetColumns, PushCsrTargetFilters, + CountCsrDegrees, CsrQueryPlanner, EdgeChanges, EdgeRows, NodeRows, PruneCsrTargetColumns, + PushCsrTargetFilters, }; use crate::csr_store; use crate::error::{GraphError, Result}; @@ -521,6 +522,7 @@ async fn build_engine_state( // Move target-only filters into CSR nodes first, so pruning can then drop // the filter columns from the take. .with_optimizer_rule(Arc::new(PushCsrTargetFilters)) + .with_optimizer_rule(Arc::new(CountCsrDegrees)) .with_optimizer_rule(Arc::new(PruneCsrTargetColumns)) .build(), ); diff --git a/docs/internals/index.md b/docs/internals/index.md index abb3059..5700e8e 100644 --- a/docs/internals/index.md +++ b/docs/internals/index.md @@ -65,7 +65,7 @@ flowchart TD S --> R["Graph rewrites
pattern_order.rs
pattern_alternatives.rs"] R --> G["Graph logical plan
logical_plan.rs
ScanByLabel, Expand, VariableLengthExpand,
ShortestPath, PathMatch, OptionalMatch, Apply"] G --> D["DataFusion logical plan
datafusion_planner/
scans + joins + extension nodes"] - D --> O["DataFusion optimizer
+ PushCsrTargetFilters
+ PruneCsrTargetColumns"] + D --> O["DataFusion optimizer
+ PushCsrTargetFilters
+ CountCsrDegrees
+ PruneCsrTargetColumns"] O --> X["Physical planning
CsrQueryPlanner → CsrExtensionPlanner"] X --> E["Execution
CsrExpandExec, NodeLookupExec, ShortestPathExec,
PathMatchExec, ApplyExec, KeyStatisticsExec"] E --> B["Arrow record batches"] @@ -80,19 +80,20 @@ DataFusion logical plan and the physical plan, and every custom operator shows u | DataFusion trait | lance-csr implementation | Role | | --- | --- | --- | -| `UserDefinedLogicalNodeCore` | `CsrExpandNode`, `NodeLookupNode`, `ShortestPathNode`, `PathMatchNode`, `ApplyNode` | Graph operators as logical plan nodes | +| `UserDefinedLogicalNodeCore` | `CsrExpandNode`, `NodeLookupNode`, `ShortestPathNode`, `PathMatchNode`, `ApplyNode`, `HopDegreesNode` | Graph operators as logical plan nodes | | `ExecutionPlan` | `CsrExpandExec`, `NodeLookupExec`, `ShortestPathExec`, `PathMatchExec`, `ApplyExec`, `KeyStatisticsExec` | Their physical operators | | `ExtensionPlanner` | `CsrExtensionPlanner` | Lowers each graph node to its `ExecutionPlan` | | `QueryPlanner` | `CsrQueryPlanner` | Installs the extension planner on the session | -| `OptimizerRule` | `PushCsrTargetFilters`, `PruneCsrTargetColumns` | Filter pushdown into, and column pruning of, extension nodes | +| `OptimizerRule` | `PushCsrTargetFilters`, `CountCsrDegrees`, `PruneCsrTargetColumns` | Filter pushdown into extension nodes, per-node counts on the CSR, and column pruning | | `PhysicalOptimizerRule` | `DeadlineJoins` (`query_deadline_joins`) | Cooperative deadlines in nested-loop and cross joins | | `TableProvider` | `KeyStatisticsProvider` | Reports distinct counts of join keys so DataFusion sizes graph joins correctly | | `ScalarUDFImpl` | `cypher_list`, `cypher_field`, Cypher math and scalar functions; vector distance and similarity functions (from [lance-graph](https://github.com/lance-format/lance-graph)) | Cypher semantics that DataFusion's built-ins don't match | Apart from the vector functions, every implementation in this table was added in lance-csr; [lance-graph](https://github.com/lance-format/lance-graph) used no DataFusion extension points. `CypherEngine` builds -one DataFusion session per snapshot (in `engine.rs`). It installs `CsrQueryPlanner` and adds `PushCsrTargetFilters` -before `PruneCsrTargetColumns`, so that pruning can drop the columns a moved filter reads. Rust callers that bring +one DataFusion session per snapshot (in `engine.rs`). It installs `CsrQueryPlanner` and adds `PushCsrTargetFilters`, +then `CountCsrDegrees`, then `PruneCsrTargetColumns`: counting in place needs the hop's filter moved in first, and +pruning can then drop the columns a moved filter reads. Rust callers that bring their own `SessionContext` must install `CsrQueryPlanner` themselves; contexts created by `CypherQuery` and `CypherEngine` already have it. diff --git a/docs/internals/operators.md b/docs/internals/operators.md index da7972a..8638d90 100644 --- a/docs/internals/operators.md +++ b/docs/internals/operators.md @@ -80,6 +80,7 @@ keeps the two hops from reusing one edge. Both plans start from `c` (Dan), not ` | Relationship filter: drops failing edges before anything else is fetched | `edge_filter=since > 2020` | `FilterExec` after the join | | Closing edge: the target is already bound (a cycle), so only neighbours equal to it are kept | `match=b__id` | Expanding everything, then filtering | | [Iterative variable-length traversal](../csr.md#variable-length-traversal): one cursor per hop over typed (label, key) frontiers, with bounded output batches | `iterative hops=1..3` | One join branch per path length, combined with `UNION ALL` | +| Degrees: one row per input node with the number of its neighbours that pass the target filter, and how many of them are the node itself, for the [two-hop count](optimizer.md#two-hop-counts-from-degrees) | `degrees=[__n1, __s1]` | A hash aggregate over one row per edge | A hop that uses several of these at once: diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 50e5baa..bbd4ed0 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -106,6 +106,50 @@ Two details keep this from repeating each row's data once per element: A list of nodes binds each element to its node through a [`NodeLookup`](operators.md#nodelookup-group-by-id-then-look-up) when the node table is in memory, rather than a hash join against the whole table. +### Two-hop counts from degrees + +`MATCH (a)-[:R]->(b)-[:S]->(c) RETURN count(*)` counts paths, and the paths through a middle node `b` number +`n1(b) × n2(b)`: its neighbours on each side that pass that side's filters. `count_ops.rs` plans the count as +`Σ 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 +[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 +into its `CsrExpand`, the `CountCsrDegrees` rule replaces it: a CSR hop with at most that filter counts each node's +neighbours in place ([`degrees=`](operators.md#what-one-expansion-can-do)), one row per node, and any other hop keeps +the aggregate. The rule runs in the optimizer, not while planning, because the far-end filter may hold query +parameters, and those are bound into the cached plan template after it is planned. On a graph of 100,000 people and 2.4 +million `FOLLOWS` edges: + +```text +MATCH (a:Person)-[r1:FOLLOWS]->(b:Person)-[r2:FOLLOWS]->(c:Person) +WHERE b.age < $age_1 AND c.age > $age_2 RETURN count(*) AS numpaths + +ProjectionExec: expr=[CASE WHEN __paths@0 IS NOT NULL THEN __paths@0 ELSE 0 END as numpaths] + AggregateExec: mode=Final, gby=[], aggr=[sum(__n1 * __n2 - __s1 * __s2) as __paths] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[sum(__n1 * __n2 - __s1 * __s2) as __paths] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(__m1@0, __m2@0)], projection=[__n1@1, __s1@2, __n2@4, __s2@5] + RepartitionExec: partitioning=Hash([__m1@0], 10), input_partitions=10 + ProjectionExec: expr=[b__id@0 as __m1, __n1@1 as __n1, __s1@2 as __s1] + CsrExpandExec: key=b__id, edges=2417738, take=a[], degrees=[__n1, __s1] + ProjectionExec: expr=[id@0 as b__id] + FilterExec: age@1 < 50, projection=[id@0] + KeyStatisticsExec: distinct=[id=100000] + DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] + RepartitionExec: partitioning=Hash([__m2@0], 10), input_partitions=10 + ProjectionExec: expr=[b__id@0 as __m2, __n2@1 as __n2, __s2@2 as __s2] + CsrExpandExec: key=b__id, edges=2417738, take=c[], filter=age@0 > 25, degrees=[__n2, __s2] + ProjectionExec: expr=[id@0 as b__id] + FilterExec: age@1 < 50, projection=[id@0] + KeyStatisticsExec: distinct=[id=100000] + DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] +``` + ## Join statistics for graph workloads Not every hop runs on the CSR, and the joins that remain need good cardinality estimates. DataFusion estimates an From 7e8fd4903261122d7456fd16cce19454173538b1 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 20:06:55 -0500 Subject: [PATCH 5/7] perf(planner): count the root hop per node before the rest of the pattern `MATCH (x)-[:FOLLOWS]->(m)-[:LIVES_IN]->(c) RETURN m.name, c.city, count(x)` carried one row per FOLLOWS edge through the second hop and into the aggregate. On the CSR those rows come out spread over every partition, so each partial aggregate built a table of every group, keyed here on four strings: about 3x the join path's time. When every aggregate is a plain count of rows and nothing else reads the root variable or its relationship, count each m's root neighbours once (a HopDegreesNode, counted in place on the CSR), plan the rest of the pattern from m, join the counts on m's id and sum them. Conjuncts over the root alone filter the counted neighbours. The join path keeps its plan. --- .../datafusion_planner/builder/basic_ops.rs | 3 + .../datafusion_planner/builder/count_ops.rs | 289 ++++++++++++++++++ docs/internals/optimizer.md | 40 +++ python/python/tests/test_engine_csr.py | 6 +- python/python/tests/test_engine_with.py | 23 +- 5 files changed, 355 insertions(+), 6 deletions(-) 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 c235aa1..e7c11b4 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs @@ -216,6 +216,9 @@ impl DataFusionPlanner { if let Some(plan) = self.try_build_two_hop_count(ctx, input, projections)? { return Ok(plan); } + if let Some(plan) = self.try_build_root_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 index d04b91f..b9e3485 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -435,6 +435,231 @@ impl DataFusionPlanner { } } +impl DataFusionPlanner { + /// `RETURN m.name, count(x)` (or `count(*)`) over `(x)-[r]->(m)...`, where nothing else + /// reads `x` or `r`: on the CSR, count each `m`'s neighbours `x` once (`n(m)`, in place, + /// see [`HopDegreesNode`]), plan the rest of the pattern from `m`, and sum `n(m)` over + /// its rows. The hop's edges are read once to count them, rather than carried, one row + /// each, through the rest of the pattern and into the aggregate (eager aggregation, + /// as in [`Self::try_build_two_hop_count`]). `None` (plan as usual) for any other query. + /// + /// Applies when the hop's mean fan-out at `m` is at least two, every aggregate is a + /// plain count of rows, and the group keys are properties or variables other than + /// `x` and `r`. `WHERE` conjuncts over `x` alone filter the counted neighbours; any + /// other conjunct that reads `x` or `r` keeps the usual plan. + pub(crate) fn try_build_root_hop_count( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + projections: &[ProjectionItem], + ) -> Result> { + if self.adjacency.is_none() { + 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; + } + // The pattern's hops, last first, down to the one from the root. + let mut hops = vec![op]; + while let LogicalOperator::Expand { input, .. } = op { + if matches!(input.as_ref(), LogicalOperator::ScanByLabel { .. }) { + break; + } + op = input; + hops.push(op); + } + let Some((first, root_op)) = Hop::of(op) else { + return Ok(None); + }; + let LogicalOperator::ScanByLabel { + variable: root, + label: root_label, + properties: root_properties, + } = root_op + else { + return Ok(None); + }; + let (x, m) = (root.as_str(), first.target); + let counted = |v: &str| v == x || v == "*"; + let hidden: HashSet<&str> = [x].into_iter().chain(first.relationship).collect(); + // `x` and `r` appear in no later hop. + for hop in &hops[..hops.len() - 1] { + let LogicalOperator::Expand { + source_variable, + target_variable, + relationship_variable, + .. + } = hop + else { + return Ok(None); + }; + let names = [ + Some(source_variable), + Some(target_variable), + relationship_variable.as_ref(), + ]; + if names + .into_iter() + .flatten() + .any(|n| hidden.contains(n.as_str())) + { + return Ok(None); + } + } + let mut counts = false; + for p in projections { + match &p.expression { + ValueExpression::AggregateFunction { + name, + args, + distinct: false, + } if name.eq_ignore_ascii_case("count") => match args.as_slice() { + [ValueExpression::Variable(v)] if counted(v) => counts = true, + // `count(x)` of a node arrives as `count(x.id)`, never null. + [ValueExpression::Property(p)] + if p.variable == x + && self.id_column(ctx, x) == Some(qualify_column(x, &p.property)) => + { + counts = true + } + _ => return Ok(None), + }, + ValueExpression::Property(p) if !hidden.contains(p.variable.as_str()) => {} + ValueExpression::Variable(v) if !hidden.contains(v.as_str()) => {} + _ => return Ok(None), + } + } + if !counts { + return Ok(None); + } + let fan_out = self.fan_out.as_ref().and_then(|f| first.fan_out(f, m)); + if !fan_out.is_some_and(|d| d >= 2.0) { + return Ok(None); + } + let Some(middle_id) = self.id_column(ctx, m) else { + return Ok(None); + }; + let (mut far, mut rest) = (vec![], vec![]); + for predicate in predicates { + let mut names = HashSet::new(); + if referenced(predicate, &mut names).is_none() { + return Ok(None); + } + if names.iter().all(|n| n == x) { + far.push(predicate.clone()); + } else if names.iter().any(|n| hidden.contains(n.as_str())) { + return Ok(None); + } else { + rest.push(predicate.clone()); + } + } + + let LogicalOperator::Expand { + target_label, + relationship_types, + direction, + properties, + target_properties, + .. + } = first.op + else { + unreachable!("Hop::of matched an Expand") + }; + let scan = LogicalOperator::ScanByLabel { + variable: m.to_string(), + label: target_label.clone(), + properties: target_properties.clone(), + }; + let reversed = match direction { + RelationshipDirection::Outgoing => RelationshipDirection::Incoming, + _ => RelationshipDirection::Outgoing, + }; + let hop = LogicalOperator::Expand { + input: Box::new(scan.clone()), + source_variable: m.to_string(), + target_variable: x.to_string(), + target_label: root_label.clone(), + relationship_types: relationship_types.clone(), + direction: reversed, + relationship_variable: None, + properties: properties.clone(), + target_properties: root_properties.clone(), + }; + // The rest of the pattern, from `m` instead of the hop. + let mut rest_op = scan; + for hop in hops[..hops.len() - 1].iter().rev() { + let mut hop = (*hop).clone(); + if let LogicalOperator::Expand { input, .. } = &mut hop { + **input = rest_op; + } + rest_op = hop; + } + 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, + }; + // Built in the original order, so each hop takes the same relationship instance. + let hop_plan = self.build_operator(ctx, &filtered(hop, far))?; + let rest_plan = self.build_operator(ctx, &filtered(rest_op, rest))?; + + let column = |name: &str| col(Column::new_unqualified(name)); + let degrees = LogicalPlanBuilder::from(hop_plan) + .aggregate( + vec![column(&middle_id).alias("__m0")], + vec![count(lit(1)).alias("__n0")], + )? + .build() + .map_err(|e| self.plan_error("Failed to count a hop per node", e))?; + let degrees = LogicalPlan::Extension(Extension { + node: std::sync::Arc::new(HopDegreesNode::new(degrees, &middle_id, "__n0", "__s0")), + }); + let joined = LogicalPlanBuilder::from(rest_plan) + .join( + degrees, + JoinType::Inner, + (vec![middle_id.as_str()], vec!["__m0"]), + None, + )? + .build() + .map_err(|e| self.plan_error("Failed to join a hop's counts", e))?; + // Each count of rows becomes the sum of the counted neighbours, 0 over no rows. + let summed = ValueExpression::ScalarFunction { + name: "coalesce".to_string(), + args: vec![ + ValueExpression::AggregateFunction { + name: "sum".to_string(), + args: vec![ValueExpression::Variable("__n0".to_string())], + distinct: false, + }, + ValueExpression::Literal(crate::ast::PropertyValue::Integer(0)), + ], + }; + let projections: Vec = projections + .iter() + .map(|p| match &p.expression { + ValueExpression::AggregateFunction { .. } => ProjectionItem { + expression: summed.clone(), + alias: Some(p.alias.clone().unwrap_or_else(|| { + super::super::expression::to_cypher_column_name(&p.expression) + })), + }, + _ => p.clone(), + }) + .collect(); + self.build_project_with_aggregates(ctx, joined, &projections, &[]) + .map(Some) + } +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -601,6 +826,70 @@ mod tests { } } + /// Rows of `query` as sorted strings, equal on the join path and the CSR; on the CSR, + /// whether the plan counts a hop in place matches `in_place`. + async fn grouped(follows: &[(i64, i64)], query: &str, in_place: bool) -> Vec { + let mut answers = Vec::new(); + for use_csr in [false, true] { + let engine = engine(follows, &[], use_csr).await; + let plan = engine.explain(query, HashMap::new()).await.unwrap(); + assert_eq!( + plan.contains("degrees="), + use_csr && in_place, + "{query}: {plan}" + ); + let result = engine.execute(query, HashMap::new()).await.unwrap(); + let mut rows: Vec = (0..result.num_rows()) + .map(|i| { + let values = result + .columns() + .iter() + .map(|c| arrow::util::display::array_value_to_string(c, i).unwrap()); + values.collect::>().join(",") + }) + .collect(); + rows.sort(); + answers.push(rows); + } + assert_eq!(answers[0], answers[1], "{query}"); + answers.pop().unwrap() + } + + #[tokio::test] + async fn root_hop_counted_per_node() { + // 1, 2 and 3 follow 2; 2 follows 3 and itself; 99 has no Person row. + let follows = [(1, 2), (3, 2), (2, 2), (99, 2), (2, 3), (2, 99)]; + let one = "MATCH (a:Person)-[:FOLLOWS]->(b:Person) "; + let query = format!("{one}RETURN b.id AS b, count(a.id) AS n"); + assert!(grouped(&follows, &query, true) + .await + .contains(&"2,3".to_string())); + let query = format!("{one}WHERE a.age > 15 AND b.age < 50 RETURN b.age, count(*)"); + assert!(grouped(&follows, &query, true) + .await + .contains(&"30,2".to_string())); + // Nothing left to count is 0, not null. + let query = format!("{one}WHERE b.age > 1000 RETURN count(*) AS n"); + assert_eq!(grouped(&follows, &query, true).await, ["0"]); + // The rest of the pattern continues from b, with b's count on each of its rows. + let two = "MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:LIVES_IN]->(c:Person) "; + let query = format!("{two}RETURN c.id AS c, b.age AS age, count(a) AS n"); + assert!(grouped(&follows, &query, true) + .await + .contains(&"1,30,3".to_string())); + // Anything else that reads a, or a different aggregate, keeps the usual plan. + for query in [ + format!("{one}RETURN a.age, count(*)"), + format!("{one}RETURN b.id, count(a.age)"), + format!("{one}RETURN b.id, count(DISTINCT a)"), + format!("{one}RETURN b.id, sum(b.age)"), + format!("{one}WHERE a.age < b.age RETURN b.id, count(*)"), + format!("{two}WHERE c.id <> a.id RETURN c.id, count(*)"), + ] { + grouped(&follows, &query, false).await; + } + } + #[tokio::test] async fn other_two_hop_shapes() { // Different types never bind the same edge: the self-loops pair up. diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index bbd4ed0..e9b9ead 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -150,6 +150,46 @@ ProjectionExec: expr=[CASE WHEN __paths@0 IS NOT NULL THEN __paths@0 ELSE 0 END DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] ``` +### Counting the first hop per node + +`MATCH (x)-[:R]->(m)-[:S]->(c) RETURN c.name, count(x)` doesn't need `x`'s rows, only how many there are for each +`m`. When every aggregate is a plain count of rows and nothing else reads `x` or the relationship, the planner +counts each `m`'s `x` neighbours once, through the same `HopDegreesNode`, plans the rest of the pattern from `m`, +joins the counts on `m`'s id, and sums them in place of the count. Conditions on `x` alone filter the counted +neighbours. It applies only with the CSR and when the hop's mean fan-out at `m` is at least two, so that the counts +are fewer than the edges. + +Without it, the first hop's edges come out of the CSR expansion one row each, and those rows reach the aggregate spread +over every partition, so each partition's partial aggregate holds a row for nearly every group. On the benchmark +graph that made `RETURN person.name, count(follower.id), city.city, city.state, city.country` over two hops about +three times slower than the join path. Counting first, the rest of the pattern and the aggregate see one row per +person: + +```text +MATCH (follower:Person)-[:FOLLOWS]->(person:Person)-[:LIVES_IN]->(city:City) +RETURN person.name AS name, count(follower.id) AS numfollowers, city.city AS city, city.state AS state, + city.country AS country ORDER BY numfollowers DESC LIMIT 1 + +SortPreservingMergeExec: [numfollowers@1 DESC], fetch=1 + SortExec: TopK(fetch=1), expr=[numfollowers@1 DESC], preserve_partitioning=[true] + ProjectionExec: expr=[person__name@0 as name, CASE WHEN __agg_0@4 IS NOT NULL THEN __agg_0@4 ELSE 0 END as numfollowers, city__city@1 as city, city__state@2 as state, city__country@3 as country] + AggregateExec: mode=FinalPartitioned, gby=[person__name@0 as person__name, city__city@1 as city__city, city__state@2 as city__state, city__country@3 as city__country], aggr=[sum(__n0) as __agg_0] + RepartitionExec: partitioning=Hash([person__name@0, city__city@1, city__state@2, city__country@3], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[person__name@0 as person__name, city__city@1 as city__city, city__state@2 as city__state, city__country@3 as city__country], aggr=[sum(__n0) as __agg_0] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(person__id@0, __m0@0)], projection=[person__name@1, city__city@2, city__state@3, city__country@4, __n0@6] + CoalescePartitionsExec + ProjectionExec: expr=[person__id@0 as person__id, person__name@1 as person__name, city__city@4 as city__city, city__state@5 as city__state, city__country@6 as city__country] + CsrExpandExec: key=person__id, edges=100000, take=city[city__city, city__state, city__country] + ProjectionExec: expr=[id@0 as person__id, name@1 as person__name] + KeyStatisticsExec: distinct=[id=100000] + DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] + ProjectionExec: expr=[person__id@0 as __m0, __n0@1 as __n0] + CsrExpandExec: key=person__id, edges=2417738, take=follower[], degrees=[__n0, __s0] + ProjectionExec: expr=[id@0 as person__id] + KeyStatisticsExec: distinct=[id=100000] + DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] +``` + ## Join statistics for graph workloads Not every hop runs on the CSR, and the joins that remain need good cardinality estimates. DataFusion estimates an diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index edb3c59..72f3e18 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -448,7 +448,11 @@ def test_node_index_rebuilt_only_when_node_version_changes(tmp_path): @pytest.mark.parametrize( ("query", "taken"), [ - ("MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN count(*) AS n", "take=b[]"), + # Counted per node in place, from b: nothing is taken but each a's row position. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN count(*) AS n", + "take=a[], degrees=", + ), ("MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.id", "take=b[b__id*]"), ("MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.name", "take=b[b__name]"), ( diff --git a/python/python/tests/test_engine_with.py b/python/python/tests/test_engine_with.py index 231584d..2f6b786 100644 --- a/python/python/tests/test_engine_with.py +++ b/python/python/tests/test_engine_with.py @@ -7,6 +7,8 @@ 2->3, 3->3 and 5->1. Expected answers are written out by hand. """ +import re + import lance import pyarrow as pa import pytest @@ -97,6 +99,12 @@ def test_camel_case_aggregate_alias_can_be_ordered_by(): ] +def hop_takes_b_columns(plan): + """Whether a CSR hop in the physical plan fetches b's name or age for every edge.""" + physical = plan.split("physical_plan")[-1] + return re.search(r"CsrExpandExec.*take=b\[[^]]*b__(name|age)", physical) is not None + + def test_grouped_node_columns_fetched_after_grouping(): # With a node index (unique integer ids), `WITH b, ...` groups by b's id and looks # its other columns up afterwards; without one it groups by all of b's columns. @@ -109,8 +117,8 @@ def test_grouped_node_columns_fetched_after_grouping(): join = CypherEngine(CONFIG, datasets) plan = csr.explain(query) assert "NodeLookupExec: key=b__id" in plan - # The hop under the aggregate takes only b's id; the lookup fetches the rest. - assert "take=b[b__id*]" in plan, plan + # No hop under the aggregate takes b's other columns; the lookup fetches them. + assert not hop_takes_b_columns(plan), plan assert "NodeLookup" not in join.explain(query) by_id = lambda t: sorted(t.to_pylist(), key=lambda r: r["id"]) # noqa: E731 assert by_id(csr.execute(query)) == by_id(join.execute(query)) @@ -141,7 +149,10 @@ def test_grouped_node_columns_fetched_after_grouping(): "RETURN b.name AS name, count(*) AS n", False, ), - ("MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.id AS id, count(*) AS n", False), + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.id AS id, count(*) AS n", + False, + ), ], ) def test_grouped_node_properties_fetched_after_grouping(query, lookup): @@ -150,9 +161,11 @@ def test_grouped_node_properties_fetched_after_grouping(query, lookup): plan = csr.explain(query) assert ("NodeLookupExec: key=b__id" in plan) is lookup, plan if lookup: - assert "take=b[b__id*]" in plan, plan + assert not hop_takes_b_columns(plan), plan rows = lambda t: sorted(map(str, t.to_pylist())) # noqa: E731 - assert rows(csr.execute(query)) == rows(CypherEngine(CONFIG, datasets).execute(query)) + assert rows(csr.execute(query)) == rows( + CypherEngine(CONFIG, datasets).execute(query) + ) @pytest.mark.parametrize("use_csr", [False, True]) From 2a3ecc48c32bd279a86a58dc255c42503d0f65fc Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 20:46:52 -0500 Subject: [PATCH 6/7] docs(csr): say when the CSR helps and how to spot a whole-relationship hop A hop that starts from every node reads every edge either way, so the CSR is only ahead there when the planner counts in place. Describe the counting rewrites, what `degrees=` means in explain() output, and that `edges=` is the relationship's edge count, not what the hop reads. --- docs/csr.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/docs/csr.md b/docs/csr.md index 948d4f5..e044ad8 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -470,11 +470,53 @@ print(open_engine().delete_stored_csr()) # 1 change. `python/benchmarks/csr_persist.py` measures the same cases on a synthetic graph, including a load with a delta. +## When it helps + +A CSR hop's cost follows the edges it reads, and it reads only the neighbours of its input rows. A hash join +reads the whole relationship table and builds a hash table on it, however few rows it starts from. So the CSR +gains most when a hop starts from few rows: a filtered node, or the result of an earlier hop. + +A hop that starts from every node of a label reads every edge either way. An in-memory hash join is already +cheap there, and a CSR hop that emits one row per edge is no faster. What the CSR adds is that it can count a +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. +- `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 + [grouping by a node](#grouping-by-a-node)). + +In `explain()` output, a hop counted in place shows `degrees=[...]` on its `CsrExpandExec` line, and its output +has one row per node rather than one per edge. A `WHERE` condition on the counted neighbour alone still applies +(`filter=` on the same line); a condition that compares it with another variable, or reads the relationship, +needs the hop's rows and keeps the plain hop. + +### Spotting a hop over the whole relationship + +`edges=` on a `CsrExpandExec` line is the relationship's edge count, not the number of edges the hop reads. +The hop reads all of them when its input is the whole node table: the lines below it are a scan of that table, +with no `FilterExec` and no earlier hop in between. From a 100,000-person graph with 2.4 million `FOLLOWS` edges, +`MATCH (follower:Person)-[:FOLLOWS]->(person:Person) RETURN person.name, max(follower.age)` plans: + +```text +CsrExpandExec: key=follower__id, edges=2417738, take=person[person__name] + ProjectionExec: expr=[id@0 as follower__id, age@1 as follower__age] + KeyStatisticsExec: distinct=[id=100000] + DataSourceExec: partitions=10, partition_sizes=[2, 2, 2, 1, 1, 1, 1, 1, 1, 1] +``` + +If such a hop isn't counted in place and the query is slower with `use_csr=True`, the join path is doing the +same work with less overhead. Narrow the hop's input (a filter on its start node) where the query allows it, +or run that workload without the CSR. + ## Should I turn it on? | Situation | Recommendation | | --- | --- | | Queries traverse relationships, especially 2+ hops or variable-length | **Turn it on** | +| Queries read whole relationships (no filter on the start node) and only count | **Turn it on**; counts are taken per node ([when it helps](#when-it-helps)) | +| Queries read whole relationships and return or group by each edge's properties | About even with the join path; measure both | | Integer node keys and edge endpoints | **Turn it on**; you get the fused target lookup too | | String or UUID keys | Little benefit; add an integer id column if traversal matters | | Only node lookups and aggregations, no hops | Leave it off; it costs build time and memory for nothing | From 1b810899186f51c319e8c7b0af83094b5b6a2a0c Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 2 Nov 2026 21:05:41 -0500 Subject: [PATCH 7/7] docs(csr): say the join fallback is chosen by capability, not cost With use_csr=True, a hop runs on the CSR whenever the CSR can take it; the fallback to a join covers hops it can't express, not slower ones, and use_csr applies to the whole engine. Say so next to the fallback statements and in the advice for whole-relationship hops. --- docs/csr.md | 11 +++++++++-- docs/index.md | 5 +++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/csr.md b/docs/csr.md index e044ad8..1efa3dc 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -143,6 +143,11 @@ When those hold, it covers: Everything else in the query (filters, aggregations, sorts and any hop the CSR can't take) is planned by DataFusion as usual. You can mix: a three-hop query might run two hops on the CSR and one as a join. +The planner chooses between the two by what the CSR can express, not by which is faster. A hop that meets the +conditions above always runs on the CSR, even when a join would be about as fast; see +[when it helps](#when-it-helps). `use_csr` is set once, for the whole engine, because it decides whether the +CSR is built at all: with `use_csr=False` every hop is a join. + Named path values and `relationships(path)` use a query-local inventory of actual edge occurrences because CSR neighbors do not carry edge identity or properties. This also preserves parallel edges when a path passes through `collect()` and `UNWIND`. The inventory materializes the participating node and relationship @@ -507,8 +512,10 @@ CsrExpandExec: key=follower__id, edges=2417738, take=person[person__name] ``` If such a hop isn't counted in place and the query is slower with `use_csr=True`, the join path is doing the -same work with less overhead. Narrow the hop's input (a filter on its start node) where the query allows it, -or run that workload without the CSR. +same work with less overhead. The planner won't switch that hop to a join for you: the fallback covers hops the +CSR can't take, not slower ones. Narrow the hop's input (a filter on its start node) where the query allows +it. If the workload is mostly such queries, measure it with `use_csr=False` too, as that setting applies to +every query on the engine. ## Should I turn it on? diff --git a/docs/index.md b/docs/index.md index ae01b96..973ed6a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -118,8 +118,9 @@ flowchart LR J --> R ``` -Every hop that the CSR can answer becomes a neighbour lookup; anything it can't express falls back to an -ordinary join, so enabling the CSR never changes results. [How it works](execution.md) walks through each +With `use_csr=True`, every hop that the CSR can answer becomes a neighbour lookup; anything it can't express +falls back to an ordinary join, so enabling the CSR never changes results. The choice is by what the CSR can +express, not by cost: see [when the CSR helps](csr.md#when-it-helps). [How it works](execution.md) walks through each stage. ## What it is not