diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 6fac611..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. /// @@ -2635,30 +2906,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/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/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 03f90fe..b9e3485 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,10 @@ 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 { if self.is_distinct_edges(ctx, predicate, pairs) { if !chain { @@ -297,14 +304,25 @@ 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()); } 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() @@ -316,54 +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: relationship_variable.clone(), - 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) }; // 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 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)?; @@ -393,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; @@ -514,6 +781,115 @@ 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 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); + } + } + + /// 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/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/csr.md b/docs/csr.md index f37a788..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 @@ -219,9 +224,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 +239,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. @@ -469,11 +475,55 @@ 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. 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? | 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 | 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 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 3cc285d..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: @@ -128,11 +129,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 +146,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/docs/internals/optimizer.md b/docs/internals/optimizer.md index 5fe00af..e9b9ead 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`. @@ -104,6 +106,90 @@ 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] +``` + +### 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 cb93ab7..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. @@ -107,12 +115,59 @@ 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 + # 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)) +@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 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) + ) + + @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)