From e127755c872949733c69435047384c9d17fb572b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:59:30 -0700 Subject: [PATCH 1/3] fix(csr): stream expansion in bounded chunks Resolve #11 by resuming neighbor traversal between output batches and storing immutable deletion occurrence offsets. Add focused traversal/snapshot tests, a reusable hub memory scenario, and documentation. Assisted by Codex. --- crates/lance-graph/src/csr_expand.rs | 439 +++++++++++++++++-------- docs/csr.md | 8 + docs/cypher.md | 2 + python/benchmarks/bounded_memory.py | 42 ++- python/python/tests/test_engine_csr.py | 35 ++ 5 files changed, 377 insertions(+), 149 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index b8b88f9..201a3cc 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -77,9 +77,9 @@ use crate::error::{GraphError, Result}; /// /// An immutable CSR base, shared between snapshots via `Arc`, plus one delta applied /// by [`Adjacency::with_changes`]: appended edges per source key, and tombstones as -/// multiset counts of `(source, neighbour)` pairs. Counting pairs instead of -/// recording edge identities keeps the base at 8 bytes per edge, and is exact -/// because identical edges are indistinguishable in query results. +/// sorted offsets into each source's base-plus-appended neighbours. Offsets let +/// readers skip deleted occurrences without copying per-source deletion state. +/// Identical edges remain distinct occurrences; the base stays at 8 bytes per edge. /// /// # Design note for future changes (read before "optimizing" the delta) /// @@ -89,6 +89,8 @@ use crate::error::{GraphError, Result}; /// in-flight queries keep using the previous one, so the delta cannot be mutated in /// place. The copy is bounded: the engine rebuilds once the delta exceeds /// `CSR_REBUILD_FRACTION` (25%) of the base. +/// Deletions additionally scan each affected source's neighbours once to resolve +/// occurrence offsets; queries can then resume without allocating deletion maps. /// /// Measured at 10M edges (`python/benchmarks/csr_delta_accumulated.py`, 20 × 1% /// appends): refresh grows from about 11 ms to about 80 ms as the delta reaches 20%, @@ -105,7 +107,7 @@ use crate::error::{GraphError, Result}; pub struct Adjacency { base: Arc, added: FastMap>, - removed: FastMap>, + removed: FastMap>, added_count: u64, removed_count: u64, } @@ -158,7 +160,7 @@ impl Adjacency { /// /// Callers must only remove edges that are present. `self` is unchanged, so /// snapshots still using it keep their answers. The base is never copied. - // shortcut: copies the accumulated delta on every call (O(delta), ~80 ms at 2M delta edges); bounded by the engine's 25% rebuild threshold. Upgrade to a persistent map if refreshes get frequent on much larger graphs. + // shortcut: copies the accumulated delta on refresh; bounded by the engine's 25% rebuild threshold. Upgrade to a persistent map if refreshes get frequent on much larger graphs. pub fn with_changes(&self, added: &[(i64, i64)], removed: &[(i64, i64)]) -> Self { let mut next = Self { base: self.base.clone(), @@ -170,39 +172,79 @@ impl Adjacency { for &(from, to) in added { next.added.entry(from).or_default().push(to as u64); } + let mut removals: FastMap> = FastMap::default(); for &(from, to) in removed { - *next - .removed + *removals .entry(from) .or_default() .entry(to as u64) .or_default() += 1; } + for (from, mut counts) in removals { + let base = next.base_neighbors(from); + let added = next.added.get(&from).map_or(&[][..], Vec::as_slice); + let prior = next.removed.get(&from).map_or(&[][..], Vec::as_slice); + let mut prior = prior.iter().copied().peekable(); + let mut positions = Vec::new(); + for (position, neighbor) in base.iter().chain(added).enumerate() { + if prior.peek() == Some(&position) { + prior.next(); + continue; + } + if let Some(count) = counts.get_mut(neighbor) { + positions.push(position); + *count -= 1; + if *count == 0 { + counts.remove(neighbor); + } + if counts.is_empty() { + break; + } + } + } + debug_assert!(counts.is_empty(), "removed edges must be present"); + let tombstones = next.removed.entry(from).or_default(); + tombstones.extend(positions); + tombstones.sort_unstable(); + } next } /// Call `f` for each target endpoint value reachable from `key`, duplicates included. - pub fn for_each_neighbor(&self, key: i64, mut f: impl FnMut(u64)) { - let base = self - .base + pub fn for_each_neighbor(&self, key: i64, f: impl FnMut(u64)) { + self.neighbors_from(key, &mut 0).for_each(f); + } + + fn base_neighbors(&self, key: i64) -> &[u64] { + self.base .dense .get(&key) - .map_or(&[][..], |&v| self.base.csr.neighbors(v)); - let added = self.added.get(&key).map_or(&[][..], |v| v.as_slice()); - let all = base.iter().chain(added).copied(); - match self.removed.get(&key) { - None => all.for_each(f), - Some(removed) => { - // Skip as many occurrences of each pair as were removed. - let mut remaining = removed.clone(); - for neighbor in all { - match remaining.get_mut(&neighbor) { - Some(n) if *n > 0 => *n -= 1, - _ => f(neighbor), - } + .map_or(&[], |&v| self.base.csr.neighbors(v)) + } + + /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + fn neighbors_from<'a>( + &'a self, + key: i64, + offset: &'a mut usize, + ) -> impl Iterator + 'a { + let base = self.base_neighbors(key); + let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); + let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); + let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; + base.iter() + .chain(added) + .enumerate() + .skip(*offset) + .filter_map(move |(position, &neighbor)| { + *offset = position + 1; + if removed.first() == Some(&position) { + removed = &removed[1..]; + None + } else { + Some(neighbor) } - } - } + }) } /// Live edges: base + appended - removed. @@ -729,6 +771,13 @@ pub struct Neighbors { reverse: Option>, } +/// Occurrence offsets are sufficient to pause even a high-degree source row. +#[derive(Default)] +struct NeighborPosition { + forward: usize, + reverse: usize, +} + impl Neighbors { pub fn directed(adjacency: Arc) -> Self { Self { @@ -744,15 +793,27 @@ impl Neighbors { } } - pub fn for_each_neighbor(&self, key: i64, mut f: impl FnMut(u64)) { - self.forward.for_each_neighbor(key, &mut f); - if let Some(reverse) = &self.reverse { - reverse.for_each_neighbor(key, |n| { - if n as i64 != key { - f(n) - } - }); - } + pub fn for_each_neighbor(&self, key: i64, f: impl FnMut(u64)) { + self.neighbors_from(key, &mut NeighborPosition::default()) + .for_each(f); + } + + fn neighbors_from<'a>( + &'a self, + key: i64, + position: &'a mut NeighborPosition, + ) -> impl Iterator + 'a { + let forward = self.forward.neighbors_from(key, &mut position.forward); + let reverse = self + .reverse + .as_ref() + .map(|r| r.neighbors_from(key, &mut position.reverse)); + forward.chain( + reverse + .into_iter() + .flatten() + .filter(move |&n| n as i64 != key), + ) } /// Edges visited over all keys (self-loops count once per direction in this estimate). @@ -1020,7 +1081,7 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { } } -/// Physical operator for [`CsrExpandNode`]; streams one output batch per input batch. +/// Physical operator for [`CsrExpandNode`]; expands input rows in bounded chunks. #[derive(Debug)] pub struct CsrExpandExec { input: Arc, @@ -1064,83 +1125,78 @@ impl CsrExpandExec { } } -/// Constraints on the expanded rows, applied before the output columns are taken. -#[derive(Clone, Copy)] -struct NeighborFilter<'a> { - /// Only rows whose target satisfies this (the moved target filter). - target_filter: Option<&'a CompiledTargetFilter>, - /// Only neighbours equal to this input column's value (a closing edge). - closing_index: Option, +/// One input batch and its current row/occurrence, independent of total fan-out. +struct BatchExpansion { + batch: RecordBatch, + keys: Int64Array, + bound: Option, + row: usize, + position: NeighborPosition, } -/// Expand one batch: repeat each input row once per neighbour and append both -/// endpoints, then the target node's columns when a lookup is fused in. -/// -/// Output comes in batches of at most `batch_size` rows, as DataFusion operators -/// expect: one input row can have thousands of neighbours. (Bounded batches are also -/// what a future work-splitting, morsel-style expansion would hand out.) -async fn expand_batch( - batch: RecordBatch, - key_index: usize, - adjacency: &Neighbors, - target: Option<&TargetLookup>, - neighbor_filter: NeighborFilter<'_>, - schema: &SchemaRef, - batch_size: usize, -) -> datafusion::common::Result> { - let keys = cast(batch.column(key_index), &DataType::Int64)?; - let keys = keys.as_primitive::(); - let bound = neighbor_filter - .closing_index - .map(|i| cast(batch.column(i), &DataType::Int64)) - .transpose()?; - let bound = bound.as_ref().map(|b| b.as_primitive::()); - let mut rows = Vec::new(); - let mut key_values = Vec::new(); - let mut neighbor_values = Vec::new(); - let mut target_positions = Vec::new(); - for i in 0..keys.len() { - if keys.is_null(i) || bound.is_some_and(|b| b.is_null(i)) { - continue; - } - let key = keys.value(i); - let required = bound.map(|b| b.value(i)); - adjacency.for_each_neighbor(key, |neighbor| { - let neighbor = neighbor as i64; - if required.is_some_and(|r| r != neighbor) { - return; - } - if let Some(target) = target { - // No target row: the inner join would drop this edge too. - let Some(position) = target.rows.position(neighbor) else { - return; - }; - target_positions.push(position); - } - rows.push(i as u32); - key_values.push(key); - neighbor_values.push(neighbor); - }); - } - let mut out = Vec::new(); - for start in (0..rows.len()).step_by(batch_size.max(1)) { - let end = (start + batch_size).min(rows.len()); - let chunk = ExpandedRows { - rows: rows[start..end].to_vec(), - key_values: key_values[start..end].to_vec(), - neighbor_values: neighbor_values[start..end].to_vec(), - target_positions: target_positions - .get(start..end) - .unwrap_or_default() - .to_vec(), +impl BatchExpansion { + fn try_new( + batch: RecordBatch, + key_index: usize, + closing_index: Option, + ) -> datafusion::common::Result { + let ints = |i| { + cast(batch.column(i), &DataType::Int64).map(|a| a.as_primitive::().clone()) }; - out.push(materialize(&batch, chunk, target, neighbor_filter, schema).await?); + let keys = ints(key_index)?; + let bound = closing_index.map(ints).transpose()?; + Ok(Self { + batch, + keys, + bound, + row: 0, + position: NeighborPosition::default(), + }) + } + + /// Collect only the next chunk; target properties are fetched after this returns. + fn next_chunk( + &mut self, + adjacency: &Neighbors, + target: Option<&TargetLookup>, + batch_size: usize, + ) -> ExpandedRows { + let mut chunk = ExpandedRows::default(); + while self.row < self.keys.len() { + let i = self.row; + if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { + let key = self.keys.value(i); + let required = self.bound.as_ref().map(|b| b.value(i)); + for neighbor in adjacency.neighbors_from(key, &mut self.position) { + let neighbor = neighbor as i64; + if required.is_some_and(|r| r != neighbor) { + continue; + } + if let Some(target) = target { + // No target row: the inner join would drop this edge too. + let Some(position) = target.rows.position(neighbor) else { + continue; + }; + chunk.target_positions.push(position); + } + chunk.rows.push(i as u32); + chunk.key_values.push(key); + chunk.neighbor_values.push(neighbor); + if chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + self.row += 1; + self.position = NeighborPosition::default(); + } + chunk } - Ok(out) } /// Expanded rows of one input batch: input row index, both endpoints, and the /// target row position when a lookup is fused in. +#[derive(Default)] struct ExpandedRows { rows: Vec, key_values: Vec, @@ -1153,7 +1209,7 @@ async fn materialize( batch: &RecordBatch, expanded: ExpandedRows, target: Option<&TargetLookup>, - neighbor_filter: NeighborFilter<'_>, + target_filter: Option<&CompiledTargetFilter>, schema: &SchemaRef, ) -> datafusion::common::Result { let ExpandedRows { @@ -1162,7 +1218,7 @@ async fn materialize( mut neighbor_values, mut target_positions, } = expanded; - if let (Some(target), Some(filter)) = (target, neighbor_filter.target_filter) { + if let (Some(target), Some(filter)) = (target, target_filter) { // Take only the filter's columns, evaluate, and keep passing rows before // taking anything else. Null counts as false, as in a WHERE clause. let taken = target @@ -1323,8 +1379,9 @@ impl ExecutionPlan for CsrExpandExec { let schema = self.schema.clone(); let out_schema = schema.clone(); let batch_size = context.session_config().batch_size(); - // `then` keeps batches in order and polls one take at a time per partition. - let stream = input.then(move |batch| { + // 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)| { let (adjacency, target, target_filter, schema) = ( adjacency.clone(), target.clone(), @@ -1332,29 +1389,31 @@ impl ExecutionPlan for CsrExpandExec { schema.clone(), ); async move { - let neighbor_filter = NeighborFilter { - target_filter: target_filter.as_deref(), - closing_index, - }; - expand_batch( - batch?, - key_index, - &adjacency, - target.as_deref(), - neighbor_filter, - &schema, - batch_size, - ) - .await + loop { + if pending.is_none() { + let Some(batch) = input.next().await else { + return Ok(None); + }; + pending = Some(BatchExpansion::try_new(batch?, key_index, closing_index)?); + } + let expansion = pending.as_mut().expect("input batch just set"); + let chunk = expansion.next_chunk(&adjacency, target.as_deref(), batch_size); + if chunk.rows.is_empty() { + pending = None; + continue; + } + let batch = materialize( + &expansion.batch, + chunk, + target.as_deref(), + target_filter.as_deref(), + &schema, + ) + .await?; + return Ok(Some((batch, (input, pending)))); + } } }); - let stream = stream.flat_map(|batches| { - let batches: Vec<_> = match batches { - Ok(batches) => batches.into_iter().map(Ok).collect(), - Err(e) => vec![Err(e)], - }; - futures::stream::iter(batches) - }); Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) } } @@ -1749,7 +1808,7 @@ mod tests { } #[tokio::test] - async fn expand_batch_repeats_rows_per_neighbor() { + async fn expansion_repeats_rows_per_neighbor_and_resumes_between_chunks() { let adj = Adjacency::build( edge_stream(vec![Some(1), Some(1)], vec![Some(2), Some(3)]), "src", @@ -1776,27 +1835,117 @@ mod tests { Field::new("r__src", DataType::Int32, true), Field::new("r__dst", DataType::Int32, true), ])); - let none = NeighborFilter { - target_filter: None, - closing_index: None, - }; let adj = Neighbors::directed(Arc::new(adj)); - let out = expand_batch(input.clone(), 0, &adj, None, none, &out_schema, 8192) + let mut expansion = BatchExpansion::try_new(input.clone(), 0, None).unwrap(); + let chunk = expansion.next_chunk(&adj, None, 8192); + let out = materialize(&input, chunk, None, None, &out_schema) .await .unwrap(); - assert_eq!(out.len(), 1); - assert_eq!(out[0].num_rows(), 2); - let dst = out[0] - .column(2) - .as_primitive::(); + assert_eq!(out.num_rows(), 2); + let dst = out.column(2).as_primitive::(); assert_eq!(dst.values().to_vec(), vec![2, 3]); + assert!(expansion.next_chunk(&adj, None, 8192).rows.is_empty()); // Output is split at `batch_size` rows. - let split = expand_batch(input, 0, &adj, None, none, &out_schema, 1) - .await - .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![2]); + assert_eq!(expansion.row, 0); + assert_eq!(expansion.position.forward, 1); + assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![3]); + assert!(expansion.next_chunk(&adj, None, 1).rows.is_empty()); + } + + #[tokio::test] + async fn high_degree_expansion_visits_only_the_first_chunk() { + let n = 100_000; + let base = Adjacency::build( + edge_stream(vec![Some(0); n], (1..=n as i64).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(); + let removed: Vec<_> = (1..=n as i64).step_by(2).map(|n| (0, n)).collect(); + let adj = Neighbors::directed(Arc::new(base.with_changes(&[], &removed))); + let input = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![0])) as ArrayRef)]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let chunk = expansion.next_chunk(&adj, None, 32); assert_eq!( - split.iter().map(|b| b.num_rows()).collect::>(), - vec![1, 1] + chunk.neighbor_values, + (2..=64).step_by(2).collect::>() + ); + assert_eq!(chunk.rows.len(), 32); + // The rest of the hub, including its tombstones, has not been visited. + assert_eq!(expansion.row, 0); + assert_eq!(expansion.position.forward, 64); + assert_eq!(expansion.next_chunk(&adj, None, 32).neighbor_values[0], 66); + assert_eq!(neighbors(&base, 0).len(), n); + } + + #[tokio::test] + async fn chunk_boundaries_preserve_closing_edges_and_missing_target_rows() { + let adjacency = Adjacency::build( + edge_stream(vec![Some(1); 4], vec![Some(2), Some(2), Some(3), Some(999)]), + "src", + "dst", + ) + .await + .unwrap(); + let adjacency = Neighbors::directed(Arc::new(adjacency)); + let input = RecordBatch::try_from_iter([ + ( + "key", + Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(1)])) as ArrayRef, + ), + ( + "bound", + Arc::new(Int64Array::from(vec![Some(2), None, Some(2), Some(3)])) as ArrayRef, + ), + ]) + .unwrap(); + let nodes = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![2])) as ArrayRef)]) + .unwrap(); + let target = TargetLookup::new( + "b".to_string(), + Arc::new(NodeRows::from_batch(nodes, "id").unwrap().unwrap()), + vec![Field::new("b__id", DataType::Int64, false)], ); + let mut expansion = BatchExpansion::try_new(input, 0, Some(1)).unwrap(); + for _ in 0..2 { + let chunk = expansion.next_chunk(&adjacency, Some(&target), 1); + assert_eq!(chunk.rows, vec![0]); + assert_eq!(chunk.neighbor_values, vec![2]); + assert_eq!(chunk.target_positions, vec![0]); + } + assert!(expansion + .next_chunk(&adjacency, Some(&target), 1) + .rows + .is_empty()); + } + + #[tokio::test] + async fn resumed_neighbors_keep_duplicate_deletions_and_self_loops_exact() { + let edges = || edge_stream(vec![Some(1); 4], vec![Some(1), Some(2), Some(2), Some(3)]); + let base = Adjacency::build(edges(), "src", "dst").await.unwrap(); + let next = base.with_changes(&[(1, 2), (1, 4)], &[(1, 2), (1, 3)]); + let later = next.with_changes(&[(1, 3)], &[(1, 2), (1, 2)]); + assert_eq!(neighbors(&later, 1), vec![1, 3, 4]); + assert_eq!(neighbors(&next, 1), vec![1, 2, 2, 4]); + assert_eq!(neighbors(&base, 1), vec![1, 2, 2, 3]); + + let reverse = Adjacency::build(edges(), "dst", "src").await.unwrap(); + let both = Neighbors::undirected(Arc::new(base), Arc::new(reverse)); + let mut position = NeighborPosition::default(); + let mut got = Vec::new(); + while let Some(n) = both.neighbors_from(1, &mut position).next() { + got.push(n); + } + assert_eq!(got, vec![1, 2, 2, 3]); + let mut position = NeighborPosition::default(); + assert_eq!(both.neighbors_from(2, &mut position).next(), Some(1)); + assert_eq!(both.neighbors_from(2, &mut position).next(), Some(1)); + assert_eq!(both.neighbors_from(2, &mut position).next(), None); } } diff --git a/docs/csr.md b/docs/csr.md index c34d0e7..f819ec6 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -91,6 +91,14 @@ Node index Expect about **70 bytes per edge** of extra peak memory, for both directions together. The index lives in memory even when the tables are read lazily from Lance. + During query execution, each CSR expansion partition keeps one input batch and collects at most one + output chunk of neighbour positions before fetching properties. It yields that chunk before visiting + the remaining neighbours, even for a single high-degree node. Deleted edge positions are read from the + shared snapshot rather than copied for each query; refresh resolves those positions by scanning the + affected sources' neighbour lists. The chunk row limit comes from DataFusion's batch + size; wide properties and multiple active partitions still increase working memory. This does not bound + the resident graph, the input data, or buffering in joins, sorts and aggregations elsewhere in the plan. + !!! warning "Experimental" The CSR is marked experimental. Results are identical with it on or off (the test suite runs each diff --git a/docs/cypher.md b/docs/cypher.md index b7e4674..d292977 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -105,6 +105,8 @@ print(rows) # 4 - Batches are computed as you read them. Scans, filters and projections stream with bounded memory; joins, sorts and aggregations may still buffer their inputs. +- CSR expansion visits neighbours and fetches their properties a chunk at a time, including when one node + has many neighbours. See [CSR memory cost](csr.md#turning-it-on) for what it keeps in memory. - Planning errors (a syntax error, an unknown label) raise immediately. Execution errors raise while reading. - The reader keeps its snapshot alive, so a concurrent [`refresh()`](lance.md#refresh) doesn't affect it. - Dropping the reader cancels the remaining work. diff --git a/python/benchmarks/bounded_memory.py b/python/benchmarks/bounded_memory.py index a2b5768..4e54971 100644 --- a/python/benchmarks/bounded_memory.py +++ b/python/benchmarks/bounded_memory.py @@ -3,7 +3,7 @@ """Peak memory of streamed Cypher scans over Lance data larger than RAM limits. -``generate`` writes a deterministic Person table with a wide unused ``payload`` +``generate`` writes a deterministic Person table with a wide ``payload`` column, in batches so the generator itself stays small. ``run`` executes a query in a fresh subprocess, consumes ``execute_batches()`` incrementally (count + checksum, nothing retained), and reports peak RSS from ``getrusage``, which includes native @@ -11,10 +11,19 @@ python benchmarks/bounded_memory.py generate --out /data/large --rows 2000000 python benchmarks/bounded_memory.py run --data /data/large --query scan + +For CSR fan-out, generate a hub and return its wide properties. Compare fresh +processes on the base and changed revisions for several row counts. The pre-query +RSS includes the resident CSR and node index; peak growth is a lower bound because +getrusage also includes allocations made during engine construction. + + python benchmarks/bounded_memory.py generate --out /data/hub --rows 100000 --hub + python benchmarks/bounded_memory.py run --data /data/hub --query hub """ import argparse import json +import os import resource import subprocess import sys @@ -27,10 +36,13 @@ "scan": "MATCH (p:Person) RETURN p.id", # Selective graph join; the KNOWS side is filtered before the join. "join": "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id < 1000 RETURN b.id", + "hub": ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = 0 RETURN b.id, b.payload" + ), } -def generate(out: Path, rows: int, batch_rows: int = 16_384) -> None: +def generate(out: Path, rows: int, batch_rows: int = 16_384, hub: bool = False) -> None: import lance import numpy as np import pyarrow as pa @@ -56,7 +68,11 @@ def person_batches(): reader = pa.RecordBatchReader.from_batches(schema, person_batches()) lance.write_dataset(reader, out / "person.lance", mode="overwrite") ids = np.arange(rows, dtype=np.int64) - knows = pa.table({"src": ids, "dst": (ids + 1) % rows}) + knows = pa.table( + {"src": np.zeros(rows, dtype=np.int64), "dst": ids} + if hub + else {"src": ids, "dst": (ids + 1) % rows} + ) lance.write_dataset(knows, out / "knows.lance", mode="overwrite") logical = rows * (8 + PAYLOAD_BYTES) print(json.dumps({"rows": rows, "logical_input_bytes": logical})) @@ -80,7 +96,14 @@ def measure(data: Path, query: str) -> dict: "Person": lance.dataset(data / "person.lance"), "KNOWS": lance.dataset(data / "knows.lance"), }, + use_csr=query == "hub", + ) + peak_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + rss_before = ( + int(subprocess.check_output(["ps", "-o", "rss=", "-p", str(os.getpid())])) + * 1024 ) + query_start = time.perf_counter() reader = engine.execute_batches(QUERIES[query]) rows = checksum = batches = 0 first_batch_s = None @@ -94,7 +117,11 @@ def measure(data: Path, query: str) -> dict: peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # Linux reports KiB, macOS reports bytes. peak_bytes = peak if sys.platform == "darwin" else peak * 1024 + peak_before_bytes = peak_before if sys.platform == "darwin" else peak_before * 1024 input_rows = lance.dataset(data / "person.lance").count_rows() + if query in {"scan", "hub"}: + assert rows == input_rows + assert checksum == input_rows * (input_rows - 1) // 2 return { "query": query, "input_rows": input_rows, @@ -102,8 +129,14 @@ def measure(data: Path, query: str) -> dict: "checksum": checksum, "batches": batches, "time_to_first_batch_s": first_batch_s, + "query_time_to_first_batch_s": first_batch_s - (query_start - start) + if first_batch_s is not None + else None, "elapsed_s": elapsed, "peak_rss_bytes": peak_bytes, + "pre_query_rss_bytes": rss_before, + "pre_query_peak_rss_bytes": peak_before_bytes, + "peak_growth_lower_bound_bytes": max(0, peak_bytes - peak_before_bytes), "plan_scan": next( ( line.strip() @@ -121,6 +154,7 @@ def main() -> None: gen = sub.add_parser("generate") gen.add_argument("--out", type=Path, required=True) gen.add_argument("--rows", type=int, required=True) + gen.add_argument("--hub", action="store_true") run = sub.add_parser("run") run.add_argument("--data", type=Path, required=True) run.add_argument("--query", choices=QUERIES, default="scan") @@ -131,7 +165,7 @@ def main() -> None: if args.cmd == "generate": args.out.mkdir(parents=True, exist_ok=True) - generate(args.out, args.rows) + generate(args.out, args.rows, hub=args.hub) elif args.cmd == "child": print(json.dumps(measure(args.data, args.query))) else: diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index 7741a55..b5a4dd7 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -7,6 +7,8 @@ sorted results must be identical. """ +from itertools import chain + import lance import pyarrow as pa import pytest @@ -197,6 +199,39 @@ def test_execute_batches_with_csr(datasets): assert rows(streamed) == rows(csr.execute(query)) +def test_high_degree_stream_keeps_its_snapshot_across_deletions(tmp_path): + n = 20_000 + content = "wide property " * 32 + people = pa.table({"id": range(n + 1), "content": [content] * (n + 1)}) + edges = pa.table({"src": [0] * n, "dst": [1, 1, *range(2, n)]}) + knows = lance.write_dataset(edges, tmp_path / "knows.lance") + person = lance.write_dataset(people, tmp_path / "person.lance") + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=True) + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = 0 RETURN b.id AS id, b.content AS content" + ) + reader = engine.execute_batches(query) + first = reader.read_next_batch() + assert 0 < first.num_rows < n + # A live reader retains the old adjacency, including both copies of 0->1. + knows.delete("dst = 1") + engine.refresh() + assert engine.csr_stats()["KNOWS"]["tombstones"] == 2 + counts = {} + for batch in chain([first], reader): + for row in batch.to_pylist(): + assert row["content"] == content + counts[row["id"]] = counts.get(row["id"], 0) + 1 + assert counts == {1: 2, **dict.fromkeys(range(2, n), 1)} + assert engine.execute(query).num_rows == n - 2 + + reader = engine.execute_batches(query) + reader.read_next_batch() + reader.close() + assert engine.execute(query).num_rows == n - 2 + + def physical_plan(engine, query): plan = engine.explain(query) return plan[plan.index("physical_plan") :] From e8ca5df46d4e13ec90581bc0cb905059373fe44b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 21:27:32 -0700 Subject: [PATCH 2/3] perf(csr): collect bounded chunks with local slice cursors --- crates/lance-graph/src/csr_expand.rs | 59 ++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 13 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 201a3cc..9954a1e 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -232,12 +232,14 @@ impl Adjacency { let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; - base.iter() - .chain(added) - .enumerate() - .skip(*offset) - .filter_map(move |(position, &neighbor)| { - *offset = position + 1; + let added_offset = offset.saturating_sub(base.len()); + base.get(*offset..) + .unwrap_or_default() + .iter() + .chain(added.get(added_offset..).unwrap_or_default()) + .filter_map(move |&neighbor| { + let position = *offset; + *offset += 1; if removed.first() == Some(&position) { removed = &removed[1..]; None @@ -772,7 +774,7 @@ pub struct Neighbors { } /// Occurrence offsets are sufficient to pause even a high-degree source row. -#[derive(Default)] +#[derive(Clone, Copy, Default)] struct NeighborPosition { forward: usize, reverse: usize, @@ -1162,12 +1164,14 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); - while self.row < self.keys.len() { - let i = self.row; + let mut position = self.position; + let mut row = self.row; + 'rows: while row < self.keys.len() { + let i = row; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); let required = self.bound.as_ref().map(|b| b.value(i)); - for neighbor in adjacency.neighbors_from(key, &mut self.position) { + for neighbor in adjacency.neighbors_from(key, &mut position) { let neighbor = neighbor as i64; if required.is_some_and(|r| r != neighbor) { continue; @@ -1183,13 +1187,15 @@ impl BatchExpansion { chunk.key_values.push(key); chunk.neighbor_values.push(neighbor); if chunk.rows.len() == batch_size.max(1) { - return chunk; + break 'rows; } } } - self.row += 1; - self.position = NeighborPosition::default(); + row += 1; + position = NeighborPosition::default(); } + self.row = row; + self.position = position; chunk } } @@ -1752,6 +1758,33 @@ mod tests { assert!(Arc::ptr_eq(&base.base, &later.base)); } + #[tokio::test] + async fn resumed_neighbors_cross_the_base_and_delta_boundary() { + let base = Adjacency::build( + edge_stream(vec![Some(1), Some(1)], vec![Some(10), Some(11)]), + "src", + "dst", + ) + .await + .unwrap(); + let next = base.with_changes(&[(1, 12), (1, 12), (1, 13)], &[(1, 11), (1, 12)]); + for (mut offset, expected) in [ + (0, vec![10, 12, 13]), + (1, vec![12, 13]), + (2, vec![12, 13]), + (3, vec![12, 13]), + (4, vec![13]), + (5, vec![]), + (6, vec![]), + ] { + assert_eq!( + next.neighbors_from(1, &mut offset).collect::>(), + expected + ); + } + assert_eq!(neighbors(&base, 1), vec![10, 11]); + } + #[tokio::test] async fn debug_output_does_not_grow_with_edges() { let n = 10_000; From d7ad39819d04acf44907213ba7532df5e93d8869 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 21:54:51 -0700 Subject: [PATCH 3/3] perf(csr): fuse borrowed cursor construction with its consumer --- crates/lance-graph/src/csr_expand.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 9954a1e..41d811a 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -215,6 +215,7 @@ impl Adjacency { self.neighbors_from(key, &mut 0).for_each(f); } + #[inline(always)] fn base_neighbors(&self, key: i64) -> &[u64] { self.base .dense @@ -223,6 +224,7 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, @@ -800,6 +802,7 @@ impl Neighbors { .for_each(f); } + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64,