diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index cc9c2c6..bddc22f 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -61,8 +61,10 @@ use futures::StreamExt; use lance::dataset::{Dataset, ProjectionRequest}; mod node_lookup; +mod path_exists; use lance_table::format::Fragment; pub use node_lookup::NodeLookupNode; +pub use path_exists::{Hop, PathExistsNode}; /// Yield after bounded CPU work without requiring a particular async runtime. pub(crate) async fn cooperative_yield() { @@ -2408,6 +2410,9 @@ impl ExtensionPlanner for CsrExtensionPlanner { if let Some(node) = node.as_any().downcast_ref::() { return Ok(Some(node_lookup::plan_node_lookup(node, physical_inputs))); } + if let Some(node) = node.as_any().downcast_ref::() { + return Ok(Some(path_exists::plan_path_exists(node, physical_inputs))); + } let Some(node) = node.as_any().downcast_ref::() else { return Ok(None); }; @@ -2610,6 +2615,9 @@ fn used_columns(plan: &LogicalPlan) -> datafusion::common::Result() { used.insert(lookup.key_column().to_string()); true + } else if let Some(exists) = ext.node.as_any().downcast_ref::() { + used.extend(exists.endpoints().map(str::to_string)); + true } else { false } diff --git a/crates/lance-csr/src/csr_expand/path_exists.rs b/crates/lance-csr/src/csr_expand/path_exists.rs new file mode 100644 index 0000000..93ce7e4 --- /dev/null +++ b/crates/lance-csr/src/csr_expand/path_exists.rs @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! A pattern predicate answered per row by walking the CSR from the row's start node. +//! +//! The relational marking of `WHERE (a)-[:R]->()-[:S]->(b)` plans everything before the +//! predicate a second time as a seed, matches the pattern from it and joins the marks +//! back, so each predicate doubles the plan above it. This node instead walks the +//! adjacencies from each input row's `start` node, stops at the first path found, and +//! appends the same mark: `true`, or null when there is no path. + +use super::*; + +/// One hop of the walk: its adjacency, and the node table every node it reaches must be +/// in, as the relational plan joins each node of the pattern to its table. +#[derive(Debug, Clone)] +pub struct Hop { + pub neighbors: Neighbors, + pub nodes: Arc, +} + +#[derive(Clone)] +pub struct PathExistsNode { + input: LogicalPlan, + start: String, + /// The bound node the last hop must reach; any node when `None`. + end: Option, + hops: Vec, + mark: String, + schema: DFSchemaRef, +} + +impl PathExistsNode { + pub fn try_new( + input: LogicalPlan, + start: String, + end: Option, + hops: Vec, + mark: String, + ) -> Result { + let added = DFSchema::try_from(Schema::new(vec![Field::new( + &mark, + DataType::Boolean, + true, + )])) + .map_err(|e| plan_error(format!("PathExists schema: {}", e)))?; + let schema = input + .schema() + .join(&added) + .map_err(|e| plan_error(format!("PathExists schema: {}", e)))?; + Ok(Self { + input, + start, + end, + hops, + mark, + schema: Arc::new(schema), + }) + } + + /// The start and end columns, which the node reads without an expression. + pub(super) fn endpoints(&self) -> impl Iterator { + std::iter::once(self.start.as_str()).chain(self.end.as_deref()) + } + + fn describe(&self) -> String { + let end = self.end.as_ref().map(|e| format!(", match={e}")); + format!( + "start={}, hops={}{}, mark={}", + self.start, + self.hops.len(), + end.unwrap_or_default(), + self.mark + ) + } +} + +/// Whether a path over `hops` leads from `node` to `end` (or anywhere). +// shortcut: no memo of visited nodes, so a long chain of high-degree hops with a bound end +// costs the product of their degrees per row; add one if longer predicates appear. +fn reaches(hops: &[Hop], node: i64, end: Option) -> bool { + let Some((hop, rest)) = hops.split_first() else { + return end.is_none_or(|end| end == node); + }; + let mut position = NeighborPosition::default(); + let found = hop + .neighbors + .neighbors_from(node, &mut position) + .any(|next| hop.nodes.position(next as i64).is_some() && reaches(rest, next as i64, end)); + found +} + +// Logical nodes must be comparable; adjacencies compare by identity. +impl PartialEq for PathExistsNode { + fn eq(&self, other: &Self) -> bool { + self.input == other.input + && self.start == other.start + && self.end == other.end + && self.mark == other.mark + && self.hops.len() == other.hops.len() + && self + .hops + .iter() + .zip(&other.hops) + .all(|(a, b)| a.neighbors.ptr_eq(&b.neighbors) && Arc::ptr_eq(&a.nodes, &b.nodes)) + } +} +impl Eq for PathExistsNode {} +impl Hash for PathExistsNode { + fn hash(&self, state: &mut H) { + self.input.hash(state); + self.start.hash(state); + self.end.hash(state); + self.mark.hash(state); + } +} +impl PartialOrd for PathExistsNode { + fn partial_cmp(&self, other: &Self) -> Option { + (&self.start, &self.mark).partial_cmp(&(&other.start, &other.mark)) + } +} + +// DataFusion formats every extension node while planning it; printing the input plan +// here would make that quadratic in the plan's depth. +impl std::fmt::Debug for PathExistsNode { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + UserDefinedLogicalNodeCore::fmt_for_explain(self, f) + } +} + +impl UserDefinedLogicalNodeCore for PathExistsNode { + fn name(&self) -> &str { + "PathExists" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + /// Filters on the mark must stay above the node. + fn prevent_predicate_push_down_columns(&self) -> HashSet { + HashSet::from([self.mark.clone()]) + } + + /// The input only needs the requested input columns plus the endpoints. + fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option>> { + let schema = self.input.schema(); + let input_width = schema.fields().len(); + let mut needed: Vec = output_columns + .iter() + .copied() + .filter(|&i| i < input_width) + .collect(); + for column in std::iter::once(&self.start).chain(&self.end) { + needed.push(schema.index_of_column_by_name(None, column)?); + } + needed.sort_unstable(); + needed.dedup(); + Some(vec![needed]) + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PathExists: {}", self.describe()) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + inputs: Vec, + ) -> datafusion::common::Result { + let input = inputs.into_iter().next().expect("PathExists has one input"); + Self::try_new( + input, + self.start.clone(), + self.end.clone(), + self.hops.clone(), + self.mark.clone(), + ) + .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string())) + } +} + +/// Executes [`PathExistsNode`] one input batch at a time. +#[derive(Debug)] +pub struct PathExistsExec { + input: Arc, + start: String, + end: Option, + hops: Arc>, + schema: SchemaRef, + properties: Arc, +} + +impl PathExistsExec { + fn new(node: &PathExistsNode, input: Arc) -> Self { + let schema: SchemaRef = Arc::new(node.schema.as_arrow().clone()); + let properties = Self::properties_for(&schema, &input); + Self { + input, + start: node.start.clone(), + end: node.end.clone(), + hops: Arc::new(node.hops.clone()), + schema, + properties, + } + } + + fn properties_for(schema: &SchemaRef, input: &Arc) -> Arc { + Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + + /// The batch with each row's mark appended. A null endpoint matches no path. + fn mark( + batch: RecordBatch, + start: &str, + end: Option<&str>, + hops: &[Hop], + schema: &SchemaRef, + ) -> datafusion::common::Result { + let external = |e: GraphError| datafusion::common::DataFusionError::External(Box::new(e)); + let starts = int64_column(&batch, start).map_err(external)?; + let ends = end + .map(|end| int64_column(&batch, end)) + .transpose() + .map_err(external)?; + let marks: arrow::array::BooleanArray = (0..batch.num_rows()) + .map(|i| { + let end = match &ends { + Some(ends) if ends.is_null(i) => return None, + Some(ends) => Some(ends.value(i)), + None => None, + }; + (starts.is_valid(i) && reaches(hops, starts.value(i), end)).then_some(true) + }) + .collect(); + let mut columns = batch.columns().to_vec(); + columns.push(Arc::new(marks)); + Ok(RecordBatch::try_new(schema.clone(), columns)?) + } +} + +impl DisplayAs for PathExistsExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + let end = self.end.as_ref().map(|e| format!(", match={e}")); + write!( + f, + "PathExistsExec: start={}, hops={}{}", + self.start, + self.hops.len(), + end.unwrap_or_default() + ) + } +} + +impl ExecutionPlan for PathExistsExec { + fn name(&self) -> &str { + "PathExistsExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::common::Result> { + let input = children + .into_iter() + .next() + .expect("PathExistsExec has one child"); + Ok(Arc::new(Self { + properties: Self::properties_for(&self.schema, &input), + input, + start: self.start.clone(), + end: self.end.clone(), + hops: self.hops.clone(), + schema: self.schema.clone(), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion::common::Result { + let input = self.input.execute(partition, context)?; + let (start, end, hops, schema) = ( + self.start.clone(), + self.end.clone(), + self.hops.clone(), + self.schema.clone(), + ); + // shortcut: one batch is marked without yielding; bound the work per poll if + // predicates over high-degree hops show up in cancellation latency. + let stream = + input.map(move |batch| Self::mark(batch?, &start, end.as_deref(), &hops, &schema)); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } +} + +pub(super) fn plan_path_exists( + node: &PathExistsNode, + physical_inputs: &[Arc], +) -> Arc { + Arc::new(PathExistsExec::new(node, physical_inputs[0].clone())) +} diff --git a/crates/lance-csr/src/datafusion_planner/analysis.rs b/crates/lance-csr/src/datafusion_planner/analysis.rs index 76ab3fe..31fa380 100644 --- a/crates/lance-csr/src/datafusion_planner/analysis.rs +++ b/crates/lance-csr/src/datafusion_planner/analysis.rs @@ -55,6 +55,36 @@ impl<'a> PlanningContext<'a> { } } + /// Skip the relationship instances that `planned` would take beyond `built`, as when + /// `built` is planned in place of `planned`, which contains it. + pub(crate) fn skip_relationship_instances( + &mut self, + planned: &crate::logical_plan::LogicalOperator, + built: &crate::logical_plan::LogicalOperator, + ) -> Result<()> { + let count = |op| -> Result> { + let mut counts = HashMap::new(); + for instance in analyze(op)?.relationship_instances { + *counts.entry(instance.rel_type).or_default() += 1; + } + Ok(counts) + }; + let built = count(built)?; + for (rel_type, n) in count(planned)? { + let skip = n - built.get(&rel_type).copied().unwrap_or(0); + if skip > 0 { + let idx = self.relationship_instance_idx.entry(rel_type); + match idx { + std::collections::hash_map::Entry::Occupied(mut e) => *e.get_mut() += skip, + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(skip - 1); + } + } + } + } + Ok(()) + } + /// Get the next relationship instance for a given type (returns a clone) pub fn next_relationship_instance(&mut self, rel_type: &str) -> Result { let idx = self @@ -301,6 +331,11 @@ fn analyze_operator( LogicalOperator::CallInput { labels } => { analysis.var_to_label.extend(labels.clone()); } + // The fallback holds the input; the builder keeps its relationship instances + // in step whichever of the two it plans. + LogicalOperator::PathExists { fallback, .. } => { + analyze_operator(fallback, analysis, rel_counter)?; + } LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } Ok(()) diff --git a/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs index 51c99bb..bcb3a92 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs @@ -6,14 +6,15 @@ use std::collections::HashSet; use arrow_schema::DataType; -use datafusion::common::Column; +use datafusion::common::{Column, ScalarValue}; use datafusion::functions::core::expr_fn::coalesce; use datafusion::functions_aggregate::array_agg::array_agg; use datafusion::functions_aggregate::count::count; -use datafusion::functions_nested::expr_fn::{array_element, array_length, gen_series, make_array}; +use datafusion::functions_nested::expr_fn::{array_length, array_slice, gen_series, make_array}; use datafusion::functions_window::expr_fn::row_number; use datafusion::logical_expr::{ - col, lit, Cast, Expr, ExprFunctionExt, ExprSchemable, JoinType, LogicalPlan, LogicalPlanBuilder, + col, lit, when, Cast, Expr, ExprFunctionExt, ExprSchemable, JoinType, LogicalPlan, + LogicalPlanBuilder, }; use super::super::analysis::PlanningContext; @@ -68,16 +69,37 @@ impl DataFusionPlanner { let length = int64(coalesce(vec![array_length(list.clone()), lit(0_u64)])); let last = datafusion::functions::core::expr_fn::greatest(vec![length.clone(), lit(1_i64)]); let positions = gen_series(lit(1_i64), last, lit(1_i64)); + // The list is unnested beside its positions, so an element row carries its item, + // not a copy of the whole list. A list read from an input column is rebuilt once + // per input row below instead of being grouped on, so unless the comprehension + // reads it, projection pushdown drops it before the unnest. + let (items, len, is_null) = ( + format!("{output}__items"), + format!("{output}__len"), + format!("{output}__null"), + ); + let iterated = match &list { + Expr::Column(c) if numbered.schema().has_column(c) => Some(c.name.clone()), + _ => None, + }; + let select = |e: Expr| datafusion::logical_expr::select_expr::SelectExpr::Expression(e); let exploded = LogicalPlanBuilder::from(numbered.clone()) .project(vec![ datafusion::logical_expr::wildcard(), - datafusion::logical_expr::select_expr::SelectExpr::Expression( - positions.alias(&pos), - ), + select(positions.alias(&pos)), + select(list.clone().alias(&items)), + select(length.alias(&len)), + select(list.is_null().alias(&is_null)), ]) - .and_then(|b| b.unnest_column(pos.as_str())) + .and_then(|b| { + b.unnest_columns_with_options( + vec![Column::from_name(&pos), Column::from_name(&items)], + Default::default(), + ) + }) .map_err(error("Failed to unnest comprehension list"))?; - let item = array_element(list, col(&pos)); + let length = col(&len); + let item = col(&items); let element_rows = match label { None => exploded .project(vec![ @@ -98,15 +120,54 @@ impl DataFusionPlanner { } })?; let nodes = self.build_scan(ctx, variable, label, &Default::default())?; - let joined = exploded + let with_ids = exploded .project(vec![ datafusion::logical_expr::wildcard(), datafusion::logical_expr::select_expr::SelectExpr::Expression( item.alias(&id_column), ), ]) - .and_then(|b| { - b.join( + .map_err(error("Failed to join comprehension elements to nodes"))?; + // An in-memory node index looks each element up, at a cost that follows the + // elements rather than the node table, as for `late_node`; Lance tables join. + let rows = self + .adjacency + .as_ref() + .and_then(|a| a.nodes(label)) + .filter(|rows| rows.in_memory()); + let prefix = format!("{}__", variable.to_lowercase()); + let wanted: Option> = rows.and_then(|rows| { + nodes + .schema() + .fields() + .iter() + .map(|f| { + let suffix = f.name().strip_prefix(&prefix)?; + let column = rows + .schema() + .fields() + .iter() + .find(|c| c.name().eq_ignore_ascii_case(suffix))?; + Some((f.name().clone(), column.name().clone())) + }) + .collect() + }); + if let (Some(rows), Some(wanted)) = (rows, wanted) { + let lookup = crate::csr_expand::NodeLookupNode::try_new( + with_ids + .build() + .map_err(error("Failed to look up comprehension elements"))?, + id_column, + variable.to_lowercase(), + rows.clone(), + &wanted, + )?; + LogicalPlan::Extension(datafusion::logical_expr::Extension { + node: std::sync::Arc::new(lookup), + }) + } else { + let joined = with_ids + .join( nodes, JoinType::Left, ( @@ -115,20 +176,20 @@ impl DataFusionPlanner { ), None, ) - }) - .and_then(|b| b.build()) - .map_err(error("Failed to join comprehension elements to nodes"))?; - let kept: Vec = joined - .schema() - .columns() - .into_iter() - .filter(|c| c.name != id_column) - .map(Expr::Column) - .collect(); - LogicalPlanBuilder::from(joined) - .project(kept) - .and_then(|b| b.build()) - .map_err(error("Failed to project comprehension elements"))? + .and_then(|b| b.build()) + .map_err(error("Failed to join comprehension elements to nodes"))?; + let kept: Vec = joined + .schema() + .columns() + .into_iter() + .filter(|c| c.name != id_column) + .map(Expr::Column) + .collect(); + LogicalPlanBuilder::from(joined) + .project(kept) + .and_then(|b| b.build()) + .map_err(error("Failed to project comprehension elements"))? + } } }; @@ -155,12 +216,26 @@ impl DataFusionPlanner { let collected = aggregate .map_err(error("Failed to aggregate comprehension values"))? .alias(output); + let mut aggregates = vec![collected]; + let rebuilt = format!("{output}__list"); + if iterated.is_some() { + // The padding row of an empty list sorts last and is sliced off. + aggregates.push( + array_agg(col(&items)) + .order_by(vec![col(&pos).sort(true, false)]) + .build() + .map_err(error("Failed to rebuild comprehension list"))? + .alias(&rebuilt), + ); + } // One group per input row: its number plus every input column, all equal within it. let names: HashSet = numbered .schema() .fields() .iter() .map(|f| f.name().clone()) + .chain([len.clone(), is_null.clone()]) + .filter(|name| Some(name) != iterated.as_ref()) .collect(); let group: Vec = elements .schema() @@ -174,8 +249,19 @@ impl DataFusionPlanner { .fields() .iter() .filter(|f| *f.name() != row) - .map(|f| Expr::Column(Column::new_unqualified(f.name()))) - .collect(); + .map(|f| { + if Some(f.name()) != iterated.as_ref() { + return Ok(Expr::Column(Column::new_unqualified(f.name()))); + } + let list = array_slice(col(&rebuilt), lit(1_i64), col(&len), None); + when(col(&is_null), lit(ScalarValue::Null)) + .otherwise(list) + .map(|e| { + datafusion::logical_expr::cast(e, f.data_type().clone()).alias(f.name()) + }) + .map_err(error("Failed to rebuild comprehension list")) + }) + .collect::>()?; let value = if count_only { col(output) } else { @@ -183,7 +269,7 @@ impl DataFusionPlanner { }; out.push(value.alias(output)); LogicalPlanBuilder::from(elements) - .aggregate(group, vec![collected]) + .aggregate(group, aggregates) .and_then(|b| b.project(out)) .and_then(|b| b.build()) .map_err(error("Failed to build list comprehension")) @@ -192,7 +278,7 @@ impl DataFusionPlanner { #[cfg(test)] mod tests { - use arrow_array::{types::Int64Type, ArrayRef, Int64Array, ListArray, RecordBatch}; + use arrow_array::{types::Int64Type, Array, ArrayRef, Int64Array, ListArray, RecordBatch}; use std::collections::HashMap; use std::sync::Arc; @@ -264,8 +350,10 @@ mod tests { .collect::>() ); let plan = query.explain(data()).await.unwrap().to_lowercase(); - assert!(plan.contains("count("), "{plan}"); - assert!(!plan.contains("array_agg("), "{plan}"); + // The unused rebuild of `p.values` is pruned before execution. + let physical = plan.split("physical_plan").nth(1).unwrap(); + assert!(physical.contains("count("), "{plan}"); + assert!(!physical.contains("array_agg("), "{plan}"); let mapped = make_query("| x") .explain(data()) .await @@ -274,4 +362,51 @@ mod tests { assert!(mapped.contains("array_agg("), "{mapped}"); } } + + #[tokio::test] + async fn iterated_list_is_returned_unchanged() { + let config = crate::GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let run = |query: &str| { + let query = crate::CypherQuery::new(query) + .unwrap() + .with_config(config.clone()); + async move { query.execute(data(), None).await.unwrap() } + }; + let expected = run("MATCH (p:Person) RETURN p.id AS id, p.values AS v ORDER BY id").await; + for comprehension in [ + "size([x IN p.values WHERE x > 1])", + "[x IN p.values WHERE x > 1 | x * 10]", + // The body reads the iterated list itself. + "[x IN p.values | size(p.values)]", + ] { + let actual = run(&format!( + "MATCH (p:Person) RETURN p.id AS id, p.values AS v, {comprehension} AS c ORDER BY id" + )) + .await; + assert_eq!(actual.column(1), expected.column(1), "{comprehension}"); + } + let reads_list = run( + "MATCH (p:Person) RETURN p.id AS id, [x IN p.values | size(p.values)] AS c ORDER BY id", + ) + .await; + let c = reads_list + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let sizes: Vec<_> = (0..c.len()) + .map(|i| { + c.value(i) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(sizes, vec![vec![4; 4], vec![], vec![], vec![1]]); + } } diff --git a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs index 3387360..318d506 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs @@ -423,6 +423,75 @@ impl DataFusionPlanner { }))) } + /// A pattern predicate walked on the CSR from each row of `input` (see + /// [`LogicalOperator::PathExists`]), or its relational `fallback` when a hop has no + /// CSR or an endpoint is not an integer id column of `input`. + #[allow(clippy::too_many_arguments)] + pub(crate) fn build_path_exists( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + start: &str, + hops: &[(String, RelationshipDirection, String)], + end: Option<&str>, + mark: &str, + fallback: &LogicalOperator, + ) -> Result { + use crate::csr_expand::{is_supported_id_type, Hop, Neighbors, PathExistsNode}; + + let neighbors: Option> = self.adjacency.as_ref().and_then(|catalog| { + hops.iter() + .map(|(rel_type, direction, label)| { + let neighbors = match direction { + RelationshipDirection::Outgoing => catalog + .get(rel_type, true) + .cloned() + .map(Neighbors::directed), + RelationshipDirection::Incoming => catalog + .get(rel_type, false) + .cloned() + .map(Neighbors::directed), + RelationshipDirection::Undirected => catalog + .get(rel_type, true) + .zip(catalog.get(rel_type, false)) + .map(|(out, inc)| Neighbors::undirected(out.clone(), inc.clone())), + }; + let nodes = catalog.nodes(label).cloned(); + neighbors + .zip(nodes) + .map(|(neighbors, nodes)| Hop { neighbors, nodes }) + }) + .collect() + }); + let start_id = self.id_column(ctx, start); + let end_id = end.map(|end| self.id_column(ctx, end)); + let (Some(neighbors), Some(start_id), None | Some(Some(_))) = + (neighbors, start_id, &end_id) + else { + return self.build_operator(ctx, fallback); + }; + let end_id = end_id.flatten(); + let instances = ctx.relationship_instance_idx.clone(); + let plan = self.build_operator(ctx, input)?; + let is_id = |column: &String| { + plan.schema() + .field_with_unqualified_name(column) + .is_ok_and(|f| is_supported_id_type(f.data_type())) + }; + if !is_id(&start_id) || !end_id.as_ref().is_none_or(is_id) { + // The fallback plans the input again, taking the same instances. + ctx.relationship_instance_idx = instances; + return self.build_operator(ctx, fallback); + } + ctx.skip_relationship_instances(fallback, input)?; + let node = PathExistsNode::try_new(plan, start_id, end_id, neighbors, mark.to_string())?; + Ok(LogicalPlan::Extension( + datafusion::logical_expr::Extension { + node: Arc::new(node), + }, + )) + } + /// Mark a freshly built CSR expansion as a closing edge on `bound_target`. fn with_closing_column(plan: LogicalPlan, bound_target: String) -> Result { use crate::csr_expand::CsrExpandNode; diff --git a/crates/lance-csr/src/datafusion_planner/builder/helpers.rs b/crates/lance-csr/src/datafusion_planner/builder/helpers.rs index f2e5721..b418918 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/helpers.rs @@ -119,9 +119,8 @@ impl DataFusionPlanner { Self::collect_variables(first, vars); } } - LogicalOperator::ListComprehension { input, .. } => { - Self::collect_variables(input, vars) - } + LogicalOperator::ListComprehension { input, .. } + | LogicalOperator::PathExists { input, .. } => Self::collect_variables(input, vars), LogicalOperator::Apply { input, labels, .. } => { Self::collect_variables(input, vars); vars.extend(labels.keys().cloned()); @@ -204,7 +203,8 @@ impl DataFusionPlanner { } LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} - LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), + LogicalOperator::ListComprehension { input, .. } + | LogicalOperator::PathExists { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Apply { input, labels: exports, diff --git a/crates/lance-csr/src/datafusion_planner/builder/mod.rs b/crates/lance-csr/src/datafusion_planner/builder/mod.rs index b559d84..e07b4d1 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/mod.rs @@ -146,6 +146,14 @@ impl DataFusionPlanner { output, *count_only, ), + LogicalOperator::PathExists { + input, + start, + hops, + end, + mark, + fallback, + } => self.build_path_exists(ctx, input, start, hops, end.as_deref(), mark, fallback), LogicalOperator::ComprehensionElement => Ok(ctx .comprehension_elements .last() diff --git a/crates/lance-csr/src/logical_plan.rs b/crates/lance-csr/src/logical_plan.rs index 5087a01..21c1f24 100644 --- a/crates/lance-csr/src/logical_plan.rs +++ b/crates/lance-csr/src/logical_plan.rs @@ -204,6 +204,20 @@ pub enum LogicalOperator { /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. ComprehensionElement, + + /// A pattern predicate that walks anonymous single-type hops from the bound node + /// `start`, ending anywhere or at the bound node `end`. Marks each row of `input` with + /// `mark` (`true`, or null when no path exists) by walking the CSR from that row when + /// every hop has one; otherwise `fallback`, the relational marking, is planned instead. + PathExists { + input: Box, + start: String, + /// Relationship type, direction and target label of each hop, in order from `start`. + hops: Vec<(String, RelationshipDirection, String)>, + end: Option, + mark: String, + fallback: Box, + }, } /// Projection item for SELECT/RETURN clauses @@ -1162,6 +1176,8 @@ impl<'a> LogicalPlanner<'a> { } => { self.exists_marks += 1; let mark = format!("__exists_{}", self.exists_marks); + let walk = self.exists_walk(patterns, where_clause.as_deref()); + let original = walk.is_some().then(|| input.clone()); let inner = MatchClause { patterns: patterns.clone(), optional: true, @@ -1203,17 +1219,106 @@ impl<'a> LogicalPlanner<'a> { projections, }), }; - *expr = BooleanExpression::IsNotNull(ValueExpression::Variable(mark)); - Ok(LogicalOperator::OptionalMatch { + *expr = BooleanExpression::IsNotNull(ValueExpression::Variable(mark.clone())); + let fallback = LogicalOperator::OptionalMatch { input, optional: Box::new(marked), keys, + }; + Ok(match (walk, original) { + (Some((start, hops, end)), Some(original)) => LogicalOperator::PathExists { + input: Box::new(original), + start, + hops, + end, + mark, + fallback: Box::new(fallback), + }, + _ => fallback, }) } _ => Ok(input), } } + /// `(start, hops, end)` when a pattern predicate is one path from a bound node through + /// anonymous nodes, over anonymous relationships of one type each, ending at an + /// anonymous or bound node, with endpoint labels that chain. Hops have distinct types, + /// so no two can bind the same relationship. + #[allow(clippy::type_complexity)] + fn exists_walk( + &self, + patterns: &[GraphPattern], + where_clause: Option<&BooleanExpression>, + ) -> Option<( + String, + Vec<(String, RelationshipDirection, String)>, + Option, + )> { + let ([GraphPattern::Path(path)], None) = (patterns, where_clause) else { + return None; + }; + let bound_label = |variable: &str| { + self.variables + .get(variable) + .filter(|label| !label.contains('|')) + }; + let fits = |node: &NodePattern, label: &str| { + node.properties.is_empty() + && node.labels.len() <= 1 + && node.labels.iter().all(|l| l.eq_ignore_ascii_case(label)) + }; + let start = path.start_node.variable.clone()?; + let mut label = bound_label(&start)?.clone(); + if path.variable.is_some() || !fits(&path.start_node, &label) { + return None; + } + let mut hops: Vec<(String, RelationshipDirection, String)> = Vec::new(); + let mut end = None; + for (i, segment) in path.segments.iter().enumerate() { + let rel = &segment.relationship; + let [rel_type] = rel.types.as_slice() else { + return None; + }; + if rel.variable.is_some() + || !rel.properties.is_empty() + || rel.length.is_some() + || hops + .iter() + .any(|(t, _, _)| t.eq_ignore_ascii_case(rel_type)) + { + return None; + } + let mapping = self.config.get_relationship_mapping(rel_type)?; + let (source, target) = ( + mapping.source_label.as_deref()?, + mapping.target_label.as_deref()?, + ); + let (from, to) = match rel.direction { + RelationshipDirection::Outgoing => (source, target), + RelationshipDirection::Incoming => (target, source), + RelationshipDirection::Undirected if source.eq_ignore_ascii_case(target) => { + (source, target) + } + RelationshipDirection::Undirected => return None, + }; + if !from.eq_ignore_ascii_case(&label) || !fits(&segment.end_node, to) { + return None; + } + label = to.to_string(); + hops.push((rel_type.clone(), rel.direction.clone(), to.to_string())); + if let Some(variable) = &segment.end_node.variable { + // Only the last node may be bound, and an unbound one would be new. + let last = i + 1 == path.segments.len(); + if !last || !bound_label(variable)?.eq_ignore_ascii_case(to) { + return None; + } + end = Some(variable.clone()); + } + } + (!hops.is_empty()).then_some((start, hops, end)) + } + /// [`Self::rewrite_exists`] for the conditions of a CASE inside a value. fn rewrite_exists_in_value( &mut self, @@ -1958,9 +2063,8 @@ impl<'a> LogicalPlanner<'a> { LogicalOperator::OptionalMatch { optional, .. } => { self.extract_variable_from_plan(optional) } - LogicalOperator::ListComprehension { input, .. } => { - self.extract_variable_from_plan(input) - } + LogicalOperator::ListComprehension { input, .. } + | LogicalOperator::PathExists { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::ComprehensionElement => Err(GraphError::PlanError { message: "a comprehension's element rows bind no pattern variable".to_string(), location: snafu::Location::new(file!(), line!(), column!()), diff --git a/crates/lance-csr/tests/test_engine.rs b/crates/lance-csr/tests/test_engine.rs index 116d64c..fe1b18d 100644 --- a/crates/lance-csr/tests/test_engine.rs +++ b/crates/lance-csr/tests/test_engine.rs @@ -287,3 +287,167 @@ async fn optional_match_where_pattern_reads_an_earlier_variable() { assert_eq!(column(2), [1, 1, 0, 2], "use_csr={use_csr}"); } } + +/// Pattern predicates walked on the CSR (`PathExists`) give the same rows as the +/// relational marking the engine plans without a CSR. +#[tokio::test] +async fn path_exists_matches_the_join_plan() { + let config = GraphConfig::builder() + .node_table("Person", "id") + .node_table("Post", "id") + .node_table("Tag", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .rel_table("CREATED", "Post", "Person", "src", "dst") + .rel_table("HAS_TAG", "Post", "Tag", "src", "dst") + .rel_table("INTEREST", "Person", "Tag", "src", "dst") + .build() + .unwrap(); + let ids = |ids: Vec| { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids))]).unwrap() + }; + let edges = |pairs: Vec<(i64, i64)>| { + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::Int64, false), + Field::new("dst", DataType::Int64, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from_iter_values(pairs.iter().map(|p| p.0))), + Arc::new(Int64Array::from_iter_values(pairs.iter().map(|p| p.1))), + ], + ) + .unwrap() + }; + // Persons 1-8 with a self-loop, a duplicate and mutual edges; person 8 knows nobody. + // Edges to missing nodes (99, 299, 98) match nothing, as nodes join to their tables. + let knows = vec![ + (1, 2), + (2, 1), + (1, 3), + (3, 3), + (3, 4), + (3, 4), + (4, 5), + (5, 6), + (6, 1), + (7, 2), + (8, 99), + ]; + let created: Vec<_> = (100..116).map(|p| (p, 1 + p % 7)).collect(); + // Posts 112-115 have no tags. + let has_tag: Vec<_> = (100..112) + .flat_map(|p| [(p, 200 + p % 5), (p, 200 + p % 3)]) + .chain([(113, 299)]) + .collect(); + let interest = vec![ + (1, 200), + (2, 201), + (2, 204), + (3, 202), + (5, 203), + (6, 200), + (6, 204), + (98, 203), + ]; + let sources = [ + ("Person", ids((1..=8).collect())), + ("Post", ids((100..116).collect())), + ("Tag", ids((200..205).collect())), + ("KNOWS", edges(knows)), + ("CREATED", edges(created)), + ("HAS_TAG", edges(has_tag)), + ("INTEREST", edges(interest)), + ]; + let queries = [ + // (query, walked on the CSR) + ( + "MATCH (a:Person) WHERE (a)-[:KNOWS]->(:Person) RETURN a.id AS a", + true, + ), + ( + "MATCH (a:Person) WHERE NOT (a)<-[:KNOWS]-() RETURN a.id AS a", + true, + ), + ( + "MATCH (a:Person), (b:Person) WHERE a.id < b.id AND NOT (a)-[:KNOWS]-(b) \ + RETURN a.id AS a, b.id AS b", + true, + ), + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE (a)<-[:KNOWS]-(b) \ + RETURN a.id AS a, b.id AS b", + true, + ), + ( + "MATCH (p:Post), (person:Person) \ + WHERE (p)-[:HAS_TAG]->(:Tag)<-[:INTEREST]-(person) RETURN p.id AS p, person.id AS q", + true, + ), + ( + "MATCH (person:Person)-[:KNOWS]->(friend:Person) WITH DISTINCT friend, person \ + OPTIONAL MATCH (friend)<-[:CREATED]-(post:Post) \ + WITH friend, person, collect(post) AS posts \ + RETURN person.id AS p, friend.id AS f, size(posts) AS n, \ + size([x IN posts WHERE (x)-[:HAS_TAG]->()<-[:INTEREST]-(person)]) AS c", + true, + ), + ( + "MATCH (person:Person)-[:KNOWS]->(friend:Person) WITH DISTINCT friend, person \ + OPTIONAL MATCH (friend)<-[:CREATED]-(post:Post) \ + WITH friend, person, post, \ + CASE WHEN (post)-[:HAS_TAG]->()<-[:INTEREST]-(person) THEN 1 END AS common \ + RETURN person.id AS p, friend.id AS f, count(post) AS n, count(common) AS c", + true, + ), + ( + "MATCH (person:Person)-[:KNOWS]->(friend:Person) WITH DISTINCT friend, person \ + OPTIONAL MATCH (friend)<-[:CREATED]-(post:Post) \ + WHERE (post)-[:HAS_TAG]->(:Tag)<-[:INTEREST]-(person) \ + RETURN person.id AS p, friend.id AS f, count(post) AS n", + true, + ), + // Not walked: one type twice, a named middle node, a WHERE inside the pattern. + ( + "MATCH (a:Person) WHERE (a)-[:KNOWS]->()-[:KNOWS]->(:Person) RETURN a.id AS a", + false, + ), + ( + "MATCH (a:Person) WHERE (a)<-[:CREATED]-(m:Post)-[:HAS_TAG]->(:Tag) RETURN a.id AS a", + false, + ), + ]; + let mut engines = Vec::new(); + for use_csr in [false, true] { + let sources = sources + .iter() + .map(|(name, batch)| (name.to_string(), EngineSource::Arrow(batch.clone()))) + .collect(); + engines.push( + CypherEngine::new(config.clone(), sources, use_csr) + .await + .unwrap(), + ); + } + for (query, walked) in queries { + let mut rows = Vec::new(); + for engine in &engines { + let batch = engine.execute(query, HashMap::new()).await.unwrap(); + let mut lines: Vec = arrow::util::pretty::pretty_format_batches(&[batch]) + .unwrap() + .to_string() + .lines() + .map(str::to_string) + .collect(); + lines.sort(); + rows.push(lines); + } + assert_eq!(rows[0], rows[1], "{query}"); + assert!(rows[0].len() > 4, "no rows from {query}"); + let plan = engines[1].explain(query, HashMap::new()).await.unwrap(); + assert_eq!(plan.contains("PathExistsExec"), walked, "{query}\n{plan}"); + let plan = engines[0].explain(query, HashMap::new()).await.unwrap(); + assert!(!plan.contains("PathExistsExec"), "{query}\n{plan}"); + } +} diff --git a/docs/csr.md b/docs/csr.md index 6e9e910..d6ed058 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -133,6 +133,7 @@ When those hold, it covers: | (a)-[:LIKES_POST|LIKES_COMMENT]->(m) | :material-check: one expansion per branch | | `OPTIONAL MATCH (a)-[:KNOWS]->(b)` | :material-check: | | `shortestPath((a)-[:KNOWS*]-(b))` | :material-check: the search walks the CSR adjacency | +| `WHERE NOT (b)-[:KNOWS]->(a)`, `(p)-[:HAS_TAG]->()<-[:INTEREST]-(person)` | :material-check: walked per row from a bound node, through anonymous nodes and relationships of distinct types ([`PathExists`](internals/operators.md#pathexists-pattern-predicates-walked-in-place)) | | `(a)-[k:KNOWS]->(b)` | :material-check: `k`'s other columns are fetched by edge row (single hops only) | | `(a)-[k:KNOWS {since: 2020}]->(b)` | :material-close: hash join | | endpoints of type `string`, `float` or `uint64` | :material-close: hash join | diff --git a/docs/internals/index.md b/docs/internals/index.md index fcb5fe7..2532ed4 100644 --- a/docs/internals/index.md +++ b/docs/internals/index.md @@ -13,7 +13,7 @@ the data structure those pieces read. | Page | What it covers | | --- | --- | -| [Custom operators](operators.md) | `CsrExpand`, `NodeLookup`, `ShortestPath`, `PathMatch`, `Apply` and `KeyStatistics`, with their plans | +| [Custom operators](operators.md) | `CsrExpand`, `NodeLookup`, `PathExists`, `ShortestPath`, `PathMatch`, `Apply` and `KeyStatistics`, with their plans | | [Optimizer rules and rewrites](optimizer.md) | The two logical optimizer rules, graph rewrites before DataFusion sees the plan, and join statistics | | [Indexes and refresh](indexes.md) | How the CSR adjacency and node indexes are laid out, and how a refresh diffs Lance versions | | [Execution plumbing](execution-plumbing.md) | Partitioning, the plan template cache, deadlines and Cypher functions | diff --git a/docs/internals/operators.md b/docs/internals/operators.md index 0095f47..3cc285d 100644 --- a/docs/internals/operators.md +++ b/docs/internals/operators.md @@ -150,7 +150,67 @@ ProjectionExec: expr=[b__name@2 as b.name, __agg_0@1 as n] Unlike a join back to the node table, its 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. +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. + +## `PathExists`: pattern predicates walked in place + +A pattern predicate such as `WHERE NOT (b)-[:KNOWS]->(a)` asks, for each row, whether a path exists. Relationally, +the planner answers it by planning everything before the predicate a second time, matching the pattern from that +copy, and left-joining a mark back onto the rows. That doubles the plan beneath every predicate, and DataFusion plans +and runs both copies. When the pattern qualifies, `PathExistsExec` instead walks the CSR from each row's start node, +stops at the first path it finds, and appends the same mark (`true`, or null for no path), so the filter or `CASE` +above it is unchanged: + +```cypher +MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE NOT (b)-[:KNOWS]->(a) RETURN a.name, b.name +``` + +=== "use_csr=True" + + ```text + ProjectionExec: expr=[a__name@0 as a.name, b__name@1 as b.name] + FilterExec: __exists_1@2 IS NULL, projection=[a__name@0, b__name@1] + ProjectionExec: expr=[a__name@1 as a__name, b__name@5 as b__name, __exists_1@6 as __exists_1] + PathExistsExec: start=b__id, hops=1, match=a__id + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + CsrExpandExec: key=a__id, edges=4, take=b[b__id*, b__name] + ProjectionExec: expr=[id@0 as a__id, name@1 as a__name] + KeyStatisticsExec: distinct=[id=4] + DataSourceExec: partitions=1, partition_sizes=[1] + ``` + +=== "use_csr=False" + + ```text + ProjectionExec: expr=[a__name@0 as a.name, b__name@1 as b.name] + FilterExec: __exists_1@2 IS NULL, projection=[a__name@0, b__name@1] + HashJoinExec: mode=CollectLeft, join_type=Left, on=[(b__id@2, __optional__b__id@0), (a__id@0, __optional__a__id@1)], ... + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(knows_1__dst@2, b__id@0)], ... + ... (the MATCH, planned once for the rows) + ProjectionExec: expr=[b__id@0 as __optional__b__id, a__id@3 as __optional__a__id, __exists_1@6 as __exists_1] + AggregateExec: mode=FinalPartitioned, gby=[b__id@0 as b__id, ..., __exists_1@6 as __exists_1], aggr=[] + ... + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(b__id@0, knows_3__src@0), (a__id@3, knows_3__dst@1)], ... + ... + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(knows_2__dst@3, b__id@0)], ... + ... (the MATCH again, as the seed of the predicate) + ``` + +A predicate qualifies when it is one path that: + +- starts at a node bound earlier in the query; +- passes through anonymous nodes, with no properties and only the label the relationship implies; +- uses anonymous relationships of one type each, all different, so no two hops can bind the same relationship, and + none variable-length; +- ends at an anonymous node or at another bound node (`match=` in the plan); +- has no `WHERE` of its own. + +Every node the walk reaches must be in its label's node index, as the relational plan joins each node to its table, so +an edge to a missing node matches nothing. A predicate that doesn't qualify, a hop or label without a CSR, or an +endpoint that isn't an integer id keeps the relational plan, which the logical operator carries as its fallback. Both +give the same rows. The walk keeps no set of visited nodes, so a long chain of high-degree hops ending at a bound node +costs the product of their degrees per row. ## `ShortestPath`: breadth-first search over the adjacency diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 5b7f722..1bec696 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -88,6 +88,22 @@ unrolled into the sequences of types whose endpoint labels chain, and each seque Combinations that the configuration can't pin down are an error rather than a guess: joining an edge type to the wrong node table silently matches wrong rows whenever two labels share ids. +### List comprehensions keep each row once + +`[p IN posts WHERE ...]` is evaluated relationally: the rows are numbered, each row's list becomes one row per element, +the filter (and any pattern predicate, through [`PathExists`](operators.md#pathexists-pattern-predicates-walked-in-place) +or its fallback) runs on those element rows, and the kept elements are counted or collected back per row number. +Two details keep this from repeating each row's data once per element: + +- The list is unnested beside its element positions, so an element row holds its own item rather than a copy of the + whole list. +- A list read from a column, such as `posts` from an earlier `collect(post)`, is not a grouping key on the way back. + It is rebuilt once per row from its elements, in order (null for a null list, empty for an empty one), so when the + comprehension doesn't read the list itself, projection pushdown drops it before the unnest. + +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. + ## 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_comprehension_count.py b/python/python/tests/test_comprehension_count.py index ad04322..9ece105 100644 --- a/python/python/tests/test_comprehension_count.py +++ b/python/python/tests/test_comprehension_count.py @@ -56,7 +56,9 @@ def test_count_matches_list_reference(engine, predicate, expected): assert engine.execute(query).to_pylist() == rows assert engine.execute_batches(query).read_all().to_pylist() == rows assert engine.execute(reference).to_pylist() == rows - assert "array_agg(" not in engine.explain(query).lower() + # The unused rebuild of `p.values` is pruned before execution. + physical = engine.explain(query).lower().split("physical_plan")[1] + assert "array_agg(" not in physical assert "array_agg(" in engine.explain(reference).lower()