From 62a5af87b4a63b8031a92ec271f1c0495c216c3e Mon Sep 17 00:00:00 2001 From: prrao87 Date: Tue, 3 Nov 2026 16:15:48 -0500 Subject: [PATCH] perf(csr): expand one-hop OPTIONAL MATCH as an outer CSR hop A single-hop OPTIONAL MATCH with no WHERE of its own now runs as one outer CsrExpand over its input, instead of planning the input again as a seed and left-joining the optional side back. Input rows with no neighbour keep one row with null hop columns. Filters never move into an outer hop, and it is never counted in place. Part of #18. --- crates/lance-csr/src/csr_expand.rs | 110 +++++++++++++++++- .../datafusion_planner/builder/expand_ops.rs | 2 +- .../builder/join_builder.rs | 67 +++++++++++ docs/execution.md | 2 +- docs/internals/operators.md | 1 + python/python/tests/test_engine_optional.py | 22 +++- 6 files changed, 200 insertions(+), 4 deletions(-) diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 770117a..cbf4a7e 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -1493,6 +1493,9 @@ pub struct CsrExpandNode { /// 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]>, + /// OPTIONAL MATCH: an input row with no neighbour is kept once, with nulls for the + /// columns this node adds (see [`CsrExpandNode::into_outer`]). + outer: bool, schema: DFSchemaRef, } @@ -1528,10 +1531,52 @@ impl CsrExpandNode { closing_column: None, traversal: None, degrees: None, + outer: false, schema: Arc::new(schema), }) } + /// `plan`, a single hop that fetches its target, as a left outer expansion; `None` + /// for any other plan (a closing edge, a walk, a moved filter, a degree count). + pub(crate) fn outer_plan(plan: &LogicalPlan) -> Option { + let LogicalPlan::Extension(ext) = plan else { + return None; + }; + let node = ext.node.as_any().downcast_ref::()?; + let simple = node.target.is_some() + && node.target_filter.is_none() + && node.edge_filter.is_none() + && node.closing_column.is_none() + && node.traversal.is_none() + && node.degrees.is_none(); + let node = node.clone().into_outer().ok().filter(|_| simple)?; + Some(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } + + /// This hop as a left outer expansion: input rows without a neighbour keep one row, + /// whose added columns are null. Replaces OPTIONAL MATCH's seed and left join for a + /// single hop. Filters must not move into an outer node: they would null a row out + /// instead of dropping it. + fn into_outer(mut self) -> Result { + let width = self.input.schema().fields().len(); + let fields: Vec = self.schema.fields()[width..] + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)) + .collect(); + let added = DFSchema::try_from(Schema::new(fields)) + .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.outer = true; + Ok(self) + } + /// `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 @@ -1556,6 +1601,7 @@ impl CsrExpandNode { }; let node = ext.node.as_any().downcast_ref::()?; let simple = node.edge.is_none() + && !node.outer && node.closing_column.is_none() && node.traversal.is_none() && node.degrees.is_none(); @@ -1622,6 +1668,9 @@ impl CsrExpandNode { if let Some(fields) = &self.degrees { node.count_degrees(fields.clone())?; } + if self.outer { + node = node.into_outer()?; + } Ok(node) } @@ -1719,6 +1768,7 @@ impl PartialEq for CsrExpandNode { _ => false, } && self.degrees == other.degrees + && self.outer == other.outer } } impl Eq for CsrExpandNode {} @@ -1740,6 +1790,7 @@ impl Hash for CsrExpandNode { .as_ref() .map(|f| [f[0].name(), f[1].name()]) .hash(state); + self.outer.hash(state); } } impl PartialOrd for CsrExpandNode { @@ -1850,6 +1901,9 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { if let Some([count, loops]) = &self.degrees { write!(f, ", degrees=[{}, {}]", count.name(), loops.name())?; } + if self.outer { + write!(f, ", outer")?; + } Ok(()) } @@ -1883,6 +1937,8 @@ pub struct CsrExpandExec { traversal: Option>, /// Count each input row's neighbours instead of emitting them. degrees: bool, + /// Keep each input row without a neighbour once, with null added columns. + outer: bool, schema: SchemaRef, properties: Arc, } @@ -1900,6 +1956,7 @@ impl CsrExpandExec { closing_column: Option, traversal: Option>, degrees: bool, + outer: bool, schema: SchemaRef, ) -> Self { let properties = Arc::new(PlanProperties::new( @@ -1921,6 +1978,7 @@ impl CsrExpandExec { closing_column, traversal, degrees, + outer, schema, properties, } @@ -1941,6 +1999,8 @@ struct BatchExpansion { edge_mask: Option>, /// Per input row, in degree mode: neighbours kept, and those equal to the key. degrees: Option<(Vec, Vec)>, + /// Per input row, in outer mode: whether it has produced a row yet. + matched: Option>, } impl BatchExpansion { @@ -1964,9 +2024,36 @@ impl BatchExpansion { edge_rows: false, edge_mask: None, degrees: None, + matched: None, }) } + /// In outer mode, the input rows that produced no row, each with null added + /// columns; `None` otherwise or when every row matched. + fn unmatched_batch( + &self, + schema: &SchemaRef, + ) -> datafusion::common::Result> { + let Some(matched) = &self.matched else { + return Ok(None); + }; + let keep = arrow::array::BooleanArray::from_iter(matched.iter().map(|&m| Some(!m))); + let rows = keep.true_count(); + if rows == 0 { + return Ok(None); + } + let mut columns = self + .batch + .columns() + .iter() + .map(|c| filter_array(c.as_ref(), &keep)) + .collect::, _>>()?; + for field in &schema.fields()[columns.len()..] { + columns.push(arrow::array::new_null_array(field.data_type(), rows)); + } + Ok(Some(RecordBatch::try_new(schema.clone(), columns)?)) + } + /// 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> { @@ -2327,6 +2414,9 @@ impl DisplayAs for CsrExpandExec { let names: Vec<_> = self.schema.fields().iter().rev().take(2).collect(); write!(f, ", degrees=[{}, {}]", names[1].name(), names[0].name())?; } + if self.outer { + write!(f, ", outer")?; + } Ok(()) } } @@ -2390,6 +2480,7 @@ impl ExecutionPlan for CsrExpandExec { self.closing_column.clone(), self.traversal.clone(), self.degrees, + self.outer, self.schema.clone(), ))) } @@ -2448,6 +2539,7 @@ impl ExecutionPlan for CsrExpandExec { let out_schema = schema.clone(); let traversal = self.traversal.clone(); let degrees = self.degrees; + let outer = self.outer; // 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)| { @@ -2472,6 +2564,7 @@ impl ExecutionPlan for CsrExpandExec { expansion.edge_rows = edge.is_some(); let rows = expansion.keys.len(); expansion.degrees = degrees.then(|| (vec![0; rows], vec![0; rows])); + expansion.matched = outer.then(|| vec![false; rows]); pending = Some(expansion); } let expansion = pending.as_mut().expect("input batch just set"); @@ -2497,6 +2590,9 @@ impl ExecutionPlan for CsrExpandExec { }; if chunk.rows.is_empty() { let done = pending.take().expect("input batch just set"); + if let Some(batch) = done.unmatched_batch(&schema)? { + return Ok(Some((batch, (input, None)))); + } if let Some(batch) = done.degree_batch(&schema)? { return Ok(Some((batch, (input, None)))); } @@ -2558,6 +2654,12 @@ impl ExecutionPlan for CsrExpandExec { } continue; } + if let Some(matched) = expansion.matched.as_mut() { + // An outer node has no moved filter, so every chunk row is output. + for &row in &chunk.rows { + matched[row as usize] = true; + } + } let batch = materialize( &expansion.batch, chunk, @@ -2649,6 +2751,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { node.closing_column.clone(), node.traversal.clone(), node.degrees.is_some(), + node.outer, Arc::new(node.schema.as_arrow().clone()), )))) } @@ -2716,7 +2819,12 @@ impl OptimizerRule for PushCsrTargetFilters { let LogicalPlan::Extension(ext) = below else { return Ok(Transformed::no(node)); }; - let Some(csr) = ext.node.as_any().downcast_ref::() else { + let Some(csr) = ext + .node + .as_any() + .downcast_ref::() + .filter(|csr| !csr.outer) + else { return Ok(Transformed::no(node)); }; let target_columns = csr.target.as_ref().map(TargetLookup::raw_columns); 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 318d506..a417339 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs @@ -58,7 +58,7 @@ impl DataFusionPlanner { /// otherwise relationship scan + joins. Shared by single hops and by every hop of /// a variable-length path. #[allow(clippy::too_many_arguments)] - fn expand_plan( + pub(crate) fn expand_plan( &self, ctx: &mut PlanningContext, left_plan: LogicalPlan, diff --git a/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs b/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs index 427932b..a109694 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs @@ -25,6 +25,9 @@ impl DataFusionPlanner { use datafusion::logical_expr::Expr; let left = self.build_operator(ctx, input)?; + if let Some(plan) = self.try_outer_hop(ctx, &left, input, optional, keys)? { + return Ok(plan); + } let right = self.build_operator(ctx, optional)?; let renamed = |name: &str| format!("__optional__{name}"); let left_names: std::collections::HashSet = left @@ -73,6 +76,70 @@ impl DataFusionPlanner { .map_err(|e| self.plan_error("Failed to project OPTIONAL MATCH", e)) } + /// OPTIONAL MATCH of one hop from its only shared variable, seeded from `input` + /// itself: one outer CSR expansion of `left` instead of planning `input` again as + /// the seed and joining back. `None` (with `ctx` untouched) when the hop does not + /// take the CSR with a fused target. + fn try_outer_hop( + &self, + ctx: &mut PlanningContext, + left: &LogicalPlan, + input: &LogicalOperator, + optional: &LogicalOperator, + keys: &[String], + ) -> Result> { + let LogicalOperator::Expand { + input: seed, + source_variable, + target_variable, + target_label, + relationship_types, + direction, + relationship_variable, + properties, + target_properties, + } = optional + else { + return Ok(None); + }; + let LogicalOperator::Distinct { input: seed } = seed.as_ref() else { + return Ok(None); + }; + let LogicalOperator::Project { + input: source, + projections, + } = seed.as_ref() + else { + return Ok(None); + }; + let from_shared = matches!( + projections.as_slice(), + [ProjectionItem { expression: crate::ast::ValueExpression::Variable(v), alias: None }] + if v == source_variable + ); + if keys.len() != 1 || !from_shared || source.as_ref() != input { + return Ok(None); + } + let saved = ctx.relationship_instance_idx.clone(); + let hop = self.expand_plan( + ctx, + left.clone(), + source_variable, + target_variable, + target_label, + relationship_types, + direction, + relationship_variable.as_deref(), + properties, + target_properties, + )?; + let outer = crate::csr_expand::CsrExpandNode::outer_plan(&hop); + if outer.is_none() { + ctx.relationship_instance_idx = saved; + } + Ok(outer) + } + /// Build a join between two logical operators /// /// The join type and keys are determined by: diff --git a/docs/execution.md b/docs/execution.md index 4f8ca27..5b7b02e 100644 --- a/docs/execution.md +++ b/docs/execution.md @@ -22,7 +22,7 @@ Source paths below are in `crates/lance-csr/src/`. Multi-label nodes and relationship alternation expand into `UNION ALL` branches. Semantics are unchanged; only the order of work is. 4. **Plan the graph** (`logical_plan.rs`). Graph operators describe the query: `ScanByLabel`, `Expand` for a - hop, `VariableLengthExpand` (unrolled into one branch per length), `OptionalMatch` as a left join, + hop, `VariableLengthExpand` (unrolled into one branch per length), `OptionalMatch` as a left join (an outer CSR hop when it is a single hop), `ShortestPath`, and `Apply` for a read-only CALL subquery. 5. **Plan for DataFusion** (`datafusion_planner/`). Each graph operator becomes DataFusion relational operators. A hop becomes either two hash joins or, when the engine has a [CSR](csr.md) for it, a diff --git a/docs/internals/operators.md b/docs/internals/operators.md index 8638d90..8c255fd 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` | +| Outer expansion for a one-hop `OPTIONAL MATCH` with no `WHERE` of its own: an input row with no neighbour keeps one row, with nulls in the columns the hop adds. No filter moves into an outer hop, because that would turn a row the filter drops into a row of nulls | `outer` | Planning the input a second time to start the optional side, then a left join back | | Degrees: one row per input node with the number of its neighbours that pass the target filter, and how many of them are the node itself, for the [two-hop count](optimizer.md#two-hop-counts-from-degrees) | `degrees=[__n1, __s1]` | A hash aggregate over one row per edge | A hop that uses several of these at once: diff --git a/python/python/tests/test_engine_optional.py b/python/python/tests/test_engine_optional.py index 16391ad..8eb4f69 100644 --- a/python/python/tests/test_engine_optional.py +++ b/python/python/tests/test_engine_optional.py @@ -57,6 +57,18 @@ "OPTIONAL MATCH (b)-[:KNOWS]->(c:Person {id: 1}) RETURN b.id, c.id", [(1, N), (2, N), (3, N)], ), + # Repeated input rows each keep their own matches. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "OPTIONAL MATCH (b)-[:KNOWS]->(c:Person) RETURN b.id, c.id", + [(1, 2), (1, 2), (1, 3), (2, 3), (2, 3), (3, 3), (3, 3), (3, 3)], + ), + # A WHERE after the OPTIONAL MATCH drops the null rows instead of nulling more. + ( + "MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS]->(b:Person) " + "WITH a, b WHERE b.age > 30 RETURN a.id, b.id", + [(1, 3), (2, 3), (3, 3)], + ), ] @@ -75,12 +87,20 @@ def test_optional_match(use_csr, query, expected): assert rows(got) == sorted(expected, key=lambda t: [(v is None, v) for v in t]) -def test_optional_side_expands_on_csr_and_left_joins(): +def test_optional_hop_is_one_outer_csr_expansion(): + # One hop: an outer expansion of the input, with no second copy of it and no join. engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=True) plan = engine.explain( "MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS]->(b:Person) RETURN a.id, b.id" ) physical = plan[plan.index("physical_plan") :] + assert "outer" in physical and "Join" not in physical + # Its own WHERE keeps the CSR expansion on the optional side and the left join. + plan = engine.explain( + "MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS]->(b:Person) WHERE b.age > 30 " + "RETURN a.id, b.id" + ) + physical = plan[plan.index("physical_plan") :] assert "CsrExpandExec" in physical and "join_type=Left" in physical