From 3129c90d0005b8b4322b8eacaf4efc2587679589 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:52:36 -0700 Subject: [PATCH 1/9] feat(graph): enforce cooperative query deadlines Add keyword-only timeout to the three Cypher execution entrypoints. Keep one deadline through streamed reads and release the owned stream while idle; preserve native Arrow reader errors and cooperative traversal cancellation. Issue #19. Submitted by David Hughes (@stugorf), with AI Collaboration. External benchmark runner adoption and performance measurements remain pending. --- crates/lance-graph-python/Cargo.toml | 2 +- crates/lance-graph-python/src/executor.rs | 71 ++++++- crates/lance-graph-python/src/graph.rs | 200 +++++++++++++++++++- crates/lance-graph/src/csr_expand.rs | 122 ++++++++++-- crates/lance-graph/src/shortest_path.rs | 142 ++++++++------ docs/api-reference.md | 17 +- docs/cypher.md | 16 ++ python/python/lance_csr/_internal/graph.pyi | 28 ++- python/python/tests/test_timeout.py | 82 ++++++++ python/python/tests/typing/public_api.py | 24 ++- 10 files changed, 605 insertions(+), 99 deletions(-) create mode 100644 python/python/tests/test_timeout.py diff --git a/crates/lance-graph-python/Cargo.toml b/crates/lance-graph-python/Cargo.toml index 3244cb3..f7a9862 100644 --- a/crates/lance-graph-python/Cargo.toml +++ b/crates/lance-graph-python/Cargo.toml @@ -21,4 +21,4 @@ lance-graph = { path = "../lance-graph" } serde = { version = "1", features = ["derive"] } serde_json = "1" pyo3 = { version = "0.28", features = ["extension-module", "abi3-py312", "py-clone"] } -tokio = { version = "1.37", features = ["rt-multi-thread", "macros"] } +tokio = { version = "1.37", features = ["rt-multi-thread", "macros", "sync", "time"] } diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index cf95aaf..6e619ad 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -13,7 +13,54 @@ // limitations under the License. use futures::Future; -use pyo3::{PyResult, Python}; +use pyo3::{ + exceptions::{PyTimeoutError, PyValueError}, + PyResult, Python, +}; +use std::time::Duration; + +pub(crate) const QUERY_TIMEOUT_MESSAGE: &str = "Query timeout exceeded"; + +/// One absolute deadline shared by planning, execution and all streamed reads. +#[derive(Clone, Copy)] +pub(crate) struct QueryDeadline(Option); + +impl QueryDeadline { + pub(crate) fn new(timeout: Option) -> PyResult { + let Some(seconds) = timeout else { + return Ok(Self(None)); + }; + let duration = Duration::try_from_secs_f64(seconds).map_err(|_| { + PyValueError::new_err("timeout must be a finite nonnegative number of seconds") + })?; + let instant = tokio::time::Instant::now() + .checked_add(duration) + .ok_or_else(|| PyValueError::new_err("timeout exceeds the supported deadline range"))?; + Ok(Self(Some(instant))) + } + + pub(crate) fn instant(self) -> Option { + self.0 + } + + pub(crate) fn check(self) -> PyResult<()> { + if self + .0 + .is_some_and(|deadline| tokio::time::Instant::now() >= deadline) + { + Err(PyTimeoutError::new_err(QUERY_TIMEOUT_MESSAGE)) + } else { + Ok(()) + } + } + + async fn expired(self) { + match self.0 { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending::<()>().await, + } + } +} pub const SIGNAL_CHECK_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100); @@ -56,7 +103,21 @@ impl BackgroundExecutor { where F::Output: Send, { - let future = Self::result_or_interrupt(future); + self.block_on_until(py, QueryDeadline(None), future) + } + + /// Apply a query deadline without changing unrelated executor callers. + pub(crate) fn block_on_until( + &self, + py: Option>, + deadline: QueryDeadline, + future: F, + ) -> PyResult + where + F::Output: Send, + { + deadline.check()?; + let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) } else { @@ -65,7 +126,7 @@ impl BackgroundExecutor { } } - async fn result_or_interrupt(future: F) -> PyResult + async fn result_or_interrupt(future: F, deadline: QueryDeadline) -> PyResult where F: Future + Send, F::Output: Send, @@ -84,7 +145,9 @@ impl BackgroundExecutor { }; tokio::select! { - result = future => Ok(result), + biased; + _ = deadline.expired() => Err(PyTimeoutError::new_err(QUERY_TIMEOUT_MESSAGE)), + result = future => { deadline.check()?; Ok(result) }, err = interrupt_future => err, } } diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 4c21596..ccd677f 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -47,6 +47,7 @@ use pyo3::{ }; use serde_json::Value as JsonValue; +use crate::executor::{QueryDeadline, QUERY_TIMEOUT_MESSAGE}; use crate::namespace::PyDirNamespace; use crate::RT; @@ -656,13 +657,16 @@ impl CypherQuery { /// >>> # Explicit strategy /// >>> from lance.graph import ExecutionStrategy /// >>> result = query.execute(datasets, strategy=ExecutionStrategy.DataFusion) - #[pyo3(signature = (datasets, strategy=None))] + #[pyo3(signature = (datasets, strategy=None, *, timeout=None))] fn execute( &self, py: Python, datasets: &Bound<'_, PyDict>, strategy: Option, + timeout: Option, ) -> PyResult> { + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; // Convert datasets to Arrow batches while holding the GIL let arrow_datasets = python_datasets_to_batches(datasets)?; @@ -674,7 +678,11 @@ impl CypherQuery { // Use RT.block_on with Some(py) like the scanner to_pyarrow method let result_batch = RT - .block_on(Some(py), inner_query.execute(arrow_datasets, rust_strategy))? + .block_on_until( + Some(py), + deadline, + inner_query.execute(arrow_datasets, rust_strategy), + )? .map_err(graph_error_to_pyerr)?; record_batch_to_python_table(py, &result_batch) @@ -1911,6 +1919,120 @@ impl RecordBatchReader for BlockingBatchReader { } } +/// A timed reader owns its stream independently of calls into Arrow's C interface. +struct TimedBatchState { + stream: Option, + expired: bool, +} + +struct DeadlineBatchReader { + schema: arrow_schema::SchemaRef, + state: Arc>, + deadline: tokio::time::Instant, + watchdog: tokio::task::JoinHandle<()>, + finished: bool, +} + +fn stream_timeout_error() -> arrow_schema::ArrowError { + arrow_schema::ArrowError::ExternalError(Box::new(std::io::Error::new( + std::io::ErrorKind::TimedOut, + QUERY_TIMEOUT_MESSAGE, + ))) +} + +impl Iterator for DeadlineBatchReader { + type Item = Result; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + let result = RT.runtime.block_on(async { + let mut state = self.state.lock().await; + if state.expired || tokio::time::Instant::now() >= self.deadline { + state.stream.take(); + state.expired = true; + return Some(Err(stream_timeout_error())); + } + let stream = state.stream.as_mut()?; + let result = tokio::select! { + biased; + _ = tokio::time::sleep_until(self.deadline) => None, + result = stream.next() => Some(result), + }; + match result { + None => { + state.stream.take(); + state.expired = true; + Some(Err(stream_timeout_error())) + } + Some(result) => { + if tokio::time::Instant::now() >= self.deadline { + state.stream.take(); + state.expired = true; + return Some(Err(stream_timeout_error())); + } + if result.is_none() { + state.stream.take(); + } + result.map(|r| { + r.map_err(|e| arrow_schema::ArrowError::ExternalError(Box::new(e))) + }) + } + } + }); + if result.as_ref().is_none_or(|r| r.is_err()) { + self.finished = true; + self.watchdog.abort(); + } + result + } +} + +impl RecordBatchReader for DeadlineBatchReader { + fn schema(&self) -> arrow_schema::SchemaRef { + self.schema.clone() + } +} + +impl Drop for DeadlineBatchReader { + fn drop(&mut self) { + self.watchdog.abort(); + if let Ok(mut state) = self.state.try_lock() { + state.stream.take(); + } + } +} + +fn reader_with_deadline( + stream: datafusion::physical_plan::SendableRecordBatchStream, + deadline: QueryDeadline, +) -> Box { + let schema = stream.schema(); + let Some(deadline) = deadline.instant() else { + return Box::new(BlockingBatchReader { schema, stream }); + }; + let state = Arc::new(tokio::sync::Mutex::new(TimedBatchState { + stream: Some(stream), + expired: false, + })); + let watched = state.clone(); + let watchdog = RT.runtime.spawn(async move { + tokio::time::sleep_until(deadline).await; + let mut state = watched.lock().await; + if state.stream.take().is_some() { + state.expired = true; + } + }); + Box::new(DeadlineBatchReader { + schema, + state, + deadline, + watchdog, + finished: false, + }) +} + /// Cypher query engine with cached catalog for efficient multi-query execution /// /// The engine registers each dataset once and reuses the catalog and DataFusion @@ -2042,6 +2164,8 @@ impl CypherEngine { /// Execute a Cypher query and return the full result as a ``pyarrow.Table``. /// /// The whole result is held in memory; use ``execute_batches`` for large outputs. + /// ``timeout`` is an optional nonnegative number of seconds for the whole query; + /// expiry raises ``TimeoutError``. Cancellation is cooperative. /// /// Parameters /// ---------- @@ -2058,18 +2182,22 @@ impl CypherEngine { /// If the query is invalid /// RuntimeError /// If query execution fails - #[pyo3(signature = (query, params = None))] + #[pyo3(signature = (query, params = None, *, timeout=None))] fn execute( &self, py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, + timeout: Option, ) -> PyResult> { + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; let state = self.snapshot(); let cypher_query = self.parse(query, params, &state)?; let result_batch = RT - .block_on( + .block_on_until( Some(py), + deadline, cypher_query .execute_with_catalog_and_context(state.catalog.clone(), state.context.clone()), )? @@ -2084,30 +2212,34 @@ impl CypherEngine { /// stream with bounded memory; joins, sorts and aggregations may still buffer. /// The reader keeps its snapshot alive, so a concurrent ``refresh()`` does not /// affect it. Planning errors raise immediately; execution errors raise while reading. - #[pyo3(signature = (query, params = None))] + /// ``timeout`` keeps one deadline including pauses. Initial expiry raises + /// ``TimeoutError``; expired reads raise ArrowInvalid: Query timeout exceeded. + #[pyo3(signature = (query, params = None, *, timeout=None))] fn execute_batches( &self, py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, + timeout: Option, ) -> PyResult> { use arrow::pyarrow::IntoPyArrow; + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; let state = self.snapshot(); let cypher_query = self.parse(query, params, &state)?; let stream = RT - .block_on( + .block_on_until( Some(py), + deadline, cypher_query.execute_stream_with_catalog_and_context( state.catalog.clone(), state.context.clone(), ), )? .map_err(graph_error_to_pyerr)?; - let reader: Box = Box::new(BlockingBatchReader { - schema: stream.schema(), - stream, - }); + deadline.check()?; + let reader = reader_with_deadline(stream, deadline); Ok(reader.into_pyarrow(py)?.unbind()) } @@ -2586,6 +2718,54 @@ mod tests { } } + fn retained_stream() -> ( + datafusion::physical_plan::SendableRecordBatchStream, + std::sync::Weak, + ) { + let values = Arc::new(arrow_array::Int64Array::from(vec![1_i64])); + let retained = Arc::downgrade(&values); + let schema = Arc::new(Schema::new(vec![arrow_schema::Field::new( + "id", + arrow_schema::DataType::Int64, + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(vec![Ok(batch)]), + ); + (Box::pin(stream), retained) + } + + #[test] + fn timed_reader_releases_unconsumed_stream_at_deadline() { + let (stream, retained) = retained_stream(); + let mut reader = reader_with_deadline(stream, QueryDeadline::new(Some(0.02)).unwrap()); + let stop = std::time::Instant::now() + std::time::Duration::from_secs(1); + while retained.upgrade().is_some() && std::time::Instant::now() < stop { + std::thread::sleep(std::time::Duration::from_millis(2)); + } + assert!( + retained.upgrade().is_none(), + "idle deadline must drop the owned stream" + ); + assert!(reader + .next() + .unwrap() + .unwrap_err() + .to_string() + .contains(QUERY_TIMEOUT_MESSAGE)); + assert!(reader.next().is_none()); + } + + #[test] + fn closing_timed_reader_releases_stream_before_deadline() { + let (stream, retained) = retained_stream(); + let reader = reader_with_deadline(stream, QueryDeadline::new(Some(10.0)).unwrap()); + drop(reader); + assert!(retained.upgrade().is_none()); + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 201a3cc..109ef0d 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -61,6 +61,21 @@ use futures::StreamExt; use lance::dataset::{Dataset, ProjectionRequest}; use lance_table::format::Fragment; +/// Yield after bounded CPU work without requiring a particular async runtime. +pub(crate) async fn cooperative_yield() { + let mut yielded = false; + std::future::poll_fn(|cx| { + if yielded { + std::task::Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await +} + /// Lance's row-address column (`lance_core::ROW_ADDR`, not re-exported by `lance`). const ROW_ADDR: &str = "_rowaddr"; @@ -247,6 +262,23 @@ impl Adjacency { }) } + /// Inspect one occurrence so cancelled work also stops through long tombstone runs. + fn next_candidate(&self, key: i64, offset: &mut usize) -> Option> { + let base = self.base_neighbors(key); + let position = *offset; + let neighbor = if position < base.len() { + base[position] + } else { + *self.added.get(&key)?.get(position - base.len())? + }; + *offset += 1; + let removed = self + .removed + .get(&key) + .is_some_and(|r| r.binary_search(&position).is_ok()); + Some((!removed).then_some(neighbor)) + } + /// Live edges: base + appended - removed. pub fn num_edges(&self) -> u64 { self.base.csr.num_edges() + self.added_count - self.removed_count @@ -773,7 +805,7 @@ pub struct Neighbors { /// Occurrence offsets are sufficient to pause even a high-degree source row. #[derive(Default)] -struct NeighborPosition { +pub(crate) struct NeighborPosition { forward: usize, reverse: usize, } @@ -816,6 +848,21 @@ impl Neighbors { ) } + /// Inspect one raw edge, including tombstones and suppressed reverse self-loops. + pub(crate) fn next_candidate( + &self, + key: i64, + position: &mut NeighborPosition, + ) -> Option> { + if let Some(candidate) = self.forward.next_candidate(key, &mut position.forward) { + return Some(candidate); + } + self.reverse + .as_ref()? + .next_candidate(key, &mut position.reverse) + .map(|candidate| candidate.filter(|&n| n as i64 != key)) + } + /// Edges visited over all keys (self-loops count once per direction in this estimate). pub fn num_edges(&self) -> u64 { self.forward.num_edges() + self.reverse.as_ref().map_or(0, |r| r.num_edges()) @@ -1155,19 +1202,27 @@ impl BatchExpansion { } /// Collect only the next chunk; target properties are fetched after this returns. - fn next_chunk( + async fn next_chunk( &mut self, adjacency: &Neighbors, target: Option<&TargetLookup>, batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + let mut work = 0_usize; 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) { + while let Some(candidate) = adjacency.next_candidate(key, &mut self.position) { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + let Some(neighbor) = candidate else { + continue; + }; let neighbor = neighbor as i64; if required.is_some_and(|r| r != neighbor) { continue; @@ -1397,7 +1452,9 @@ impl ExecutionPlan for CsrExpandExec { 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); + let chunk = expansion + .next_chunk(&adjacency, target.as_deref(), batch_size) + .await; if chunk.rows.is_empty() { pending = None; continue; @@ -1837,21 +1894,58 @@ 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); + let chunk = expansion.next_chunk(&adj, None, 8192).await; let out = materialize(&input, chunk, None, None, &out_schema) .await .unwrap(); 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()); + assert!(expansion.next_chunk(&adj, None, 8192).await.rows.is_empty()); // Output is split at `batch_size` rows. 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.next_chunk(&adj, None, 1).await.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()); + assert_eq!( + expansion.next_chunk(&adj, None, 1).await.neighbor_values, + vec![3] + ); + assert!(expansion.next_chunk(&adj, None, 1).await.rows.is_empty()); + } + + #[tokio::test] + async fn expansion_yields_before_scanning_all_rejected_neighbors() { + let adj = Neighbors::directed(Arc::new( + Adjacency::build( + edge_stream(vec![Some(1); 100_000], (2..100_002).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("bound", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(Int64Array::from(vec![-1])), + ], + ) + .unwrap(); + let mut expansion = BatchExpansion::try_new(batch, 0, Some(1)).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adj, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.position.forward <= 1024); + assert!(expansion.position.forward > 0); } #[tokio::test] @@ -1870,7 +1964,7 @@ mod tests { 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); + let chunk = expansion.next_chunk(&adj, None, 32).await; assert_eq!( chunk.neighbor_values, (2..=64).step_by(2).collect::>() @@ -1879,7 +1973,10 @@ mod tests { // 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!( + expansion.next_chunk(&adj, None, 32).await.neighbor_values[0], + 66 + ); assert_eq!(neighbors(&base, 0).len(), n); } @@ -1914,13 +2011,14 @@ mod tests { ); let mut expansion = BatchExpansion::try_new(input, 0, Some(1)).unwrap(); for _ in 0..2 { - let chunk = expansion.next_chunk(&adjacency, Some(&target), 1); + let chunk = expansion.next_chunk(&adjacency, Some(&target), 1).await; 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) + .await .rows .is_empty()); } diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 9c37dec..c756a4f 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -349,6 +349,7 @@ impl ExecutionPlan for ShortestPathExec { &spec, &schema, ) + .await } }); Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) @@ -386,7 +387,7 @@ async fn build_adjacency( /// Breadth-first search from `source` until every node in `targets` is reached (or the /// frontier empties, or `max_hops`). Returns each reached node's depth and, for every node /// at depth d > 0, its predecessors at depth d - 1. -fn search( +async fn search( adjacency: &Neighbors, source: i64, targets: &HashSet, @@ -400,7 +401,7 @@ fn search( let mut level = 0; while !frontier.is_empty() && !remaining.is_empty() && level < max_hops { level += 1; - let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents); + let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents).await; // Finish the level before stopping, so every shortest path is recorded. for n in &next { remaining.remove(n); @@ -411,7 +412,7 @@ fn search( } /// Finish the whole level so parallel edges and alternative shortest paths survive. -fn expand_level( +async fn expand_level( adjacency: &Neighbors, frontier: &[i64], level: u32, @@ -419,8 +420,17 @@ fn expand_level( parents: &mut FastMap>, ) -> Vec { let mut next = Vec::new(); + let mut work = 0_usize; for &node in frontier { - adjacency.for_each_neighbor(node, |n| { + let mut position = crate::csr_expand::NeighborPosition::default(); + while let Some(candidate) = adjacency.next_candidate(node, &mut position) { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + let Some(n) = candidate else { + continue; + }; let n = n as i64; match depth.get(&n) { Some(&d) if d == level => parents.entry(n).or_default().push(node), @@ -431,14 +441,14 @@ fn expand_level( next.push(n); } } - }); + } } next } /// Search a single pair using complete layers from the smaller frontier. At the /// first intersection every meeting node lies on the same depth cut of a shortest path. -fn search_bidirectional( +async fn search_bidirectional( forward: &Neighbors, reverse: &Neighbors, source: i64, @@ -459,7 +469,7 @@ fn search_bidirectional( while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops { if frontier.len() <= reverse_frontier.len() { level += 1; - frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents); + frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents).await; } else { reverse_level += 1; reverse_frontier = expand_level( @@ -468,7 +478,8 @@ fn search_bidirectional( reverse_level, &mut reverse_depth, &mut reverse_parents, - ); + ) + .await; } let mut pending: Vec<_> = frontier .iter() @@ -480,7 +491,12 @@ fn search_bidirectional( } // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths. let mut visited = HashSet::new(); + let mut work = 0_usize; while let Some(node) = pending.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } if !visited.insert(node) { continue; } @@ -499,10 +515,15 @@ fn search_bidirectional( /// Shortest paths from the search source to `target`, as node ids from source to target: /// all of them, or the first found. -fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { +async fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { let mut out = Vec::new(); let mut stack = vec![vec![target]]; + let mut work = 0_usize; while let Some(path) = stack.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } let head = *path.last().expect("non-empty"); match parents.get(&head) { None => { @@ -526,7 +547,7 @@ fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec = by_source - .iter() - .map(|(&s, ts)| { - let result = match reverse_adjacency.filter(|_| ts.len() == 1) { - Some(reverse) => search_bidirectional( + let mut searches = FastMap::default(); + for (&source, targets) in &by_source { + let result = match reverse_adjacency.filter(|_| targets.len() == 1) { + Some(reverse) => { + search_bidirectional( adjacency, reverse, - s, - *ts.iter().next().expect("one target"), + source, + *targets.iter().next().unwrap(), spec.max_hops, - ), - None => search(adjacency, s, ts, spec.max_hops), - }; - (s, result) - }) - .collect(); + ) + .await + } + None => search(adjacency, source, targets, spec.max_hops).await, + }; + searches.insert(source, result); + } let mut rows = Vec::new(); let mut lengths = Vec::new(); let mut nodes = ListBuilder::new(Int64Builder::new()); @@ -580,7 +602,7 @@ fn search_batch( if d < spec.min_hops || d > spec.max_hops { continue; } - for path in paths_to(t, parents, spec.all) { + for path in paths_to(t, parents, spec.all).await { rows.push(i as u32); lengths.push(d as i64); nodes.values().append_slice(&path); @@ -628,7 +650,7 @@ fn plan_error(message: &str) -> GraphError { mod tests { use super::*; - fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { + async fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { let schema = Arc::new(Schema::new(vec![ Field::new("src", DataType::Int64, false), Field::new("dst", DataType::Int64, false), @@ -652,40 +674,33 @@ mod tests { futures::stream::once(async move { Ok(batch) }), )) as SendableRecordBatchStream }; - let rt = tokio::runtime::Builder::new_current_thread() - .build() - .unwrap(); - let out = rt - .block_on(Adjacency::build(stream(), "src", "dst")) - .unwrap(); - let inc = rt - .block_on(Adjacency::build(stream(), "dst", "src")) - .unwrap(); + let out = Adjacency::build(stream(), "src", "dst").await.unwrap(); + let inc = Adjacency::build(stream(), "dst", "src").await.unwrap(); (Arc::new(out), Arc::new(inc)) } - fn neighbors(edges: &[(i64, i64)]) -> Neighbors { - let (out, inc) = adjacencies(edges); + async fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + let (out, inc) = adjacencies(edges).await; Neighbors::undirected(out, inc) } - #[test] - fn finds_every_shortest_path_and_only_shortest_ones() { + #[tokio::test] + async fn finds_every_shortest_path_and_only_shortest_ones() { // Two 2-hop routes 1-2-4 and 1-3-4, and a longer 1-5-6-4. - let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]); - let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10); + let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]).await; + let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10).await; assert_eq!(depth[&4], 2); - let mut all = paths_to(4, &parents, true); + let mut all = paths_to(4, &parents, true).await; all.sort(); assert_eq!(all, vec![vec![1, 2, 4], vec![1, 3, 4]]); - assert_eq!(paths_to(4, &parents, false).len(), 1); + assert_eq!(paths_to(4, &parents, false).await.len(), 1); // Beyond the hop limit, nothing is reached. - let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1); + let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1).await; assert!(!depth.contains_key(&4)); } - #[test] - fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { + #[tokio::test] + async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { // A wide first layer makes the target's side expand. Two meeting nodes, // including parallel edges on both sides of the cut, yield six paths. let (out, inc) = adjacencies(&[ @@ -700,11 +715,12 @@ mod tests { (4, 5), (5, 5), (5, 1), - ]); + ]) + .await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3); + let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3).await; assert_eq!(depth[&5], 3); - let mut all = paths_to(5, &parents, true); + let mut all = paths_to(5, &parents, true).await; all.sort(); assert_eq!( all, @@ -713,16 +729,17 @@ mod tests { .chain(vec![vec![1, 3, 4, 5]; 2]) .collect::>() ); - let one = paths_to(5, &parents, false); + let one = paths_to(5, &parents, false).await; assert_eq!(one.len(), 1); assert!(all.contains(&one[0])); assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2) + .await .0 .contains_key(&5)); } - #[test] - fn bidirectional_matches_forward_search_for_small_graphs() { + #[tokio::test] + async fn bidirectional_matches_forward_search_for_small_graphs() { // All directed graphs on three nodes, including disconnected pairs and cycles. let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]; for mask in 0..(1 << edges.len()) { @@ -731,7 +748,7 @@ mod tests { .enumerate() .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge)) .collect(); - let (out, inc) = adjacencies(&selected); + let (out, inc) = adjacencies(&selected).await; let directed = Neighbors::directed(out.clone()); let reverse = Neighbors::directed(inc.clone()); let undirected = Neighbors::undirected(out, inc); @@ -740,17 +757,18 @@ mod tests { for target in 0..3 { for max_hops in 0..=3 { let (expected_depth, expected_parents) = - search(forward, source, &HashSet::from([target]), max_hops); + search(forward, source, &HashSet::from([target]), max_hops).await; let (depth, parents) = - search_bidirectional(forward, reverse, source, target, max_hops); + search_bidirectional(forward, reverse, source, target, max_hops) + .await; assert_eq!( depth.get(&target), expected_depth.get(&target), "graph {mask}, {source}->{target}, max {max_hops}" ); if depth.contains_key(&target) { - let mut expected = paths_to(target, &expected_parents, true); - let mut actual = paths_to(target, &parents, true); + let mut expected = paths_to(target, &expected_parents, true).await; + let mut actual = paths_to(target, &parents, true).await; expected.sort(); actual.sort(); assert_eq!( @@ -765,8 +783,8 @@ mod tests { } } - #[test] - fn bidirectional_avoids_expanding_the_wide_middle_layer() { + #[tokio::test] + async fn bidirectional_avoids_expanding_the_wide_middle_layer() { let mut edges = Vec::new(); for branch in 1..=100 { edges.push((0, branch)); @@ -774,13 +792,13 @@ mod tests { edges.push((branch, 101 + (branch - 1) * 100 + leaf)); } } - let (out, inc) = adjacencies(&edges); + let (out, inc) = adjacencies(&edges).await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2); - assert_eq!(paths_to(101, &parents, true), vec![vec![0, 1, 101]]); + let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2).await; + assert_eq!(paths_to(101, &parents, true).await, vec![vec![0, 1, 101]]); assert_eq!(depth.len(), 102); // Root, first layer and stitched target. assert_eq!( - search(&forward, 0, &HashSet::from([101]), 2).0.len(), + search(&forward, 0, &HashSet::from([101]), 2).await.0.len(), 10_101 ); } diff --git a/docs/api-reference.md b/docs/api-reference.md index 4fe1dfb..e85af6f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -68,6 +68,17 @@ Keys match labels case-insensitively. --- +`CypherEngine.execute`, `CypherEngine.execute_batches` and `CypherQuery.execute` accept keyword-only +`timeout` in seconds (`None` means no deadline). Zero expires immediately; negative, nonfinite or +unrepresentable durations raise `ValueError`. One deadline includes planning, every batch and pauses +between reads. Before returning a reader, expiry raises `TimeoutError`. Reading an expired native Arrow +reader raises `pyarrow.ArrowInvalid` containing `Query timeout exceeded`. Earlier batches remain valid. +Expiry or reader close releases the stream even while the caller is idle; the engine remains usable. + +Cancellation is cooperative. Traversal loops yield at bounded work intervals, but synchronous parsing, +Python data providers and foreign blocking calls cannot be forcibly preempted. Deadlines are checked +between planning stages and before accepting completed results; this is not a hard real-time limit. + ## `GraphConfig` An immutable mapping from tables to a property graph. Create one with `GraphConfig.builder()`. @@ -128,12 +139,12 @@ Raises `ValueError` if `datasets` is empty, a configured key or endpoint column was opened through a namespace client and can't be reopened by URI. Raises `RuntimeError` if a dataset can't be opened. -`#!python execute(query: str, params: dict | None = None) -> pyarrow.Table` +`#!python execute(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, `bool`, `None`, lists, string-keyed nested maps, `datetime.datetime` and `datetime.date`. Integers must fit signed 64-bit values; nesting beyond 128 levels, cycles and unsupported values raise `ValueError`. -`#!python execute_batches(query: str, params: dict | None = None) -> pyarrow.RecordBatchReader` +`#!python execute_batches(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.RecordBatchReader` : Streams the result. Planning errors raise immediately; execution errors raise while reading. The reader keeps its snapshot alive across `refresh()`. @@ -178,7 +189,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde `#!python with_parameter(key: str, value: object) -> CypherQuery` : Binds `$key` to `value`. -`#!python execute(datasets: dict, strategy: ExecutionStrategy | None = None) -> pyarrow.Table` +`#!python execute(datasets: dict, strategy: ExecutionStrategy | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs the query. Datasets are read into memory for this call. `#!python execute_with_namespace(namespace: DirNamespace, strategy: ExecutionStrategy | None = None) -> pyarrow.Table` diff --git a/docs/cypher.md b/docs/cypher.md index e14dee7..14039bb 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -124,6 +124,22 @@ print(rows) # 4 - The reader keeps its snapshot alive, so a concurrent [`refresh()`](lance.md#refresh) doesn't affect it. - Dropping the reader cancels the remaining work. +### Query deadlines + +Pass `timeout=5.0` to `execute`, `execute_batches` or `CypherQuery.execute` to allow five seconds from the +call's start. Streams keep that same deadline while you pause between reads. The default is `None`. + +```python +try: + engine.execute("RETURN 1 AS n", timeout=0) +except TimeoutError: + print("Query timeout exceeded") # Query timeout exceeded +``` + +Initial expiry raises `TimeoutError`. Reading an expired native Arrow reader raises `pyarrow.ArrowInvalid` +with `Query timeout exceeded`; partial output already read remains valid. Close readers when you stop +consuming them. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. + ### `explain(query, params=None)` Returns the plan as text: the graph logical plan, DataFusion's logical plan and the physical plan. diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index 1c6b13a..e6544f7 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -10,9 +10,19 @@ class _ArrowStream(Protocol): _Dataset = Table | RecordBatch | RecordBatchReader | LanceDataset | _ArrowStream _DatasetT = TypeVar("_DatasetT", bound=_Dataset) -_Parameter = int | float | str | bool | date | datetime | None +_Parameter = ( + int + | float + | str + | bool + | date + | datetime + | None + | list["_Parameter"] + | dict[str, "_Parameter"] +) _ParameterT = TypeVar("_ParameterT", bound=_Parameter) -_Scalar = int | float | str | bool | None +_Scalar = int | float | str | bool | None | list["_Scalar"] | dict[str, "_Scalar"] class _NodeIndexStats(TypedDict): rows: int @@ -88,6 +98,8 @@ class CypherQuery: self, datasets: dict[str, _DatasetT], strategy: ExecutionStrategy | None = None, + *, + timeout: float | None = None, ) -> Table: ... def execute_with_namespace( self, namespace: DirNamespace, strategy: ExecutionStrategy | None = None @@ -106,10 +118,18 @@ class CypherEngine: self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False ) -> None: ... def execute( - self, query: str, params: dict[str, _ParameterT] | None = None + self, + query: str, + params: dict[str, _ParameterT] | None = None, + *, + timeout: float | None = None, ) -> Table: ... def execute_batches( - self, query: str, params: dict[str, _ParameterT] | None = None + self, + query: str, + params: dict[str, _ParameterT] | None = None, + *, + timeout: float | None = None, ) -> RecordBatchReader: ... def explain( self, query: str, params: dict[str, _ParameterT] | None = None diff --git a/python/python/tests/test_timeout.py b/python/python/tests/test_timeout.py new file mode 100644 index 0000000..87db8f9 --- /dev/null +++ b/python/python/tests/test_timeout.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""A single deadline covers execution and idle time in native Arrow readers.""" + +import time + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +CONFIG = GraphConfig.builder().node_table("Person", "id").build() +TABLES = {"Person": pa.table({"id": list(range(20_000))})} +QUERY = "MATCH (p:Person) RETURN p.id AS id" + + +@pytest.fixture +def engine(): + return CypherEngine(CONFIG, TABLES) + + +@pytest.mark.parametrize("timeout", [-1, float("nan"), float("inf"), 1e300]) +@pytest.mark.parametrize("method", ["execute", "execute_batches", "query"]) +def test_invalid_timeout(engine, timeout, method): + with pytest.raises(ValueError, match="timeout"): + if method == "query": + CypherQuery(QUERY).with_config(CONFIG).execute(TABLES, timeout=timeout) + else: + getattr(engine, method)(QUERY, timeout=timeout) + + +@pytest.mark.parametrize("method", ["execute", "execute_batches", "query"]) +def test_zero_timeout_then_engine_reuse(engine, method): + with pytest.raises(TimeoutError, match="Query timeout exceeded"): + if method == "query": + CypherQuery(QUERY).with_config(CONFIG).execute(TABLES, timeout=0) + else: + getattr(engine, method)(QUERY, timeout=0) + assert engine.execute("RETURN 1 AS n").to_pylist() == [{"n": 1}] + + +def test_native_reader_deadline_expires_while_idle(engine): + reader = engine.execute_batches(QUERY, timeout=1.0) + assert isinstance(reader, pa.RecordBatchReader) + time.sleep(1.05) + with pytest.raises(pa.ArrowInvalid, match="Query timeout exceeded"): + reader.read_next_batch() + reader.close() + assert engine.execute("RETURN 1 AS n", timeout=10).to_pylist() == [{"n": 1}] + + +def test_stream_keeps_one_deadline_between_reads(tmp_path): + dataset = lance.write_dataset( + TABLES["Person"], tmp_path / "people.lance", max_rows_per_file=4000 + ) + engine = CypherEngine(CONFIG, {"Person": dataset}) + reader = engine.execute_batches(QUERY, timeout=1.0) + first = reader.read_next_batch() + assert 0 < first.num_rows < 20_000 + time.sleep(1.05) + with pytest.raises(pa.ArrowInvalid, match="Query timeout exceeded"): + reader.read_all() + + +def test_no_deadline_and_normal_execution_errors(engine): + assert engine.execute(QUERY, timeout=None).num_rows == 20_000 + assert engine.execute_batches(QUERY, timeout=10).read_all().num_rows == 20_000 + with pytest.raises(ValueError): + engine.execute("RETURN unknown", timeout=10) + + +def test_expensive_query_stops_and_engine_stays_usable(engine): + started = time.monotonic() + with pytest.raises(TimeoutError, match="Query timeout exceeded"): + engine.execute( + "MATCH (a:Person), (b:Person), (c:Person) " + "WHERE a.id + b.id + c.id > 30000 RETURN count(*) AS n", + timeout=0.02, + ) + assert time.monotonic() - started < 2.0 + assert engine.execute("RETURN 2 AS n").to_pylist() == [{"n": 2}] diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py index 1edc328..c4e3704 100644 --- a/python/python/tests/typing/public_api.py +++ b/python/python/tests/typing/public_api.py @@ -1,7 +1,7 @@ """Strict static consumer; checked by make lint, not executed by pytest.""" from datetime import date -from typing import assert_type +from typing import TYPE_CHECKING, TypeAlias, assert_type from lance import LanceDataset from lance_csr import ( @@ -20,6 +20,19 @@ ) from pyarrow import RecordBatch, RecordBatchReader, Table +if TYPE_CHECKING: + from collections.abc import Mapping + +ParameterResult: TypeAlias = ( + int + | float + | str + | bool + | None + | list["ParameterResult"] + | dict[str, "ParameterResult"] +) + def public_api( table: Table, batch: RecordBatch, reader: RecordBatchReader, dataset: LanceDataset @@ -56,9 +69,14 @@ def public_api( one = CypherQuery(query).with_config(config).with_parameter("id", 1) assert_type(one, CypherQuery) + assert_type(one.with_parameter("ids", [1, 2]), CypherQuery) + assert_type(engine.execute(query, timeout=1.0), Table) + assert_type(engine.execute_batches(query, timeout=1.0), RecordBatchReader) + assert_type(one.execute(datasets, timeout=1.0), Table) one.with_parameter("day", date(2026, 1, 1)) assert_type(one.query_text(), str) - assert_type(one.parameters(), dict[str, int | float | str | bool | None]) + parameters: Mapping[str, ParameterResult] = one.parameters() + assert_type(len(parameters), int) assert_type(one.variables(), list[str]) assert_type(one.node_labels(), list[str]) assert_type(one.relationship_types(), list[str]) @@ -91,5 +109,5 @@ def public_api( # Rejected calls are checked too: removing a type restriction makes these unused. engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] vector.top_k("5") # pyright: ignore[reportArgumentType] - one.with_parameter("ids", [1, 2]) # pyright: ignore[reportArgumentType] + one.with_parameter("ids", {1, 2}) # pyright: ignore[reportArgumentType] namespace.base_uri = "elsewhere" # pyright: ignore[reportAttributeAccessIssue] From 63a4928202455955f079ebcbf9672075f75ad5ca Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:36:54 -0700 Subject: [PATCH 2/9] fix(python): release timed streams after execution errors Drop the owned DataFusion stream when its first error finishes the Arrow reader, preserving the original error. A real CAST failure and Weak reference to a later Arrow batch fail before the fix and pass after it. All 15 native binding tests pass; the Python reader-error regression and complete native rebuild are queued. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-python/src/graph.rs | 45 +++++++++++++++++++++++++- python/python/tests/test_timeout.py | 10 ++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index ccd677f..240d206 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -1972,7 +1972,7 @@ impl Iterator for DeadlineBatchReader { state.expired = true; return Some(Err(stream_timeout_error())); } - if result.is_none() { + if result.as_ref().is_none_or(|result| result.is_err()) { state.stream.take(); } result.map(|r| { @@ -2766,6 +2766,49 @@ mod tests { assert!(retained.upgrade().is_none()); } + #[test] + fn timed_reader_releases_stream_after_an_execution_error() { + let (stream, retained) = RT.runtime.block_on(async { + let values = Arc::new(arrow_array::StringArray::from(vec!["2"])); + let retained = Arc::downgrade(&values); + let schema = Arc::new(Schema::new(vec![arrow_schema::Field::new( + "value", + arrow_schema::DataType::Utf8, + false, + )])); + let bad = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::StringArray::from(vec!["invalid"]))], + ) + .unwrap(); + let remaining = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "input", + Arc::new(MemTable::try_new(schema, vec![vec![bad, remaining]]).unwrap()), + ) + .unwrap(); + let stream = ctx + .sql("SELECT CAST(value AS BIGINT) AS value FROM input") + .await + .unwrap() + .execute_stream() + .await + .unwrap(); + (stream, retained) + }); + let mut reader = reader_with_deadline(stream, QueryDeadline::new(Some(10.0)).unwrap()); + assert!(retained.upgrade().is_some()); + let error = reader.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("invalid"), "{error}"); + assert!(!error.contains(QUERY_TIMEOUT_MESSAGE), "{error}"); + assert!( + retained.upgrade().is_none(), + "error must release the owned stream" + ); + assert!(reader.next().is_none()); + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/python/python/tests/test_timeout.py b/python/python/tests/test_timeout.py index 87db8f9..9717815 100644 --- a/python/python/tests/test_timeout.py +++ b/python/python/tests/test_timeout.py @@ -70,6 +70,16 @@ def test_no_deadline_and_normal_execution_errors(engine): engine.execute("RETURN unknown", timeout=10) +def test_execution_error_in_a_timed_stream_is_preserved(engine): + reader = engine.execute_batches("MATCH (p:Person) RETURN 1 / p.id AS n", timeout=10) + with pytest.raises(pa.ArrowInvalid) as raised: + reader.read_next_batch() + assert "Query timeout exceeded" not in str(raised.value) + with pytest.raises(StopIteration): + reader.read_next_batch() + assert engine.execute("RETURN 3 AS n").to_pylist() == [{"n": 3}] + + def test_expensive_query_stops_and_engine_stays_usable(engine): started = time.monotonic() with pytest.raises(TimeoutError, match="Query timeout exceeded"): From 49c1e15875f5bb3041ca31cb5ca3ebf24ad248ec Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:41:42 -0700 Subject: [PATCH 3/9] fix(query): stop Cartesian work at deadline checkpoints --- crates/lance-graph-python/src/executor.rs | 4 + crates/lance-graph/src/csr_expand.rs | 94 +++++++++------- crates/lance-graph/src/query.rs | 102 +++++++++++++++++ crates/lance-graph/src/query/deadline.rs | 127 ++++++++++++++++++++++ crates/lance-graph/src/shortest_path.rs | 2 +- docs/api-reference.md | 4 +- docs/csr.md | 4 +- docs/cypher.md | 3 +- 8 files changed, 299 insertions(+), 41 deletions(-) create mode 100644 crates/lance-graph/src/query/deadline.rs diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index 6e619ad..1fe65ac 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -117,6 +117,10 @@ impl BackgroundExecutor { F::Output: Send, { deadline.check()?; + let future = lance_graph::query::with_execution_deadline( + deadline.instant().map(tokio::time::Instant::into_std), + future, + ); let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 109ef0d..af997f8 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -243,6 +243,15 @@ impl Adjacency { key: i64, offset: &'a mut usize, ) -> impl Iterator + 'a { + self.candidates_from(key, offset).flatten() + } + + /// Keep tombstones visible to cancellation without repeating lookups for every edge. + fn candidates_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); @@ -251,7 +260,7 @@ impl Adjacency { .chain(added) .enumerate() .skip(*offset) - .filter_map(move |(position, &neighbor)| { + .map(move |(position, &neighbor)| { *offset = position + 1; if removed.first() == Some(&position) { removed = &removed[1..]; @@ -262,23 +271,6 @@ impl Adjacency { }) } - /// Inspect one occurrence so cancelled work also stops through long tombstone runs. - fn next_candidate(&self, key: i64, offset: &mut usize) -> Option> { - let base = self.base_neighbors(key); - let position = *offset; - let neighbor = if position < base.len() { - base[position] - } else { - *self.added.get(&key)?.get(position - base.len())? - }; - *offset += 1; - let removed = self - .removed - .get(&key) - .is_some_and(|r| r.binary_search(&position).is_ok()); - Some((!removed).then_some(neighbor)) - } - /// Live edges: base + appended - removed. pub fn num_edges(&self) -> u64 { self.base.csr.num_edges() + self.added_count - self.removed_count @@ -835,34 +827,28 @@ impl Neighbors { key: i64, position: &'a mut NeighborPosition, ) -> impl Iterator + 'a { - let forward = self.forward.neighbors_from(key, &mut position.forward); + self.candidates_from(key, position).flatten() + } + + /// Borrow each adjacency once while exposing tombstones and suppressed self-loops as work. + pub(crate) fn candidates_from<'a>( + &'a self, + key: i64, + position: &'a mut NeighborPosition, + ) -> impl Iterator> + 'a { + let forward = self.forward.candidates_from(key, &mut position.forward); let reverse = self .reverse .as_ref() - .map(|r| r.neighbors_from(key, &mut position.reverse)); + .map(|r| r.candidates_from(key, &mut position.reverse)); forward.chain( reverse .into_iter() .flatten() - .filter(move |&n| n as i64 != key), + .map(move |candidate| candidate.filter(|&n| n as i64 != key)), ) } - /// Inspect one raw edge, including tombstones and suppressed reverse self-loops. - pub(crate) fn next_candidate( - &self, - key: i64, - position: &mut NeighborPosition, - ) -> Option> { - if let Some(candidate) = self.forward.next_candidate(key, &mut position.forward) { - return Some(candidate); - } - self.reverse - .as_ref()? - .next_candidate(key, &mut position.reverse) - .map(|candidate| candidate.filter(|&n| n as i64 != key)) - } - /// Edges visited over all keys (self-loops count once per direction in this estimate). pub fn num_edges(&self) -> u64 { self.forward.num_edges() + self.reverse.as_ref().map_or(0, |r| r.num_edges()) @@ -1215,7 +1201,7 @@ impl BatchExpansion { 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)); - while let Some(candidate) = adjacency.next_candidate(key, &mut self.position) { + for candidate in adjacency.candidates_from(key, &mut self.position) { work += 1; if work.is_multiple_of(1024) { cooperative_yield().await; @@ -1948,6 +1934,40 @@ mod tests { assert!(expansion.position.forward > 0); } + #[tokio::test] + async fn expansion_yields_and_resumes_through_a_long_tombstone_run() { + let n = 5_000; + let base = Adjacency::build( + edge_stream(vec![Some(1); n], (2..n as i64 + 2).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(); + let removed = (2..n as i64 + 1).map(|id| (1, id)).collect::>(); + let adjacency = Neighbors::directed(Arc::new(base.with_changes(&[], &removed))); + let input = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert_eq!(expansion.position.forward, 1024); + assert_eq!( + expansion + .next_chunk(&adjacency, None, 1) + .await + .neighbor_values, + vec![n as i64 + 1] + ); + assert!(expansion + .next_chunk(&adjacency, None, 1) + .await + .rows + .is_empty()); + } + #[tokio::test] async fn high_degree_expansion_visits_only_the_first_chunk() { let n = 100_000; diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index f617431..288c3d9 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -3,6 +3,11 @@ //! High-level Cypher query interface for Lance datasets +mod deadline; + +#[doc(hidden)] +pub use deadline::with_execution_deadline; + use crate::ast::CypherQuery as CypherAST; use crate::ast::ReadingClause; use crate::config::GraphConfig; @@ -625,6 +630,7 @@ impl CypherQuery { ctx: datafusion::execution::context::SessionContext, ) -> Result { let (_logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; + let ctx = deadline::context(ctx); ctx.execute_logical_plan(df_logical_plan) .await .map_err(|e| GraphError::ExecutionError { @@ -916,6 +922,7 @@ impl CypherQuery { let (logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; // Phase 4: DataFusion Physical Plan + let ctx = deadline::context(ctx.clone()); let df = ctx .execute_logical_plan(df_logical_plan.clone()) .await @@ -1333,6 +1340,101 @@ mod tests { use super::*; use crate::config::GraphConfig; + #[tokio::test] + async fn cartesian_stream_yields_and_drops_before_full_collection() { + use arrow_array::{ArrayRef, Int64Array}; + use datafusion::datasource::{DefaultTableSource, MemTable}; + use datafusion::execution::context::{SessionConfig, SessionContext}; + use futures::TryStreamExt; + use lance_graph_catalog::InMemoryCatalog; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef, + )]) + .unwrap(); + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap()); + let catalog = Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(table))), + ); + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(1) + .with_round_robin_repartition(false), + ); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let query = CypherQuery::new( + "MATCH (a:Person), (b:Person) RETURN a.id AS left_id, b.id AS right_id", + ) + .unwrap() + .with_config(config.clone()); + let stream = with_execution_deadline( + Some(std::time::Instant::now() + std::time::Duration::from_secs(60)), + query.execute_stream_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap(); + // A fresh Tokio task supplies its own budget; no inherited source yield can pass this test. + tokio::spawn(async move { + let visited = Arc::new(AtomicUsize::new(0)); + let progress = visited.clone(); + let mut read = Box::pin(stream.try_fold(0_usize, move |count, batch| { + progress.fetch_add(batch.num_rows(), Ordering::SeqCst); + futures::future::ready(Ok(count + batch.num_rows())) + })); + assert!(futures::poll!(read.as_mut()).is_pending()); + let stopped = visited.load(Ordering::SeqCst); + assert!(stopped > 0 && stopped < 1_000_000); + drop(read); + tokio::task::yield_now().await; + assert_eq!(visited.load(Ordering::SeqCst), stopped); + }) + .await + .unwrap(); + let result = query + .execute_with_catalog_and_context(catalog.clone(), ctx.clone()) + .await + .unwrap(); + assert_eq!(result.num_rows(), 1_000_000); + assert_eq!(result.num_columns(), 2); + + let filtered = CypherQuery::new( + "MATCH (a:Person), (b:Person) WHERE a.id + b.id >= 999 RETURN count(*) AS n", + ) + .unwrap() + .with_config(config); + let expired = with_execution_deadline( + Some(std::time::Instant::now() - std::time::Duration::from_secs(1)), + filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap_err(); + assert!(expired.to_string().contains("Query timeout exceeded")); + for timed in [false, true] { + let result = with_execution_deadline( + timed.then(|| std::time::Instant::now() + std::time::Duration::from_secs(60)), + filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 500_500 + ); + } + } + #[test] fn test_parse_simple_cypher_query() { let query = CypherQuery::new("MATCH (n:Person) RETURN n.name").unwrap(); diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs new file mode 100644 index 0000000..0633eff --- /dev/null +++ b/crates/lance-graph/src/query/deadline.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Internal deadline propagation from the Python bridge into physical join work. + +use std::future::Future; +use std::sync::Arc; +use std::time::Instant; + +use arrow_schema::{DataType, Field, Schema}; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{DataFusionError, Result, ScalarValue}; +use datafusion::execution::context::SessionContext; +use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::logical_expr::{create_udf, ColumnarValue, Operator, Volatility}; +use datafusion::physical_expr::expressions::BinaryExpr; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::coop::CooperativeExec; +use datafusion::physical_plan::joins::utils::JoinFilter; +use datafusion::physical_plan::joins::{ + CrossJoinExec, NestedLoopJoinExec, NestedLoopJoinExecBuilder, +}; +use datafusion::physical_plan::ExecutionPlan; + +tokio::task_local! { + static EXECUTION_DEADLINE: Option; +} + +/// Internal cross-crate bridge plumbing; no deadline is published into an engine snapshot. +#[doc(hidden)] +pub async fn with_execution_deadline(deadline: Option, future: F) -> F::Output { + EXECUTION_DEADLINE.scope(deadline, future).await +} + +pub(super) fn context(ctx: SessionContext) -> SessionContext { + let Some(deadline) = EXECUTION_DEADLINE + .try_with(|deadline| *deadline) + .ok() + .flatten() + else { + return ctx; + }; + SessionContext::new_with_state( + SessionStateBuilder::new_from_existing(ctx.state()) + .with_physical_optimizer_rule(Arc::new(DeadlineJoins { deadline })) + .build(), + ) +} + +#[derive(Debug)] +struct DeadlineJoins { + deadline: Instant, +} + +impl PhysicalOptimizerRule for DeadlineJoins { + fn name(&self) -> &str { + "query_deadline_joins" + } + + fn schema_check(&self) -> bool { + true + } + + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + let deadline = self.deadline; + let function = create_udf( + "__query_deadline", + vec![], + DataType::Boolean, + Volatility::Volatile, + Arc::new(move |_| { + if Instant::now() >= deadline { + return Err(DataFusionError::Execution("Query timeout exceeded".into())); + } + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))) + }), + ); + let guard: Arc = Arc::new(ScalarFunctionExpr::new( + "__query_deadline", + Arc::new(function), + vec![], + Arc::new(Field::new("__query_deadline", DataType::Boolean, false)), + Arc::new(config.clone()), + )); + plan.transform_up(|plan| { + if let Some(join) = plan.as_any().downcast_ref::() { + // Output-only yielding cannot interrupt long runs of rejected join candidates. + let filter = match join.filter() { + Some(filter) => JoinFilter::new( + Arc::new(BinaryExpr::new( + guard.clone(), + Operator::And, + filter.expression().clone(), + )), + filter.column_indices().to_vec(), + filter.schema().clone(), + ), + None => JoinFilter::new(guard.clone(), vec![], Arc::new(Schema::empty())), + }; + let joined = NestedLoopJoinExecBuilder::new( + join.left().clone(), + join.right().clone(), + *join.join_type(), + ) + .with_filter(Some(filter)) + .with_projection_ref(join.projection().clone()) + .build()?; + Ok(Transformed::yes( + Arc::new(CooperativeExec::new(Arc::new(joined))) as Arc, + )) + } else if plan.as_any().is::() { + Ok(Transformed::yes( + Arc::new(CooperativeExec::new(plan)) as Arc + )) + } else { + Ok(Transformed::no(plan)) + } + }) + .map(|plan| plan.data) + } +} diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index c756a4f..133888c 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -423,7 +423,7 @@ async fn expand_level( let mut work = 0_usize; for &node in frontier { let mut position = crate::csr_expand::NeighborPosition::default(); - while let Some(candidate) = adjacency.next_candidate(node, &mut position) { + for candidate in adjacency.candidates_from(node, &mut position) { work += 1; if work.is_multiple_of(1024) { crate::csr_expand::cooperative_yield().await; diff --git a/docs/api-reference.md b/docs/api-reference.md index e85af6f..97eed26 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -75,7 +75,9 @@ between reads. Before returning a reader, expiry raises `TimeoutError`. Reading reader raises `pyarrow.ArrowInvalid` containing `Query timeout exceeded`. Earlier batches remain valid. Expiry or reader close releases the stream even while the caller is idle; the engine remains usable. -Cancellation is cooperative. Traversal loops yield at bounded work intervals, but synchronous parsing, +Cancellation is cooperative. Traversal loops yield at bounded work intervals, and timed cross joins +yield between output batches. Timed nested-loop joins check the deadline while evaluating candidates, +including batches that produce no matches. Synchronous parsing, Python data providers and foreign blocking calls cannot be forcibly preempted. Deadlines are checked between planning stages and before accepting completed results; this is not a hard real-time limit. diff --git a/docs/csr.md b/docs/csr.md index 0d3ae5a..4199320 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -95,7 +95,9 @@ Node index 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 + affected sources' neighbour lists. Cancellation checkpoints count deleted edges and suppressed + reverse self-loops as work, using borrowed adjacency slices while a chunk is being collected. + 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. diff --git a/docs/cypher.md b/docs/cypher.md index 14039bb..da3e389 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -138,7 +138,8 @@ except TimeoutError: Initial expiry raises `TimeoutError`. Reading an expired native Arrow reader raises `pyarrow.ArrowInvalid` with `Query timeout exceeded`; partial output already read remains valid. Close readers when you stop -consuming them. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. +consuming them. Timed cross joins yield between batches, and nested-loop joins check the deadline +while evaluating candidates, including when no rows match. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. ### `explain(query, params=None)` From 2d9ef765450099f412aa19ffd8b3cf5d7eaebd15 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:46:43 -0700 Subject: [PATCH 4/9] fix(query): wire existing runtime and physical plan APIs --- crates/lance-graph-python/Cargo.lock | 1 + crates/lance-graph/Cargo.toml | 1 + crates/lance-graph/src/query/deadline.rs | 4 ++-- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/lance-graph-python/Cargo.lock b/crates/lance-graph-python/Cargo.lock index 248dee7..45e9ed2 100644 --- a/crates/lance-graph-python/Cargo.lock +++ b/crates/lance-graph-python/Cargo.lock @@ -3983,6 +3983,7 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.9", + "tokio", ] [[package]] diff --git a/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index b054f60..5276c0e 100644 --- a/crates/lance-graph/Cargo.toml +++ b/crates/lance-graph/Cargo.toml @@ -40,6 +40,7 @@ nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" snafu = "0.8" +tokio = { version = "1.37", features = ["rt"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs index 0633eff..d882c44 100644 --- a/crates/lance-graph/src/query/deadline.rs +++ b/crates/lance-graph/src/query/deadline.rs @@ -89,7 +89,7 @@ impl PhysicalOptimizerRule for DeadlineJoins { Arc::new(config.clone()), )); plan.transform_up(|plan| { - if let Some(join) = plan.as_any().downcast_ref::() { + if let Some(join) = plan.downcast_ref::() { // Output-only yielding cannot interrupt long runs of rejected join candidates. let filter = match join.filter() { Some(filter) => JoinFilter::new( @@ -114,7 +114,7 @@ impl PhysicalOptimizerRule for DeadlineJoins { Ok(Transformed::yes( Arc::new(CooperativeExec::new(Arc::new(joined))) as Arc, )) - } else if plan.as_any().is::() { + } else if plan.is::() { Ok(Transformed::yes( Arc::new(CooperativeExec::new(plan)) as Arc )) From 0189c99aac4de1ff7ea85a9032bd9824f965f1b8 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:49:41 -0700 Subject: [PATCH 5/9] fix(query): share one monotonic clock through deadline planning --- crates/lance-graph-python/src/executor.rs | 5 +---- crates/lance-graph/Cargo.toml | 2 +- crates/lance-graph/src/query.rs | 6 +++--- crates/lance-graph/src/query/deadline.rs | 2 +- 4 files changed, 6 insertions(+), 9 deletions(-) diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index 1fe65ac..08c16d7 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -117,10 +117,7 @@ impl BackgroundExecutor { F::Output: Send, { deadline.check()?; - let future = lance_graph::query::with_execution_deadline( - deadline.instant().map(tokio::time::Instant::into_std), - future, - ); + let future = lance_graph::query::with_execution_deadline(deadline.instant(), future); let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) diff --git a/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index 5276c0e..f0b9ad5 100644 --- a/crates/lance-graph/Cargo.toml +++ b/crates/lance-graph/Cargo.toml @@ -40,7 +40,7 @@ nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" snafu = "0.8" -tokio = { version = "1.37", features = ["rt"] } +tokio = { version = "1.37", features = ["rt", "time"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 288c3d9..5f0b1b7 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1374,7 +1374,7 @@ mod tests { .unwrap() .with_config(config.clone()); let stream = with_execution_deadline( - Some(std::time::Instant::now() + std::time::Duration::from_secs(60)), + Some(tokio::time::Instant::now() + std::time::Duration::from_secs(60)), query.execute_stream_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await @@ -1409,7 +1409,7 @@ mod tests { .unwrap() .with_config(config); let expired = with_execution_deadline( - Some(std::time::Instant::now() - std::time::Duration::from_secs(1)), + Some(tokio::time::Instant::now() - std::time::Duration::from_secs(1)), filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await @@ -1417,7 +1417,7 @@ mod tests { assert!(expired.to_string().contains("Query timeout exceeded")); for timed in [false, true] { let result = with_execution_deadline( - timed.then(|| std::time::Instant::now() + std::time::Duration::from_secs(60)), + timed.then(|| tokio::time::Instant::now() + std::time::Duration::from_secs(60)), filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs index d882c44..cb6b26c 100644 --- a/crates/lance-graph/src/query/deadline.rs +++ b/crates/lance-graph/src/query/deadline.rs @@ -5,7 +5,7 @@ use std::future::Future; use std::sync::Arc; -use std::time::Instant; +use tokio::time::Instant; use arrow_schema::{DataType, Field, Schema}; use datafusion::common::config::ConfigOptions; From 4f2a901fb6981b889d6e0a9808983e4960bcb385 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 21:52:53 -0700 Subject: [PATCH 6/9] perf(csr): scan cancellation quanta with local cursor state --- crates/lance-graph/src/csr_expand.rs | 93 ++++++++++++++++++++++------ 1 file changed, 73 insertions(+), 20 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index af997f8..91feeb9 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -256,12 +256,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) - .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()) + .map(move |&neighbor| { + let position = *offset; + *offset += 1; if removed.first() == Some(&position) { removed = &removed[1..]; None @@ -796,7 +798,7 @@ pub struct Neighbors { } /// Occurrence offsets are sufficient to pause even a high-degree source row. -#[derive(Default)] +#[derive(Clone, Copy, Default)] pub(crate) struct NeighborPosition { forward: usize, reverse: usize, @@ -1195,17 +1197,37 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + loop { + if self.fill_chunk(adjacency, target, batch_size, &mut chunk) { + cooperative_yield().await; + } + if self.row == self.keys.len() || chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + + /// Keep coroutine state outside the scan while counting deleted occurrences as work. + fn fill_chunk( + &mut self, + adjacency: &Neighbors, + target: Option<&TargetLookup>, + batch_size: usize, + chunk: &mut ExpandedRows, + ) -> bool { + let mut position = self.position; + let mut row = self.row; let mut work = 0_usize; - while self.row < self.keys.len() { - let i = 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 candidate in adjacency.candidates_from(key, &mut self.position) { + for candidate in adjacency + .candidates_from(key, &mut position) + .take(1024 - work) + { work += 1; - if work.is_multiple_of(1024) { - cooperative_yield().await; - } let Some(neighbor) = candidate else { continue; }; @@ -1214,24 +1236,28 @@ impl BatchExpansion { 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 { + let Some(target_row) = target.rows.position(neighbor) else { continue; }; - chunk.target_positions.push(position); + chunk.target_positions.push(target_row); } 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; + break 'rows; } } + if work == 1024 { + break; + } } - self.row += 1; - self.position = NeighborPosition::default(); + row += 1; + position = NeighborPosition::default(); } - chunk + self.row = row; + self.position = position; + work == 1024 } } @@ -1795,6 +1821,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 cd7020ef8ac723651c52f77e8ff6f5390c627f54 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 21:55:15 -0700 Subject: [PATCH 7/9] perf(csr): inline the complete raw and live cursor boundary --- crates/lance-graph/src/csr_expand.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 91feeb9..920a9e8 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -230,6 +230,7 @@ impl Adjacency { self.neighbors_from(key, &mut 0).for_each(f); } + #[inline(always)] fn base_neighbors(&self, key: i64) -> &[u64] { self.base .dense @@ -238,6 +239,7 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, @@ -247,6 +249,7 @@ impl Adjacency { } /// Keep tombstones visible to cancellation without repeating lookups for every edge. + #[inline(always)] fn candidates_from<'a>( &'a self, key: i64, @@ -824,6 +827,7 @@ impl Neighbors { .for_each(f); } + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, @@ -833,6 +837,7 @@ impl Neighbors { } /// Borrow each adjacency once while exposing tombstones and suppressed self-loops as work. + #[inline(always)] pub(crate) fn candidates_from<'a>( &'a self, key: i64, From b84e62866aa94effa07bb0742a4b26c3dc83c9d5 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 22:46:53 -0700 Subject: [PATCH 8/9] fix(query): cancel traversal work that produces no edges --- crates/lance-graph/src/csr_expand.rs | 33 +++++++++++++++++++++++++ crates/lance-graph/src/shortest_path.rs | 31 +++++++++++++++++++++++ docs/csr.md | 2 ++ 3 files changed, 66 insertions(+) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 920a9e8..07903b5 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -1225,6 +1225,7 @@ impl BatchExpansion { let mut work = 0_usize; 'rows: while row < self.keys.len() { let i = row; + let previous_work = work; 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)); @@ -1257,8 +1258,14 @@ impl BatchExpansion { break; } } + if work == previous_work { + work += 1; + } row += 1; position = NeighborPosition::default(); + if work == 1024 { + break; + } } self.row = row; self.position = position; @@ -1961,6 +1968,32 @@ mod tests { assert!(expansion.next_chunk(&adj, None, 1).await.rows.is_empty()); } + #[tokio::test] + async fn cancellation_without_edges_yields_while_skipping_source_rows() { + let adjacency = Neighbors::directed(Arc::new( + Adjacency::build(edge_stream(vec![Some(1)], vec![Some(2)]), "src", "dst") + .await + .unwrap(), + )); + for missing in [Some(99), None] { + let input = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![missing; 100_000])) as ArrayRef, + )]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.row > 0 && expansion.row <= 1024); + assert!(expansion + .next_chunk(&adjacency, None, 1) + .await + .rows + .is_empty()); + } + } + #[tokio::test] async fn expansion_yields_before_scanning_all_rejected_neighbors() { let adj = Neighbors::directed(Arc::new( diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 133888c..305e782 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -422,6 +422,7 @@ async fn expand_level( let mut next = Vec::new(); let mut work = 0_usize; for &node in frontier { + let previous_work = work; let mut position = crate::csr_expand::NeighborPosition::default(); for candidate in adjacency.candidates_from(node, &mut position) { work += 1; @@ -442,6 +443,12 @@ async fn expand_level( } } } + if work == previous_work { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + } } next } @@ -699,6 +706,30 @@ mod tests { assert!(!depth.contains_key(&4)); } + #[tokio::test] + async fn cancellation_without_edges_yields_while_scanning_frontier() { + let (out, _) = adjacencies(&[]).await; + let adjacency = Neighbors::directed(out); + let frontier = vec![99; 100_000]; + let mut depth = FastMap::default(); + let mut parents = FastMap::default(); + let mut next = Box::pin(expand_level( + &adjacency, + &frontier, + 1, + &mut depth, + &mut parents, + )); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(depth.is_empty() && parents.is_empty()); + assert!( + expand_level(&adjacency, &frontier, 1, &mut depth, &mut parents) + .await + .is_empty() + ); + } + #[tokio::test] async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { // A wide first layer makes the target's side expand. Two meeting nodes, diff --git a/docs/csr.md b/docs/csr.md index 4199320..82190e5 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -243,3 +243,5 @@ The 25% threshold is fixed: lookups pay a little for every pending delta edge, a | Very large edge tables and tight memory | Budget ~70 bytes per edge before turning it on | The [benchmarks](benchmarks.md) were all run with `use_csr=True`. + +Cancellation checkpoints count source and frontier rows that have no surviving edges, as well as deleted edge occurrences. Queries producing no rows therefore remain cooperative. From a641cf25cbe34db25245ce2543140af5b735d30e Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 23:56:10 -0700 Subject: [PATCH 9/9] perf(csr): align fixed-hop work quanta with output chunks Refs #19. Preserve bounded cancellation through rejected rows and tombstones while measuring the scheduling trade-off independently. --- crates/lance-graph/src/csr_expand.rs | 16 ++++++++-------- docs/csr.md | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 07903b5..9a2aa33 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -1231,7 +1231,7 @@ impl BatchExpansion { let required = self.bound.as_ref().map(|b| b.value(i)); for candidate in adjacency .candidates_from(key, &mut position) - .take(1024 - work) + .take(8192 - work) { work += 1; let Some(neighbor) = candidate else { @@ -1254,7 +1254,7 @@ impl BatchExpansion { break 'rows; } } - if work == 1024 { + if work == 8192 { break; } } @@ -1263,13 +1263,13 @@ impl BatchExpansion { } row += 1; position = NeighborPosition::default(); - if work == 1024 { + if work == 8192 { break; } } self.row = row; self.position = position; - work == 1024 + work == 8192 } } @@ -1985,7 +1985,7 @@ mod tests { let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); assert!(futures::poll!(next.as_mut()).is_pending()); drop(next); - assert!(expansion.row > 0 && expansion.row <= 1024); + assert!(expansion.row > 0 && expansion.row <= 8192); assert!(expansion .next_chunk(&adjacency, None, 1) .await @@ -2021,13 +2021,13 @@ mod tests { let mut next = Box::pin(expansion.next_chunk(&adj, None, 1)); assert!(futures::poll!(next.as_mut()).is_pending()); drop(next); - assert!(expansion.position.forward <= 1024); + assert!(expansion.position.forward <= 8192); assert!(expansion.position.forward > 0); } #[tokio::test] async fn expansion_yields_and_resumes_through_a_long_tombstone_run() { - let n = 5_000; + let n = 20_000; let base = Adjacency::build( edge_stream(vec![Some(1); n], (2..n as i64 + 2).map(Some).collect()), "src", @@ -2044,7 +2044,7 @@ mod tests { let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); assert!(futures::poll!(next.as_mut()).is_pending()); drop(next); - assert_eq!(expansion.position.forward, 1024); + assert_eq!(expansion.position.forward, 8192); assert_eq!( expansion .next_chunk(&adjacency, None, 1) diff --git a/docs/csr.md b/docs/csr.md index 82190e5..de29463 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -244,4 +244,4 @@ The 25% threshold is fixed: lookups pay a little for every pending delta edge, a The [benchmarks](benchmarks.md) were all run with `use_csr=True`. -Cancellation checkpoints count source and frontier rows that have no surviving edges, as well as deleted edge occurrences. Queries producing no rows therefore remain cooperative. +Cancellation checkpoints count source and frontier rows that have no surviving edges, as well as deleted edge occurrences. Queries producing no rows therefore remain cooperative. Fixed-hop expansion checks at most 8,192 raw candidates or missing source rows before yielding, matching the maximum output-chunk size; empty-frontier shortest-path checks retain their smaller bound.