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