diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 62aa08e..0add730 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -909,6 +909,41 @@ impl fmt::Debug for TargetScanState { // shortcut: fixed share and no memory budget (the engine sets no DataFusion memory pool); add a reservation and fall back to per-chunk fetches if one is configured. const TARGET_SCAN_FRACTION: u64 = 8; +/// Fetched rows of `rows` after which [`TargetScanState::read_once`] reads the columns +/// whole. Fetching at most one batch makes at most one take: never worth a scan. +fn read_once_after(rows: &NodeRows, batch_size: usize) -> u64 { + (rows.num_rows() as u64 / TARGET_SCAN_FRACTION).max(batch_size as u64 + 1) +} + +impl TargetScanState { + /// Count `fetched` more rows of `rows`, read `columns` of the whole table once the + /// count reaches `scan_after`, and return that read when there is one. A batch that + /// crosses the threshold already uses the read. `CsrExpandExec` inlines the same + /// steps: calling this per chunk measured about 2% slower on a hop over 1.6M edges. + async fn read_once( + &self, + rows: &NodeRows, + columns: &[usize], + fetched: u64, + scan_after: u64, + ) -> datafusion::common::Result>> { + use std::sync::atomic::Ordering::Relaxed; + if self.columns.get().is_none() + && self.fetched.fetch_add(fetched, Relaxed) + fetched >= scan_after + { + let _building = self.building.lock().await; + if self.columns.get().is_none() { + let columns = rows + .scan_columns(columns) + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?; + let _ = self.columns.set(columns.map(Arc::new)); + } + } + Ok(self.columns.get().cloned().flatten()) + } +} + /// Fetch a target's columns from the read-once columns when present, else by row id. async fn take_target( target: &TargetLookup, @@ -2531,10 +2566,10 @@ impl ExecutionPlan for CsrExpandExec { .unwrap_or_default(), ); let batch_size = context.session_config().batch_size(); - // A hop that fetches at most one batch makes at most one take: never worth a scan. - let scan_after = self.target.as_ref().map_or(1, |t| { - (t.rows.num_rows() as u64 / TARGET_SCAN_FRACTION).max(batch_size as u64 + 1) - }); + let scan_after = self + .target + .as_ref() + .map_or(1, |t| read_once_after(&t.rows, batch_size)); let schema = self.schema.clone(); let out_schema = schema.clone(); let traversal = self.traversal.clone(); @@ -3028,19 +3063,44 @@ fn prune_within( plan: LogicalPlan, used: &HashSet, ) -> datafusion::common::Result> { + // A projection outputs only its expressions, so lookup columns right under it that + // they don't reference are unused, whatever else shares their names (the same + // node's columns from another lookup). + if let LogicalPlan::Projection(projection) = &plan { + if let LogicalPlan::Extension(ext) = projection.input.as_ref() { + if let Some(lookup) = ext.node.as_any().downcast_ref::() { + let referenced: HashSet = projection + .expr + .iter() + .flat_map(|e| e.column_refs()) + .map(|c| c.name.clone()) + .collect(); + if let Some(pruned) = lookup.pruned(&referenced) { + let input = match pruned { + Ok(node) => LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + Err(input) => input, + }; + let plan = + LogicalPlan::Projection(datafusion::logical_expr::Projection::try_new( + projection.expr.clone(), + Arc::new(input), + )?); + return Ok(Transformed::yes(prune_within(plan, used)?.data)); + } + } + } + } if let LogicalPlan::Extension(ext) = &plan { if let Some(lookup) = ext.node.as_any().downcast_ref::() { // The columns a lookup fetches shadow any of the same name below it (a CSR // hop's `take` of `b__name` under `NodeLookup: take=b[b__name]`), so its input // needs only what the input itself outputs. Find that set only when some hop // below fetches one of them, as it takes a walk of the input. - let shadows = lookup.input().exists(|node| { - Ok(matches!(node, LogicalPlan::Extension(ext) - if ext.node.as_any().downcast_ref::().is_some_and(|csr| { - csr.added_columns() - .any(|c| lookup.fields().iter().any(|f| f.name() == c)) - }))) - })?; + let shadows = hop_fetches(lookup.input(), |c| { + lookup.fields().iter().any(|f| f.name() == c) + }); let input = if shadows { prune_unused_columns(lookup.input().clone())? } else { @@ -3086,6 +3146,48 @@ fn prune_within( }) } +/// Whether a CSR hop in `plan` adds a column for which `wanted` holds (`b__name` for a +/// hop that fetches `b`'s name). +pub(crate) fn hop_fetches(plan: &LogicalPlan, wanted: impl Fn(&str) -> bool) -> bool { + plan.exists(|node| { + Ok(matches!(node, LogicalPlan::Extension(ext) + if ext.node.as_any().downcast_ref::() + .is_some_and(|csr| csr.added_columns().any(|c| wanted(c))))) + }) + .unwrap_or(false) +} + +/// Whether a CSR hop in `plan` fetches a column for which `wanted` holds and nothing +/// else in `plan` reads one, other than a projection passing it through. Pruning then +/// leaves the hop fetching none of them, once the columns are only needed above `plan`. +pub(crate) fn only_hop_fetches(plan: &LogicalPlan, wanted: impl Fn(&str) -> bool) -> bool { + let passes_through = |e: &Expr| match e { + Expr::Column(_) => true, + Expr::Alias(alias) => matches!(alias.expr.as_ref(), Expr::Column(_)), + _ => false, + }; + let mut fetched = false; + let read = plan.exists(|node| { + if let LogicalPlan::Extension(ext) = node { + if let Some(csr) = ext.node.as_any().downcast_ref::() { + fetched |= csr.added_columns().any(|c| wanted(c)); + return Ok(false); + } + } + // `DISTINCT` compares every column without naming them. + if let LogicalPlan::Distinct(_) = node { + return Ok(node.schema().fields().iter().any(|f| wanted(f.name()))); + } + let projection = matches!(node, LogicalPlan::Projection(_)); + Ok(node + .expressions() + .iter() + .filter(|e| !(projection && passes_through(e))) + .any(|e| e.column_refs().iter().any(|c| wanted(&c.name)))) + }); + fetched && read.is_ok_and(|read| !read) +} + fn used_columns(plan: &LogicalPlan) -> datafusion::common::Result> { let mut used: HashSet = plan .schema() diff --git a/crates/lance-csr/src/csr_expand/node_lookup.rs b/crates/lance-csr/src/csr_expand/node_lookup.rs index 217422b..26f2245 100644 --- a/crates/lance-csr/src/csr_expand/node_lookup.rs +++ b/crates/lance-csr/src/csr_expand/node_lookup.rs @@ -229,9 +229,9 @@ impl UserDefinedLogicalNodeCore for NodeLookupNode { } /// 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. +/// wanted columns. From a Lance table, once the looked-up rows pass the share a CSR +/// hop's target fetch switches at, the columns are read whole once and gathered from +/// memory, like the hop's. #[derive(Debug)] pub struct NodeLookupExec { input: Arc, @@ -241,6 +241,8 @@ pub struct NodeLookupExec { columns: Vec, schema: SchemaRef, properties: Arc, + /// Shared by the partitions: rows looked up so far, and the columns once read whole. + scan_state: Arc, } impl NodeLookupExec { @@ -260,6 +262,7 @@ impl NodeLookupExec { columns: node.columns.clone(), schema, properties, + scan_state: Arc::default(), } } @@ -270,6 +273,8 @@ impl NodeLookupExec { rows: &NodeRows, columns: &[usize], schema: &SchemaRef, + scan_state: &TargetScanState, + scan_after: u64, ) -> datafusion::common::Result { let external = |e: GraphError| datafusion::common::DataFusionError::External(Box::new(e)); let ids = int64_column(&batch, key_column).map_err(external)?; @@ -285,7 +290,14 @@ impl NodeLookupExec { Some(positions.len() as u32 - 1) }) .collect(); - let fetched = rows.take(&positions, columns).await.map_err(external)?; + let scanned = scan_state + .read_once(rows, columns, positions.len() as u64, scan_after) + .await?; + let fetched = match scanned { + Some(scanned) => scanned.take(&positions, columns), + None => rows.take(&positions, columns).await, + } + .map_err(external)?; let fetched = if positions.len() == ids.len() { fetched } else { @@ -353,6 +365,7 @@ impl ExecutionPlan for NodeLookupExec { columns: self.columns.clone(), schema: self.schema.clone(), properties, + scan_state: Arc::default(), })) } @@ -361,14 +374,27 @@ impl ExecutionPlan for NodeLookupExec { partition: usize, context: Arc, ) -> datafusion::common::Result { + let scan_after = read_once_after(&self.rows, context.session_config().batch_size()); 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 (key_column, scan_state) = (self.key_column.clone(), self.scan_state.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 } + let (key_column, scan_state) = (key_column.clone(), scan_state.clone()); + async move { + let batch = batch?; + Self::lookup( + batch, + &key_column, + &rows, + &columns, + &schema, + &scan_state, + scan_after, + ) + .await + } }); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema.clone(), @@ -412,7 +438,8 @@ mod tests { Field::new("b__id", DataType::Int64, true), Field::new("b__name", DataType::Utf8, true), ])); - let out = NodeLookupExec::lookup(input, "b__id", &rows, &[1], &schema) + let state = TargetScanState::default(); + let out = NodeLookupExec::lookup(input, "b__id", &rows, &[1], &schema, &state, u64::MAX) .await .unwrap(); let names = out.column(1).as_string::(); 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 25c0cbb..7e09bde 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use crate::ast::ValueExpression; use crate::case_insensitive::qualify_column; -use crate::csr_expand::{NodeLookupNode, NodeRows}; +use crate::csr_expand::{only_hop_fetches, NodeLookupNode, NodeRows}; use crate::datafusion_planner::analysis::PlanningContext; use crate::datafusion_planner::DataFusionPlanner; use crate::error::Result; @@ -112,28 +112,30 @@ fn count_values_without_ids( impl DataFusionPlanner { /// For `WITH b, ...` grouped by a node, or `RETURN b.id, b.name, ...` grouped by its - /// id and other properties: when `b`'s label has an in-memory node index (unique - /// integer ids), group by the id alone and look the other columns up afterwards, - /// once per node. Rows before the grouping can outnumber nodes many times over (an - /// expansion), and carrying each node's columns through the aggregate costs more - /// than the lookup. With unique ids, grouping by the id and by all the columns are - /// the same. A Lance node table keeps the plain grouping: the input rows already hold - /// the node's columns, read by an earlier step, so grouping by them reads nothing - /// more, while a lookup afterwards would read them from disk a second time. + /// id and other properties: when `b`'s label has a 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 is looked up only when a CSR hop in `input` fetches the + /// columns and nothing else in `input` reads them: the hop then fetches only the id, + /// and the lookup reads the columns once per node instead of once per expanded row. + /// Columns read by a scan or a join, or used before the grouping (a `WHERE` or + /// `DISTINCT` on them), keep the plain grouping, as a lookup afterwards would read + /// them a second time. /// /// `names` are the grouped columns of `var` (`b__id`, `b__name`); `None` unless they /// include the id and at least one other column of the node table. fn late_node<'a>( &self, ctx: &PlanningContext, + input: &LogicalPlan, var: &str, names: impl IntoIterator, ) -> Option { 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(); @@ -154,6 +156,9 @@ impl DataFusionPlanner { others.push((name.to_string(), field.name().clone())); } } + if !rows.in_memory() && !only_hop_fetches(input, |c| others.iter().any(|(n, _)| n == c)) { + return None; + } (found_id && !others.is_empty()).then(|| LateNode { variable: var.to_lowercase(), rows, @@ -167,6 +172,7 @@ impl DataFusionPlanner { fn late_properties( &self, ctx: &PlanningContext, + input: &LogicalPlan, projections: &[ProjectionItem], ) -> Vec { let mut grouped: Vec<(String, Vec)> = Vec::new(); @@ -187,7 +193,9 @@ impl DataFusionPlanner { grouped .into_iter() .filter(|(var, _)| !whole(var)) - .filter_map(|(var, names)| self.late_node(ctx, &var, names.iter().map(String::as_str))) + .filter_map(|(var, names)| { + self.late_node(ctx, input, &var, names.iter().map(String::as_str)) + }) .collect() } @@ -214,7 +222,7 @@ impl DataFusionPlanner { }) .collect(); let projections = bound.as_slice(); - let late_properties = self.late_properties(ctx, projections); + let late_properties = self.late_properties(ctx, &input_plan, projections); let mut late = Vec::new(); for p in projections { // Grouping by a node groups by all its columns (`WITH b, count(*) AS n`). @@ -229,7 +237,9 @@ impl DataFusionPlanner { _ => None, }) .collect(); - if let Some(node) = names.and_then(|names| self.late_node(ctx, var, names)) { + if let Some(node) = + names.and_then(|names| self.late_node(ctx, &input_plan, var, names)) + { group_exprs.push(col(Column::new_unqualified(&node.id))); late.push(node); continue; diff --git a/docs/csr.md b/docs/csr.md index 53bbace..a06c95d 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -225,7 +225,7 @@ cache across queries or dataset versions. ### Grouping by a node `WITH b, count(*) AS n` groups by the node `b`, and `RETURN b.id, b.name, count(*) AS n` by its id and name. -When `b`'s label has a node index and its table is in memory, either grouping uses `b`'s id alone and looks up +When `b`'s label has a node index, either grouping can use `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 @@ -240,10 +240,15 @@ print([line.strip(" |") for line in plan.splitlines() if "NodeLookupExec" in lin ``` The grouping needs the id: `RETURN b.name, count(*)` groups by the name, because two people can share one. -It applies only when the node table is held in memory. For a Lance dataset read lazily, grouping by a node -still groups by all of its properties, as before. By the time the query groups by `b`, its rows already hold -`b`'s properties, read from the dataset by an earlier step, so grouping by them reads nothing more. Looking -them up after the grouping would read them from disk a second time, which costs more than keeping them. +An in-memory node table always uses the lookup. A Lance dataset read lazily uses it only when a hop fetched +`b`'s properties for the rows being grouped, as in the example above, and nothing before the grouping reads +them. That hop then reads only `b`'s id, and the lookup reads each property once per node instead of once per +edge. When `b`'s properties come from somewhere else, such as the scan that matched `b` in the first place, or +when a `WHERE` or `WITH DISTINCT` before the grouping needs them, the rows already hold them by the time they +reach the grouping, so grouping by them reads nothing more, while a lookup afterwards would read them from +disk a second time. Those groupings keep all of `b`'s properties, as before. A lookup that covers a large +share of a Lance table reads the properties in one sequential pass rather than row by row, the same switch a +[fused target lookup](#the-fused-target-lookup) makes. ### Patterns start at the most selective node diff --git a/docs/internals/operators.md b/docs/internals/operators.md index 8c255fd..d260280 100644 --- a/docs/internals/operators.md +++ b/docs/internals/operators.md @@ -154,9 +154,13 @@ ProjectionExec: expr=[b__name@2 as b.name, __agg_0@1 as n] The expansion below takes only `b`'s id: the lookup's `b__name` is a column of its own, so [`PruneCsrTargetColumns`](optimizer.md#prunecsrtargetcolumns) doesn't count it as a use of the expansion's. Grouping by `b.name` without `b.id` keeps the plain grouping, since two nodes may share a name. Unlike a join back to the node -table, the lookup's cost follows the number of groups, not the size of the table. It applies to -in-memory node tables only; see [grouping by a node](../csr.md#grouping-by-a-node) for why lazily read Lance tables -keep the old plan. A [list comprehension](optimizer.md#list-comprehensions-keep-each-row-once) over a list of nodes uses +table, the lookup's cost follows the number of groups, not the size of the table. A lazily read Lance node table +gets the lookup only when a `CsrExpand` in the aggregate's input fetches the grouped columns and no other node in that +input reads them (a projection passing them through doesn't count; a `Filter` or `Distinct` does). Pruning then drops +them from the expansion, so the lookup replaces that per-edge fetch instead of reading the columns a second time; see +[grouping by a node](../csr.md#grouping-by-a-node). From Lance, once the looked-up rows pass 1/8 of the table, the +lookup reads its columns in one sequential pass and gathers from memory, with the same threshold and shared state as +`CsrExpand`'s target columns. 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 diff --git a/docs/internals/optimizer.md b/docs/internals/optimizer.md index 3e11c26..7ea39b8 100644 --- a/docs/internals/optimizer.md +++ b/docs/internals/optimizer.md @@ -34,6 +34,10 @@ narrows each `CsrExpand`'s target and relationship columns, and each `NodeLookup columns are all unused is removed. The columns a `NodeLookup` fetches shadow any of the same name below it, so the rule prunes the plan under each lookup again, against only what that plan outputs: a hop under `NodeLookup: take=b[b__name]` doesn't fetch `b__name` for every edge just because the lookup's `b__name` is used above. +A `NodeLookup` directly under a projection is also narrowed to the columns that projection references, whatever else +in the plan uses the same names. When a query groups by the same node twice, the lookup after the first grouping +would otherwise keep fetching `b__name`, used only after the second, even where the plan above it, such as one side of +an `OPTIONAL MATCH` join on `b`'s id, keeps nothing but the id. "Used" is conservative: expression references, the plan's output columns, each CSR node's key and closing column, and all inputs of nodes whose column use is implicit, such as `DISTINCT` and `UNION`. diff --git a/python/python/tests/test_engine_with.py b/python/python/tests/test_engine_with.py index 2f6b786..d44a606 100644 --- a/python/python/tests/test_engine_with.py +++ b/python/python/tests/test_engine_with.py @@ -180,10 +180,63 @@ def test_unmatched_optional_node_groups_with_null_columns(use_csr): assert rows == [(0, None, 2), (1, "a", 1), (2, "b", 2), (3, "c", 3)] +def test_lookup_projected_away_is_dropped(): + # The OPTIONAL MATCH joins on b's id alone, so the lookup under the join's left + # side fetches nothing anyone reads, even though b.name is used elsewhere. + datasets = {"Person": PEOPLE, "KNOWS": KNOWS} + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, collect(a.id) AS xs " + "OPTIONAL MATCH (b)-[:KNOWS]->(c:Person) WHERE c.id IN xs " + "WITH b, count(c) AS m RETURN b.name AS name, m" + ) + csr = CypherEngine(CONFIG, datasets, use_csr=True) + assert csr.explain(query).count("NodeLookupExec") == 2 + got = sorted((r["name"], r["m"]) for r in csr.execute(query).to_pylist()) + assert got == [("a", 0), ("b", 0), ("c", 1)] + + @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. +@pytest.mark.parametrize( + ("query", "lance_lookup"), + [ + # More groups than one batch: from Lance, the columns are read whole once. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH b, collect(a.age) AS ages " + "RETURN b.id AS id, b.name AS name, size(ages) AS n", + True, + ), + # A few groups: from Lance, looked up by row id. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id < 100 " + "WITH b, collect(a.age) AS ages " + "RETURN b.id AS id, b.name AS name, size(ages) AS n", + True, + ), + # b's columns come from a scan, not a hop: from Lance, a lookup would read + # them twice, so it keeps the plain grouping. + ( + "MATCH (b:Person)-[:KNOWS]->(a:Person) WITH b, collect(a.age) AS ages " + "RETURN b.id AS id, b.name AS name, size(ages) AS n", + False, + ), + # The hop's b columns are also read before the grouping, by a WHERE or a + # DISTINCT, so the hop keeps fetching them: from Lance, no second read. + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE b.age > 10 " + "WITH b, collect(a.age) AS ages " + "RETURN b.id AS id, b.name AS name, size(ages) AS n", + False, + ), + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WITH DISTINCT a, b " + "WITH b, collect(a.age) AS ages " + "RETURN b.id AS id, b.name AS name, size(ages) AS n", + False, + ), + ], +) +def test_node_lookup_over_many_groups(tmp_path, lazy, query, lance_lookup): + # From Lance, with deleted rows, the lookup must still match the join path. n = 20000 people = pa.table( { @@ -203,13 +256,9 @@ def test_node_lookup_over_many_groups(tmp_path, lazy): 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 + plan = csr.explain(query) + assert ("NodeLookupExec" in plan) is (lance_lookup or not lazy), plan 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)