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-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..08c16d7 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,22 @@ 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 = 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)) } else { @@ -65,7 +127,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 +146,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..240d206 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.as_ref().is_none_or(|result| result.is_err()) { + 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,97 @@ 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 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/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index b054f60..f0b9ad5 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", "time"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 201a3cc..9a2aa33 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"; @@ -215,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 @@ -223,21 +239,34 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + #[inline(always)] fn neighbors_from<'a>( &'a self, 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. + #[inline(always)] + 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); let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; - base.iter() - .chain(added) - .enumerate() - .skip(*offset) - .filter_map(move |(position, &neighbor)| { - *offset = position + 1; + let added_offset = offset.saturating_sub(base.len()); + base.get(*offset..) + .unwrap_or_default() + .iter() + .chain(added.get(added_offset..).unwrap_or_default()) + .map(move |&neighbor| { + let position = *offset; + *offset += 1; if removed.first() == Some(&position) { removed = &removed[1..]; None @@ -772,8 +801,8 @@ pub struct Neighbors { } /// Occurrence offsets are sufficient to pause even a high-degree source row. -#[derive(Default)] -struct NeighborPosition { +#[derive(Clone, Copy, Default)] +pub(crate) struct NeighborPosition { forward: usize, reverse: usize, } @@ -798,21 +827,32 @@ impl Neighbors { .for_each(f); } + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, position: &'a mut NeighborPosition, ) -> impl Iterator + 'a { - let forward = self.forward.neighbors_from(key, &mut position.forward); + self.candidates_from(key, position).flatten() + } + + /// 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, + 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)), ) } @@ -1155,42 +1195,81 @@ 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(); - while self.row < self.keys.len() { - let i = self.row; + 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; + '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)); - for neighbor in adjacency.neighbors_from(key, &mut self.position) { + for candidate in adjacency + .candidates_from(key, &mut position) + .take(8192 - work) + { + work += 1; + let Some(neighbor) = candidate else { + continue; + }; let neighbor = neighbor as i64; if required.is_some_and(|r| r != neighbor) { continue; } if let Some(target) = target { - // No target row: the inner join would drop this edge too. - let Some(position) = target.rows.position(neighbor) else { + 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 == 8192 { + break; + } + } + if work == previous_work { + work += 1; + } + row += 1; + position = NeighborPosition::default(); + if work == 8192 { + break; } - self.row += 1; - self.position = NeighborPosition::default(); } - chunk + self.row = row; + self.position = position; + work == 8192 } } @@ -1397,7 +1476,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; @@ -1752,6 +1833,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; @@ -1837,21 +1945,118 @@ 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 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 <= 8192); + 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( + 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 <= 8192); + assert!(expansion.position.forward > 0); + } + + #[tokio::test] + async fn expansion_yields_and_resumes_through_a_long_tombstone_run() { + let n = 20_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, 8192); + 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] @@ -1870,7 +2075,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 +2084,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 +2122,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/query.rs b/crates/lance-graph/src/query.rs index f617431..5f0b1b7 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(tokio::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(tokio::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(|| tokio::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..cb6b26c --- /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 tokio::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.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.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 9c37dec..305e782 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,18 @@ 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 previous_work = work; + let mut position = crate::csr_expand::NeighborPosition::default(); + for candidate in adjacency.candidates_from(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 +442,20 @@ fn expand_level( next.push(n); } } - }); + } + if work == previous_work { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + } } 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 +476,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 +485,8 @@ fn search_bidirectional( reverse_level, &mut reverse_depth, &mut reverse_parents, - ); + ) + .await; } let mut pending: Vec<_> = frontier .iter() @@ -480,7 +498,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 +522,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 +554,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 +609,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 +657,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 +681,57 @@ 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 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, // including parallel edges on both sides of the cut, yield six paths. let (out, inc) = adjacencies(&[ @@ -700,11 +746,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 +760,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 +779,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 +788,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 +814,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 +823,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..97eed26 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -68,6 +68,19 @@ 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, 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. + ## `GraphConfig` An immutable mapping from tables to a property graph. Create one with `GraphConfig.builder()`. @@ -128,12 +141,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 +191,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/csr.md b/docs/csr.md index 0d3ae5a..de29463 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. @@ -241,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. 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. diff --git a/docs/cypher.md b/docs/cypher.md index e14dee7..da3e389 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -124,6 +124,23 @@ 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. 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)` 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..9717815 --- /dev/null +++ b/python/python/tests/test_timeout.py @@ -0,0 +1,92 @@ +# 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_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"): + 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]