perf(planner): look up grouped node columns from lance tables

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<Option<Arc<TargetColumns>>> {
+        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<String>,
 ) -> datafusion::common::Result<Transformed<LogicalPlan>> {
+    // 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::<NodeLookupNode>() {
+                let referenced: HashSet<String> = 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::<NodeLookupNode>() {
             // 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::<CsrExpandNode>().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::<CsrExpandNode>()
+                .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::<CsrExpandNode>() {
+                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<HashSet<String>> {
     let mut used: HashSet<String> = plan
         .schema()

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<dyn ExecutionPlan>,
@@ -241,6 +241,8 @@ pub struct NodeLookupExec {
     columns: Vec<usize>,
     schema: SchemaRef,
     properties: Arc<PlanProperties>,
+    /// Shared by the partitions: rows looked up so far, and the columns once read whole.
+    scan_state: Arc<TargetScanState>,
 }
 
 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<RecordBatch> {
         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<TaskContext>,
     ) -> datafusion::common::Result<SendableRecordBatchStream> {
+        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::<i32>();

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<Item = &'a str>,
     ) -> Option<LateNode> {
         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<LateNode> {
         let mut grouped: Vec<(String, Vec<String>)> = 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;

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