diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index da229fe..885ef7c 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -59,7 +59,10 @@ use datafusion::physical_plan::{ use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; use futures::StreamExt; use lance::dataset::{Dataset, ProjectionRequest}; + +mod node_lookup; use lance_table::format::Fragment; +pub use node_lookup::NodeLookupNode; /// Yield after bounded CPU work without requiring a particular async runtime. pub(crate) async fn cooperative_yield() { @@ -691,6 +694,14 @@ impl NodeRows { self.positions.len() } + pub(crate) fn schema(&self) -> &SchemaRef { + &self.schema + } + + pub(crate) fn in_memory(&self) -> bool { + matches!(self.data, TableData::Memory(_)) + } + pub(crate) fn position(&self, id: i64) -> Option { self.positions.get(&id).copied() } @@ -2386,6 +2397,9 @@ impl ExtensionPlanner for CsrExtensionPlanner { physical_inputs, ))); } + if let Some(node) = node.as_any().downcast_ref::() { + return Ok(Some(node_lookup::plan_node_lookup(node, physical_inputs))); + } let Some(node) = node.as_any().downcast_ref::() else { return Ok(None); }; @@ -2520,7 +2534,8 @@ impl OptimizerRule for PushCsrTargetFilters { } } -/// Optimizer rule: drop fused target columns that nothing in the plan uses. +/// Optimizer rule: drop fused target columns, and node lookup columns, that nothing +/// in the plan uses. /// /// DataFusion's projection pushdown prunes an extension node's inputs but not its /// own outputs, so without this every target column would be fetched per row. @@ -2550,6 +2565,16 @@ impl OptimizerRule for PruneCsrTargetColumns { }))); } } + 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)) }) @@ -2569,14 +2594,16 @@ fn used_columns(plan: &LogicalPlan) -> datafusion::common::Result { - match ext.node.as_any().downcast_ref::() { - Some(csr) => { - // Read by the node itself, not through an expression. - used.insert(csr.key_column.clone()); - used.extend(csr.closing_column.clone()); - true - } - None => false, + if let Some(csr) = ext.node.as_any().downcast_ref::() { + // Read by the node itself, not through an expression. + used.insert(csr.key_column.clone()); + used.extend(csr.closing_column.clone()); + true + } else if let Some(lookup) = ext.node.as_any().downcast_ref::() { + used.insert(lookup.key_column().to_string()); + true + } else { + false } } LogicalPlan::Projection(_) diff --git a/crates/lance-csr/src/csr_expand/node_lookup.rs b/crates/lance-csr/src/csr_expand/node_lookup.rs new file mode 100644 index 0000000..019c633 --- /dev/null +++ b/crates/lance-csr/src/csr_expand/node_lookup.rs @@ -0,0 +1,408 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! A node's columns fetched by id through its node index, for the rows of `input`. +//! +//! The planner groups `WITH b, ...` by `b`'s id alone when `b`'s label has a node index +//! (unique ids), then adds `b`'s other columns with this node: one lookup per group +//! instead of carrying them through the aggregate for every input row. Unlike a join +//! back to the node table, the cost follows the rows looked up, not the table's size. + +use super::*; + +/// `input` plus the columns of the node whose id is in `key_column`. A null or unknown +/// id (an unmatched OPTIONAL MATCH) gets nulls. +#[derive(Debug, Clone)] +pub struct NodeLookupNode { + input: LogicalPlan, + key_column: String, + variable: String, + rows: Arc, + /// Output fields (`b__name`), nullable. + fields: Vec, + /// For each output field, its column index in `rows` (ascending). + columns: Vec, + schema: DFSchemaRef, +} + +impl NodeLookupNode { + /// `wanted` pairs each output name (`b__name`) with the node table's column name. + pub fn try_new( + input: LogicalPlan, + key_column: String, + variable: String, + rows: Arc, + wanted: &[(String, String)], + ) -> Result { + let mut columns: Vec<(usize, String)> = wanted + .iter() + .map(|(name, column)| { + rows.schema + .index_of(column) + .map(|c| (c, name.clone())) + .map_err(|e| plan_error(format!("NodeLookup column: {}", e))) + }) + .collect::>()?; + columns.sort_unstable(); + let fields = columns + .iter() + .map(|(c, name)| Field::new(name, rows.schema.field(*c).data_type().clone(), true)) + .collect(); + Self::with_fields( + input, + key_column, + variable, + rows, + fields, + columns.into_iter().map(|(c, _)| c).collect(), + ) + } + + fn with_fields( + input: LogicalPlan, + key_column: String, + variable: String, + rows: Arc, + fields: Vec, + columns: Vec, + ) -> Result { + let added = DFSchema::try_from(Schema::new(fields.clone())) + .map_err(|e| plan_error(format!("NodeLookup schema: {}", e)))?; + let schema = input + .schema() + .join(&added) + .map_err(|e| plan_error(format!("NodeLookup schema: {}", e)))?; + Ok(Self { + input, + key_column, + variable, + rows, + fields, + columns, + schema: Arc::new(schema), + }) + } + + /// The same lookup with only the fetched columns in `used`; `Some(input)` when none + /// remain, so the lookup can be dropped. `None` when every column is used. + pub(super) fn pruned( + &self, + used: &HashSet, + ) -> Option> { + let keep: Vec = (0..self.fields.len()) + .filter(|&i| used.contains(self.fields[i].name())) + .collect(); + if keep.len() == self.fields.len() { + return None; + } + if keep.is_empty() { + return Some(Err(self.input.clone())); + } + Self::with_fields( + self.input.clone(), + self.key_column.clone(), + self.variable.clone(), + self.rows.clone(), + keep.iter().map(|&i| self.fields[i].clone()).collect(), + keep.iter().map(|&i| self.columns[i]).collect(), + ) + .ok() + .map(Ok) + } + + pub(super) fn key_column(&self) -> &str { + &self.key_column + } + + fn describe(&self) -> String { + let names: Vec<&str> = self.fields.iter().map(|f| f.name().as_str()).collect(); + format!( + "key={}, take={}[{}]", + self.key_column, + self.variable, + names.join(", ") + ) + } +} + +// Logical nodes must be comparable; the node index compares by identity. +impl PartialEq for NodeLookupNode { + fn eq(&self, other: &Self) -> bool { + self.input == other.input + && self.key_column == other.key_column + && self.variable == other.variable + && self.columns == other.columns + && Arc::ptr_eq(&self.rows, &other.rows) + } +} +impl Eq for NodeLookupNode {} +impl Hash for NodeLookupNode { + fn hash(&self, state: &mut H) { + self.input.hash(state); + self.key_column.hash(state); + self.variable.hash(state); + self.columns.hash(state); + } +} +impl PartialOrd for NodeLookupNode { + fn partial_cmp(&self, other: &Self) -> Option { + (&self.key_column, &self.variable).partial_cmp(&(&other.key_column, &other.variable)) + } +} + +impl UserDefinedLogicalNodeCore for NodeLookupNode { + fn name(&self) -> &str { + "NodeLookup" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + /// Filters on the fetched columns must stay above the lookup. + fn prevent_predicate_push_down_columns(&self) -> HashSet { + self.fields.iter().map(|f| f.name().clone()).collect() + } + + /// The input only needs the requested input columns plus the key. + fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option>> { + let input_width = self.input.schema().fields().len(); + let key = self + .input + .schema() + .index_of_column_by_name(None, &self.key_column)?; + let mut needed: Vec = output_columns + .iter() + .copied() + .filter(|&i| i < input_width) + .chain(std::iter::once(key)) + .collect(); + needed.sort_unstable(); + needed.dedup(); + Some(vec![needed]) + } + + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "NodeLookup: {}", self.describe()) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + inputs: Vec, + ) -> datafusion::common::Result { + let input = inputs.into_iter().next().expect("NodeLookup has one input"); + Self::with_fields( + input, + self.key_column.clone(), + self.variable.clone(), + self.rows.clone(), + self.fields.clone(), + self.columns.clone(), + ) + .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string())) + } +} + +/// Executes [`NodeLookupNode`]: per batch, the key's node rows, then one take of the +/// wanted columns. +// shortcut: no read-once switch like `CsrExpandExec`'s targets, as the planner only looks +// up in-memory node tables; add one if Lance node tables are looked up. +#[derive(Debug)] +pub struct NodeLookupExec { + input: Arc, + key_column: String, + variable: String, + rows: Arc, + columns: Vec, + schema: SchemaRef, + properties: Arc, +} + +impl NodeLookupExec { + fn new(node: &NodeLookupNode, input: Arc) -> Self { + let schema: SchemaRef = Arc::new(node.schema.as_arrow().clone()); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + input, + key_column: node.key_column.clone(), + variable: node.variable.clone(), + rows: node.rows.clone(), + columns: node.columns.clone(), + schema, + properties, + } + } + + /// The batch's rows with their node's columns appended. + async fn lookup( + batch: RecordBatch, + key_column: &str, + rows: &NodeRows, + columns: &[usize], + schema: &SchemaRef, + ) -> datafusion::common::Result { + let external = |e: GraphError| datafusion::common::DataFusionError::External(Box::new(e)); + let ids = int64_column(&batch, key_column).map_err(external)?; + let mut positions = Vec::with_capacity(ids.len()); + // Row -> its slot in `positions`, or None for a null or unknown id. + let slots: Vec> = (0..ids.len()) + .map(|i| { + let position = ids + .is_valid(i) + .then(|| rows.position(ids.value(i))) + .flatten()?; + positions.push(position); + Some(positions.len() as u32 - 1) + }) + .collect(); + let fetched = rows.take(&positions, columns).await.map_err(external)?; + let fetched = if positions.len() == ids.len() { + fetched + } else { + let slots = UInt32Array::from(slots); + fetched + .iter() + .map(|array| take(array.as_ref(), &slots, None)) + .collect::>()? + }; + let mut output = batch.columns().to_vec(); + output.extend(fetched); + Ok(RecordBatch::try_new(schema.clone(), output)?) + } +} + +impl DisplayAs for NodeLookupExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + let names: Vec<&str> = self.schema.fields() + [self.schema.fields().len() - self.columns.len()..] + .iter() + .map(|f| f.name().as_str()) + .collect(); + write!( + f, + "NodeLookupExec: key={}, take={}[{}]", + self.key_column, + self.variable, + names.join(", ") + ) + } +} + +impl ExecutionPlan for NodeLookupExec { + fn name(&self) -> &str { + "NodeLookupExec" + } + + 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("NodeLookupExec has one child"); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(self.schema.clone()), + input.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input, + key_column: self.key_column.clone(), + variable: self.variable.clone(), + rows: self.rows.clone(), + columns: self.columns.clone(), + schema: self.schema.clone(), + properties, + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion::common::Result { + let input = self.input.execute(partition, context)?; + let (rows, columns, schema) = + (self.rows.clone(), self.columns.clone(), self.schema.clone()); + let key_column = self.key_column.clone(); + let stream = input.then(move |batch| { + let (rows, columns, schema) = (rows.clone(), columns.clone(), schema.clone()); + let key_column = key_column.clone(); + async move { Self::lookup(batch?, &key_column, &rows, &columns, &schema).await } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } +} + +pub(super) fn plan_node_lookup( + node: &NodeLookupNode, + physical_inputs: &[Arc], +) -> Arc { + Arc::new(NodeLookupExec::new(node, physical_inputs[0].clone())) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::StringArray; + + #[tokio::test] + async fn looks_up_columns_with_nulls_for_unknown_ids() { + let people = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from(vec![10, 20, 30])) as ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + ), + ]) + .unwrap(); + let rows = NodeRows::from_batch(people, "id").unwrap().unwrap(); + let input = RecordBatch::try_from_iter([( + "b__id", + Arc::new(Int64Array::from(vec![Some(30), None, Some(99), Some(10)])) as ArrayRef, + )]) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("b__id", DataType::Int64, true), + Field::new("b__name", DataType::Utf8, true), + ])); + let out = NodeLookupExec::lookup(input, "b__id", &rows, &[1], &schema) + .await + .unwrap(); + let names = out.column(1).as_string::(); + let names: Vec> = (0..names.len()) + .map(|i| names.is_valid(i).then(|| names.value(i))) + .collect(); + assert_eq!(names, vec![Some("c"), None, None, Some("a")]); + } +} 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 b91f287..c3f70e0 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs @@ -3,11 +3,26 @@ //! Aggregation operations: Projection with aggregates and grouping +use std::sync::Arc; + use crate::ast::ValueExpression; +use crate::csr_expand::{NodeLookupNode, NodeRows}; +use crate::datafusion_planner::analysis::PlanningContext; use crate::datafusion_planner::DataFusionPlanner; use crate::error::Result; use crate::logical_plan::*; -use datafusion::logical_expr::{col, LogicalPlan, LogicalPlanBuilder}; +use datafusion::common::Column; +use datafusion::logical_expr::{col, Expr, Extension, LogicalPlan, LogicalPlanBuilder}; + +/// A node grouped by its id alone, whose other columns are looked up after the aggregate. +struct LateNode { + variable: String, + rows: Arc, + /// The id column, as in the input (`b__id`). + id: String, + /// The node's other columns (`b__name`), with their names in the node table. + columns: Vec<(String, String)>, +} /// Replace each aggregate inside `expr` with a reference to a new column `__agg_`, /// pushing the aggregate to `aggregates`, so `head(collect(x))` or `count(x) + 1` aggregate @@ -95,8 +110,60 @@ 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 + /// 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( + &self, + ctx: &PlanningContext, + expr: &ValueExpression, + columns: &[Expr], + ) -> 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() { + return None; + } + let id = self.id_column(ctx, var)?; + 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 { + found_id = true; + continue; + } + let suffix = alias.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())); + } + (found_id && !others.is_empty()).then(|| LateNode { + variable: var.to_lowercase(), + rows, + id, + columns: others, + }) + } + pub(crate) fn build_project_with_aggregates( &self, + ctx: &PlanningContext, input_plan: LogicalPlan, projections: &[ProjectionItem], ordering: &[datafusion::logical_expr::SortExpr], @@ -117,9 +184,15 @@ impl DataFusionPlanner { }) .collect(); let projections = bound.as_slice(); + 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) { + group_exprs.push(col(Column::new_unqualified(&node.id))); + late.push(node); + continue; + } group_exprs.extend(columns.into_iter().map(|c| match c { datafusion::logical_expr::Expr::Alias(a) => *a.expr, other => other, @@ -191,10 +264,13 @@ impl DataFusionPlanner { } } - let aggregate = LogicalPlanBuilder::from(input_plan) + let mut aggregate = LogicalPlanBuilder::from(input_plan) .aggregate(group_exprs, agg_exprs) .and_then(|b| b.build()) .map_err(|e| self.plan_error("Failed to build aggregate", e))?; + for node in late { + aggregate = self.fetch_late_node(aggregate, node)?; + } let final_projection = final_projection .into_iter() .map(|expr| self.resolve_expr(expr, aggregate.schema())) @@ -207,6 +283,18 @@ impl DataFusionPlanner { } } +impl DataFusionPlanner { + /// `plan` with `node`'s other columns looked up by its id, under their original names. + /// A null node (an unmatched OPTIONAL MATCH) gets null columns. + fn fetch_late_node(&self, plan: LogicalPlan, node: LateNode) -> Result { + let lookup = + NodeLookupNode::try_new(plan, node.id, node.variable, node.rows, &node.columns)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(lookup), + })) + } +} + #[cfg(test)] mod tests { use crate::ast::ValueExpression; 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 0f08502..c235aa1 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs @@ -235,10 +235,10 @@ impl DataFusionPlanner { }) }) .collect::>>()?; - return self.build_project_with_aggregates(input_plan, projections, &ordering); + return self.build_project_with_aggregates(ctx, input_plan, projections, &ordering); } let input_plan = self.build_operator(ctx, input)?; - self.build_project_with_aggregates(input_plan, projections, &[]) + self.build_project_with_aggregates(ctx, input_plan, projections, &[]) } /// The columns of a node or relationship variable (`b` -> `b__id`, `b__name`, ...), 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 3fdaa50..03f90fe 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/count_ops.rs @@ -135,7 +135,7 @@ fn conjuncts<'a>(expr: &'a BooleanExpression, out: &mut Vec<&'a BooleanExpressio impl DataFusionPlanner { /// The node id column of `variable` (`b` -> `b__id`). - fn id_column(&self, ctx: &PlanningContext, variable: &str) -> Option { + pub(crate) fn id_column(&self, ctx: &PlanningContext, variable: &str) -> Option { let label = ctx.analysis.var_to_label.get(variable)?; let mapping = self.config.get_node_mapping(label)?; Some(qualify_column(variable, &mapping.id_field)) diff --git a/docs/csr.md b/docs/csr.md index e24d8be..b824e62 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -212,6 +212,28 @@ projection, its filtered values are reused within that output batch. Target keys neighbour IDs, including when a filter reads them. This reuse ends with the batch and keeps no property 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 +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: + +```python +plan = engine.explain( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, count(*) AS n RETURN b.name, n" +) +print([line.strip(" |") for line in plan.splitlines() if "NodeLookupExec" in line]) +# ['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 +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. + ### Patterns start at the most selective node The planner rewrites each `MATCH` so traversal starts at its most selective node, walking edges backwards diff --git a/python/python/tests/test_engine_with.py b/python/python/tests/test_engine_with.py index 46fcfc4..cb93ab7 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 lance +import pyarrow as pa import pytest from lance_csr import CypherEngine from test_engine_csr import CONFIG, KNOWS, PEOPLE @@ -93,3 +95,67 @@ def test_camel_case_aggregate_alias_can_be_ordered_by(): {"id": 2, "edgecount": 2}, {"id": 1, "edgecount": 1}, ] + + +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. + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, count(*) AS n " + "RETURN b.id AS id, b.name AS name, b.age AS age, n" + ) + 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) + 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("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) + got = engine.execute( + "MATCH (a:Person) OPTIONAL MATCH (a)-[:KNOWS]->(b:Person) " + "WITH b, count(*) AS n RETURN b.id AS id, b.name AS name, n" + ).to_pylist() + rows = sorted((r["id"] or 0, r["name"], r["n"]) for r in got) + # Persons 4 and 6 match nothing: one null b for both, with null columns. + assert rows == [(0, None, 2), (1, "a", 1), (2, "b", 2), (3, "c", 3)] + + +@pytest.mark.parametrize("lazy", [False, True]) +def test_node_lookup_over_many_groups(tmp_path, lazy): + # More groups than one batch. From Lance, with deleted rows, the plain grouping + # must still match the join path. + n = 20000 + people = pa.table( + { + "id": list(range(n)), + "name": [None if i % 11 == 0 else f"p{i}" for i in range(n)], + "age": [i % 90 for i in range(n)], + } + ) + knows = pa.table( + { + "src": pa.array([i % n for i in range(2 * n)], pa.int64()), + "dst": pa.array([(i * 7 + 3) % n for i in range(2 * n)], pa.int64()), + } + ) + if lazy: + lance.write_dataset(people, tmp_path / "p.lance", max_rows_per_file=3000) + lance.dataset(tmp_path / "p.lance").delete("id % 13 = 0") + people = lance.dataset(tmp_path / "p.lance") + datasets = {"Person": people, "KNOWS": knows} + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, count(*) AS n " + "RETURN b.id AS id, b.name AS name, b.age AS age, n" + ) + csr = CypherEngine(CONFIG, datasets, use_csr=True) + # Only in-memory node tables are looked up; Lance ones keep the plain grouping. + assert ("NodeLookupExec" in csr.explain(query)) is not lazy + by_id = lambda t: sorted(t.to_pylist(), key=lambda r: r["id"]) # noqa: E731 + assert by_id(csr.execute(query)) == by_id( + CypherEngine(CONFIG, datasets).execute(query) + )