From e4cdf726bb54d8c212f5e4050169dad90c50744d Mon Sep 17 00:00:00 2001 From: prrao87 Date: Fri, 30 Oct 2026 14:40:32 -0400 Subject: [PATCH 1/2] feat(csr): report how far the base CSR trails the data csr_stats() now reports base_version, the relationship dataset version of the last full build, and unindexed_edges, the live edges held outside the base (appended since that build and not since removed). delta_edges counts every append, including ones later deleted, so it can't answer that. Tombstones that remove appended edges are counted separately in Adjacency to make the live count exact. --- crates/lance-csr-python/src/graph.rs | 2 ++ crates/lance-csr/src/csr_expand.rs | 15 +++++++++++++++ crates/lance-csr/src/engine.rs | 11 +++++++++++ docs/csr.md | 10 +++++++--- python/python/lance_csr/_internal/graph.pyi | 2 ++ python/python/lance_csr/_types.py | 2 ++ python/python/lance_csr/cypher.py | 4 ++++ python/python/tests/test_engine_csr_delta.py | 7 ++++++- 8 files changed, 49 insertions(+), 4 deletions(-) diff --git a/crates/lance-csr-python/src/graph.rs b/crates/lance-csr-python/src/graph.rs index 70fe26b..d5b4a8a 100644 --- a/crates/lance-csr-python/src/graph.rs +++ b/crates/lance-csr-python/src/graph.rs @@ -1136,8 +1136,10 @@ impl CypherEngine { item.set_item("edges", entry.edges)?; item.set_item("delta_edges", entry.delta_edges)?; item.set_item("tombstones", entry.tombstones)?; + item.set_item("unindexed_edges", entry.unindexed_edges)?; item.set_item("update", entry.update)?; item.set_item("version", entry.version)?; + item.set_item("base_version", entry.base_version)?; item.set_item("build_ms", entry.build_ms)?; stats.set_item(rel_type, item)?; } diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 85b2cf0..500995e 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -129,6 +129,8 @@ pub struct Adjacency { removed: FastMap>, added_count: u64, removed_count: u64, + /// Tombstones that remove appended edges rather than base ones. + removed_added_count: u64, } struct CsrBase { @@ -181,6 +183,7 @@ impl Adjacency { removed: FastMap::default(), added_count: 0, removed_count: 0, + removed_added_count: 0, }) } @@ -199,6 +202,7 @@ impl Adjacency { removed: self.removed.clone(), added_count: self.added_count + added.len() as u64, removed_count: self.removed_count, + removed_added_count: self.removed_added_count, }; if !removed.is_empty() { // A row doesn't say which source holds it, so every live slot is checked. @@ -229,6 +233,8 @@ impl Adjacency { } if !positions.is_empty() { next.removed_count += positions.len() as u64; + next.removed_added_count += + positions.iter().filter(|&&p| p >= base_rows.len()).count() as u64; let tombstones = next.removed.entry(key).or_default(); tombstones.extend(positions); tombstones.sort_unstable(); @@ -328,6 +334,11 @@ impl Adjacency { pub fn tombstones(&self) -> u64 { self.removed_count } + + /// Live edges outside the base: appended and not since removed. + pub fn unindexed_edges(&self) -> u64 { + self.added_count - self.removed_added_count + } } // Summary only: DataFusion formats extension nodes with `{:?}` during physical @@ -2513,9 +2524,11 @@ mod tests { (next.num_edges(), next.delta_edges(), next.tombstones()), (4, 2, 1) ); + assert_eq!(next.unindexed_edges(), 2); // Removing an appended edge cancels it; the original is untouched. let later = next.with_changes(&[], &rows([4])); assert!(neighbors(&later, 9).is_empty()); + assert_eq!((later.tombstones(), later.unindexed_edges()), (2, 1)); assert_eq!(neighbors(&base, 1), vec![2, 2, 3]); assert!(Arc::ptr_eq(&base.base, &later.base)); } @@ -2532,6 +2545,8 @@ mod tests { let next = base .with_changes(&[(1, 12, 2), (1, 12, 3), (1, 13, 4)], &rows([])) .with_changes(&[], &rows([1, 2])); + // Row 1 is a base edge, row 2 an appended one. + assert_eq!((next.tombstones(), next.unindexed_edges()), (2, 2)); for (mut offset, expected) in [ (0, vec![10, 12, 13]), (1, vec![12, 13]), diff --git a/crates/lance-csr/src/engine.rs b/crates/lance-csr/src/engine.rs index 4fb8c20..c3c9ec8 100644 --- a/crates/lance-csr/src/engine.rs +++ b/crates/lance-csr/src/engine.rs @@ -168,11 +168,16 @@ pub struct CsrStats { pub delta_edges: u64, /// Removed edges held outside the base CSR since its last full build. pub tombstones: u64, + /// Live edges outside the base CSR: appended since its last full build and not + /// since removed. Lance's `num_unindexed_rows`, counted in edges. + pub unindexed_edges: u64, /// How the last refresh produced it: `build` (first), `rebuild` (full), `delta` /// (changes applied) or `reused` (no edge change). pub update: &'static str, /// Relationship dataset version it reflects (`None` for Arrow input). pub version: Option, + /// Relationship dataset version of the last full build (`None` for Arrow input). + pub base_version: Option, /// Wall time of that build or delta, both directions. pub build_ms: f64, } @@ -264,6 +269,8 @@ struct CsrEntry { incoming: Arc, /// Rows of `source`, for relationship columns the adjacencies don't hold. edges: Arc, + /// Version of the snapshot the base was fully built from. + base_version: Option, /// Wall time of the last build or delta application. build_ms: f64, /// How the last refresh produced this entry: build, rebuild, delta or reused. @@ -342,6 +349,7 @@ async fn try_apply_csr_delta( outgoing: Arc::new(outgoing), incoming: Arc::new(incoming), edges: source.edge_rows(), + base_version: old.base_version, build_ms: elapsed(), update: "delta", }) @@ -797,6 +805,7 @@ async fn build_csr_entries( outgoing: Arc::new(outgoing), incoming: Arc::new(incoming), edges: source.edge_rows(), + base_version: source.lance_version(), build_ms: started.elapsed().as_secs_f64() * 1e3, update: if previous_entry.is_some() { "rebuild" @@ -1057,8 +1066,10 @@ impl CypherEngine { edges: entry.outgoing.num_edges(), delta_edges: entry.outgoing.delta_edges(), tombstones: entry.outgoing.tombstones(), + unindexed_edges: entry.outgoing.unindexed_edges(), update: entry.update, version: entry.source.lance_version(), + base_version: entry.base_version, build_ms: entry.build_ms, }; (rel_type.clone(), stats) diff --git a/docs/csr.md b/docs/csr.md index 9dc6b6b..4784792 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -234,8 +234,9 @@ One entry per relationship type that has a CSR: ```python print(engine.csr_stats()["KNOWS"]) -# {'edges': 4, 'delta_edges': 0, 'tombstones': 0, 'update': 'build', -# 'version': None, 'build_ms': 0.41} (build_ms varies) +# {'edges': 4, 'delta_edges': 0, 'tombstones': 0, 'unindexed_edges': 0, +# 'update': 'build', 'version': None, 'base_version': None, +# 'build_ms': 0.41} (build_ms varies) ``` | Field | Meaning | @@ -243,8 +244,10 @@ print(engine.csr_stats()["KNOWS"]) | `edges` | Live indexed edges. Edges with a null endpoint are skipped. | | `delta_edges` | Appended edges held outside the base CSR since its last full build. | | `tombstones` | Removed edges masked out since the last full build. | +| `unindexed_edges` | Live edges outside the base CSR: appended since its last full build and not since removed. Lance index statistics call this `num_unindexed_rows`. | | `update` | How the last build or `refresh()` produced it: `build` (first), `rebuild` (full), `delta` (changes applied) or `reused` (no edge change). | | `version` | The relationship dataset's Lance version, or `None` for PyArrow input. | +| `base_version` | The version of the last full build, or `None` for PyArrow input. It trails `version` while deltas apply. | | `build_ms` | Wall time of that build or delta, both directions. | A relationship type missing from `csr_stats()` has no CSR (its endpoints aren't integers) and always joins. @@ -295,7 +298,8 @@ The engine doesn't watch your datasets for changes. Keeping the CSR in step with after all of them commit (see [snapshots are per dataset](lance.md#lazy-version-pinned-reads)). A refresh with nothing new reuses the existing CSR, so calling it on a timer or before each batch of queries is fine. 3. **Check what it did with [`csr_stats()`](#csr_stats)**: `update` says whether the CSR was `reused`, updated - with a `delta` or fully rebuilt (`rebuild`). + with a `delta` or fully rebuilt (`rebuild`). `base_version` and `unindexed_edges` say how far the base + CSR trails the data. Answers are exact either way: delta edges are part of every lookup. === "Python" diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index 5244151..2fe3f1d 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -33,8 +33,10 @@ class _CsrStats(TypedDict): edges: int delta_edges: int tombstones: int + unindexed_edges: int update: Literal["build", "rebuild", "delta", "reused"] version: int | None + base_version: int | None build_ms: float @final diff --git a/python/python/lance_csr/_types.py b/python/python/lance_csr/_types.py index ded83d9..34a9c5a 100644 --- a/python/python/lance_csr/_types.py +++ b/python/python/lance_csr/_types.py @@ -53,8 +53,10 @@ class CsrStats(TypedDict): edges: int delta_edges: int tombstones: int + unindexed_edges: int update: Literal["build", "rebuild", "delta", "reused"] version: int | None + base_version: int | None build_ms: float diff --git a/python/python/lance_csr/cypher.py b/python/python/lance_csr/cypher.py index 6a8c0bf..0d512fc 100644 --- a/python/python/lance_csr/cypher.py +++ b/python/python/lance_csr/cypher.py @@ -287,10 +287,14 @@ def csr_stats(self) -> dict[str, CsrStats]: - `edges`: live indexed edges (null endpoints are skipped). - `delta_edges`, `tombstones`: edges appended and removed since the last full build, held outside the base CSR. + - `unindexed_edges`: live edges outside the base CSR, appended since its + last full build and not since removed. - `update`: how the last refresh produced it: `"build"` (first), `"rebuild"` (full), `"delta"` (changes applied) or `"reused"` (no edge change). - `version`: the relationship dataset version it reflects (`None` for PyArrow). + - `base_version`: the relationship dataset version of the last full + build (`None` for PyArrow). It trails `version` while deltas apply. - `build_ms`: wall time of that build or delta, both directions. """ return self._inner.csr_stats() diff --git a/python/python/tests/test_engine_csr_delta.py b/python/python/tests/test_engine_csr_delta.py index cd5b54f..2962622 100644 --- a/python/python/tests/test_engine_csr_delta.py +++ b/python/python/tests/test_engine_csr_delta.py @@ -89,6 +89,7 @@ def test_initial_build_reports_build(graph): engine, _ = graph() s = stats(engine) assert (s["update"], s["delta_edges"], s["tombstones"]) == ("build", 0, 0) + assert s["unindexed_edges"] == 0 and s["base_version"] == s["version"] def test_append_applied_as_delta(graph): @@ -97,7 +98,8 @@ def test_append_applied_as_delta(graph): engine.refresh() s = stats(engine) assert (s["update"], s["delta_edges"], s["tombstones"]) == ("delta", 3, 0) - assert s["edges"] == 203 + assert s["edges"] == 203 and s["unindexed_edges"] == 3 + assert s["base_version"] < s["version"] check(engine, uri) @@ -159,6 +161,8 @@ def test_sequential_deltas_accumulate(graph): engine.refresh() s = stats(engine) assert s["update"] == "delta" and s["delta_edges"] == 3 and s["tombstones"] == 2 + # One tombstone removes an appended edge, so two of three appends are still live. + assert s["unindexed_edges"] == 2 check(engine, uri) @@ -217,6 +221,7 @@ def test_large_delta_triggers_rebuild(graph): engine.refresh() s = stats(engine) assert (s["update"], s["edges"], s["delta_edges"]) == ("rebuild", 300, 0) + assert s["unindexed_edges"] == 0 and s["base_version"] == s["version"] check(engine, uri) From ad96134c1aa7a28d9bb8466999753d0b4b0a4534 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Fri, 30 Oct 2026 14:48:35 -0400 Subject: [PATCH 2/2] perf(csr): read target columns once when a hop fetches much of a table The fused target lookup fetched the target's columns with take_rows for every chunk. When a hop expands a large share of a Lance node table, each chunk of scattered rows touches most of the column's pages, so the hop pays about one column read per chunk while a hash join reads the table once. Count the target rows an expansion fetches, shared across partitions as the relationship-filter mask does. Once they pass 1/8 of the node table, scan the needed columns (output plus target-filter columns) once with row ids and gather from memory for the rest of the query. Row ids are mapped explicitly, so fragments, deletions and stable row ids stay exact. Small hops and in-memory tables keep their per-chunk gathers. --- crates/lance-csr/src/csr_expand.rs | 273 +++++++++++++++++++++++-- python/python/tests/test_engine_csr.py | 30 +++ 2 files changed, 290 insertions(+), 13 deletions(-) diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 500995e..49a8f90 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -709,6 +709,138 @@ impl NodeRows { .take(&self.schema, positions, columns, false) .await } + + /// `columns` (ascending schema indices) of every row, from one sequential read. + /// `None` for an in-memory table, which already gathers without I/O. + async fn scan_columns(&self, columns: &[usize]) -> Result> { + let TableData::Lance(dataset) = &self.data else { + return Ok(None); + }; + let scan_error = |e: lance::Error| exec_error(format!("Failed to scan nodes: {}", e)); + let names: Vec<&str> = columns + .iter() + .map(|&c| self.schema.field(c).name().as_str()) + .collect(); + let mut scanner = dataset.scan(); + scanner.project(&names).map_err(scan_error)?; + scanner.with_row_id(); + let mut stream = scanner.try_into_stream().await.map_err(scan_error)?; + let mut index = FastMap::default(); + let mut parts: Vec> = vec![Vec::new(); columns.len()]; + while let Some(batch) = stream.next().await { + let batch = batch.map_err(scan_error)?; + let row_ids = batch + .column_by_name(lance::dataset::ROW_ID) + .ok_or_else(|| exec_error("Node scan returned no row ids".into()))? + .as_primitive::(); + // Row ids are not offsets (fragments, deletions, stable ids), so map them. + for &row_id in row_ids.values() { + let next = index.len() as u64; + index.insert(row_id, next); + } + for (part, name) in parts.iter_mut().zip(&names) { + part.push( + batch + .column_by_name(name) + .cloned() + .ok_or_else(|| exec_error("Node scan lost a column".into()))?, + ); + } + cooperative_yield().await; + } + let concat_error = |e: arrow::error::ArrowError| exec_error(e.to_string()); + let arrays = parts + .iter() + .zip(columns) + .map(|(part, &c)| { + let part: Vec<&dyn Array> = part.iter().map(|a| a.as_ref()).collect(); + if part.is_empty() { + Ok((c, new_empty_array(self.schema.field(c).data_type()))) + } else { + arrow::compute::concat(&part) + .map(|a| (c, a)) + .map_err(concat_error) + } + }) + .collect::>()?; + Ok(Some(TargetColumns { + index, + columns: arrays, + })) + } +} + +/// A target table's columns read once, sequentially, for a hop that fetches a large +/// share of it; later fetches gather from memory instead of reading pages again. +struct TargetColumns { + /// Lance row id -> position in `columns`' arrays. + index: FastMap, + /// Schema index -> the column's values for every row. + columns: Vec<(usize, ArrayRef)>, +} + +impl TargetColumns { + /// Same contract as [`NodeRows::take`], for columns this read covers. + fn take(&self, positions: &[u64], columns: &[usize]) -> Result> { + let indices = positions + .iter() + .map(|p| self.index.get(p).copied()) + .collect::>>() + .ok_or_else(|| exec_error("Node row missing from its own snapshot".into()))?; + let indices = UInt64Array::from(indices); + columns + .iter() + .map(|c| { + let (_, array) = self + .columns + .iter() + .find(|(column, _)| column == c) + .ok_or_else(|| exec_error("Target column was not read".into()))?; + take(array.as_ref(), &indices, None) + .map_err(|e| exec_error(format!("Failed to take rows: {}", e))) + }) + .collect() + } +} + +/// Shared by an operator's partitions: target rows fetched so far, and the target +/// columns once that count made one sequential read worthwhile. `None` once set +/// means the table is in memory and gathers directly. +#[derive(Default)] +struct TargetScanState { + fetched: std::sync::atomic::AtomicU64, + columns: std::sync::OnceLock>>, + building: futures::lock::Mutex<()>, +} + +impl fmt::Debug for TargetScanState { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("TargetScanState") + .field("scanned", &self.columns.get().is_some_and(Option::is_some)) + .finish() + } +} + +/// Switch from fetching target columns per chunk to one read of the needed columns +/// once the fetched target rows pass 1/`TARGET_SCAN_FRACTION` of the node table and +/// one batch (a scan's setup outweighs one small take), as +/// [`EDGE_MASK_FRACTION`] does for relationship filters. The read columns stay in +/// memory for the rest of the query. +// 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; + +/// Fetch a target's columns from the read-once columns when present, else by row id. +async fn take_target( + target: &TargetLookup, + scanned: Option<&TargetColumns>, + positions: &[u64], + columns: &[usize], +) -> datafusion::common::Result> { + match scanned { + Some(scanned) if !columns.is_empty() => scanned.take(positions, columns), + _ => target.rows.take(positions, columns).await, + } + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e))) } // Summary only, for the same reason as `Adjacency`. @@ -1577,6 +1709,8 @@ pub struct CsrExpandExec { /// The moved relationship filter, if any, and its shared progress. edge_filter: Option, edge_state: Arc, + /// Target rows fetched so far, and the target columns once read whole. + target_state: Arc, closing_column: Option, traversal: Option>, schema: SchemaRef, @@ -1612,6 +1746,7 @@ impl CsrExpandExec { target_filter, edge_filter, edge_state: Arc::default(), + target_state: Arc::default(), closing_column, traversal, schema, @@ -1823,10 +1958,12 @@ impl ExpandedRows { } /// Apply the moved target filter to `expanded`, then build the output batch. +#[allow(clippy::too_many_arguments)] async fn materialize( batch: &RecordBatch, expanded: ExpandedRows, target: Option<&TargetLookup>, + scanned: Option<&TargetColumns>, edge: Option<&EdgeLookup>, target_filter: Option<&CompiledFilter>, edge_filter: Option<&CompiledFilter>, @@ -1859,11 +1996,8 @@ async fn materialize( .copied() .filter(|&c| c != target.rows.id_index) .collect(); - let mut taken = target - .rows - .take(&target_positions, &fetch) - .await - .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))? + let mut taken = take_target(target, scanned, &target_positions, &fetch) + .await? .into_iter(); let arrays = filter .columns @@ -1923,11 +2057,8 @@ async fn materialize( .copied() .filter(|&c| c != id_index && !reused.iter().any(|(column, _)| *column == c)) .collect(); - let mut fetched = target - .rows - .take(&target_positions, &fetch) - .await - .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))? + let mut fetched = take_target(target, scanned, &target_positions, &fetch) + .await? .into_iter(); for (&c, field) in target.columns.iter().zip(&target.fields) { // The target id equals the neighbour value, so it is never fetched. @@ -2067,9 +2198,33 @@ impl ExecutionPlan for CsrExpandExec { let edge_filter = self.edge_filter.clone().map(Arc::new); let edge_state = self.edge_state.clone(); let mask_after = (self.adjacency.num_edges() / EDGE_MASK_FRACTION).max(1); + let target_state = self.target_state.clone(); + // Every non-id column the target fetches: its output and its filter's columns. + let scan_columns: Arc> = Arc::new( + self.target + .as_ref() + .map(|t| { + let filter = self.target_filter.iter().flat_map(|f| &f.columns); + let mut columns: Vec = t + .columns + .iter() + .chain(filter) + .copied() + .filter(|&c| c != t.rows.id_index) + .collect(); + columns.sort_unstable(); + columns.dedup(); + columns + }) + .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 schema = self.schema.clone(); let out_schema = schema.clone(); - let batch_size = context.session_config().batch_size(); let traversal = self.traversal.clone(); // Keep one input and one materialization in flight per partition. A hub's // remaining neighbours stay in the immutable adjacency until the next poll. @@ -2082,6 +2237,7 @@ impl ExecutionPlan for CsrExpandExec { schema.clone(), ); let (edge_filter, edge_state) = (edge_filter.clone(), edge_state.clone()); + let (target_state, scan_columns) = (target_state.clone(), scan_columns.clone()); let traversal = traversal.clone(); async move { loop { @@ -2142,10 +2298,30 @@ impl ExecutionPlan for CsrExpandExec { filter = None; } } + if let Some(target) = target.as_deref().filter(|_| !scan_columns.is_empty()) { + // Same decision as the edge mask: this chunk already uses the read. + use std::sync::atomic::Ordering::Relaxed; + let fetched = chunk.target_positions.len() as u64; + if target_state.columns.get().is_none() + && target_state.fetched.fetch_add(fetched, Relaxed) + fetched + >= scan_after + { + let _building = target_state.building.lock().await; + if target_state.columns.get().is_none() { + let columns = + target.rows.scan_columns(&scan_columns).await.map_err(|e| { + datafusion::common::DataFusionError::External(Box::new(e)) + })?; + let _ = target_state.columns.set(columns.map(Arc::new)); + } + } + } + let scanned = target_state.columns.get().cloned().flatten(); let batch = materialize( &expansion.batch, chunk, target.as_deref(), + scanned.as_deref(), edge.as_deref(), target_filter.as_deref(), filter, @@ -2652,7 +2828,7 @@ mod tests { let adj = Neighbors::directed(Arc::new(adj)); let mut expansion = BatchExpansion::try_new(input.clone(), 0, None).unwrap(); let chunk = expansion.next_chunk(&adj, None, 8192).await; - let out = materialize(&input, chunk, None, None, None, None, &out_schema) + let out = materialize(&input, chunk, None, None, None, None, None, &out_schema) .await .unwrap(); assert_eq!(out.num_rows(), 2); @@ -3019,6 +3195,7 @@ mod tests { expanded, Some(&target), None, + None, Some(&filter), None, &schema, @@ -3131,6 +3308,7 @@ mod tests { &input, chunk, None, + None, Some(&edge), None, Some(&filter), @@ -3148,7 +3326,7 @@ mod tests { expansion.edge_mask = Some(mask); let chunk = expansion.next_chunk(&adjacency, None, 8192).await; assert_eq!(chunk.edge_rows, vec![1, 2]); - let out = materialize(&input, chunk, None, Some(&edge), None, None, &schema) + let out = materialize(&input, chunk, None, None, Some(&edge), None, None, &schema) .await .unwrap(); assert_eq!(since(&out), vec![20, 30]); @@ -3171,4 +3349,73 @@ mod tests { assert!(!mask.contains(first) && mask.contains(first + 1)); assert!(mask.contains(second) && mask.contains(second + 1)); } + + #[tokio::test] + async fn target_columns_read_once_match_fetches_by_row_id() { + use arrow::array::{RecordBatchIterator, StringArray}; + + let nodes = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from((0..7).collect::>())) as ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec![ + Some("a"), + None, + Some("c"), + Some("d"), + Some("e"), + Some("f"), + Some("g"), + ])) as ArrayRef, + ), + ( + "age", + Arc::new(Int64Array::from(vec![70, 71, 72, 73, 74, 75, 76])) as ArrayRef, + ), + ]) + .unwrap(); + assert!(NodeRows::from_batch(nodes.clone(), "id") + .unwrap() + .unwrap() + .scan_columns(&[1]) + .await + .unwrap() + .is_none()); + // Three fragments and a deletion, so row ids are neither dense nor offsets. + for stable in [false, true] { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("nodes.lance"); + let reader = RecordBatchIterator::new(vec![Ok(nodes.clone())], nodes.schema()); + let params = lance::dataset::WriteParams { + max_rows_per_file: 3, + enable_stable_row_ids: stable, + ..Default::default() + }; + let mut dataset = Dataset::write(reader, path.to_str().unwrap(), Some(params)) + .await + .unwrap(); + dataset.delete("id = 1").await.unwrap(); + let rows = NodeRows::from_lance(Arc::new(dataset), "id") + .await + .unwrap() + .unwrap(); + let scanned = rows.scan_columns(&[1, 2]).await.unwrap().unwrap(); + let positions: Vec = [6, 0, 4, 6, 3] + .iter() + .map(|&id| rows.position(id).unwrap()) + .collect(); + for columns in [vec![1, 2], vec![2], vec![]] { + let expected = rows.take(&positions, &columns).await.unwrap(); + let got = scanned.take(&positions, &columns).unwrap(); + assert_eq!(got.len(), expected.len()); + for (got, expected) in got.iter().zip(&expected) { + assert_eq!(got.to_data(), expected.to_data(), "stable ids: {stable}"); + } + } + assert!(scanned.take(&[u64::MAX], &[2]).is_err()); + } + } } diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index 0d2fc0b..248ce60 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -112,6 +112,36 @@ def datasets(request, tmp_path): } +def test_target_columns_read_once_over_many_batches(tmp_path): + # Fetches more target rows than one batch and an eighth of the node table, so the + # expansion switches from per-batch takes to one read of the target columns. + n = 3000 + 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(6 * n)], pa.int64()), + "dst": pa.array([(i * 7 + 3) % n for i in range(6 * n)], pa.int64()), + } + ) + person = lance.write_dataset(people, tmp_path / "p.lance", max_rows_per_file=700) + person.delete("id % 13 = 0") + datasets = {"Person": lance.dataset(tmp_path / "p.lance"), "KNOWS": knows} + csr = CypherEngine(CONFIG, datasets, use_csr=True) + join = CypherEngine(CONFIG, datasets) + for query in [ + "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN a.id, b.id, b.name, b.age", + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE b.age > 40 RETURN a.id, b.name", + "MATCH (a:Person)<-[:KNOWS]-(b:Person) RETURN count(DISTINCT b.name) AS n", + ]: + assert rows(csr.execute(query)) == rows(join.execute(query)), query + + @pytest.mark.parametrize("query", QUERIES) def test_csr_matches_join_path(datasets, query): csr = CypherEngine(CONFIG, datasets, use_csr=True)