From 99353d49a2194dde2a4956e61a2768f123e99b8d Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:33:22 +0000 Subject: [PATCH 1/2] fix: spawn legacy fragment opens during scans --- rust/lance/src/io/exec/scan.rs | 38 ++++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/io/exec/scan.rs b/rust/lance/src/io/exec/scan.rs index a8c3b2e3dc3..1a2fa21c5b1 100644 --- a/rust/lance/src/io/exec/scan.rs +++ b/rust/lance/src/io/exec/scan.rs @@ -66,6 +66,18 @@ async fn open_file( Ok(reader) } +fn spawn_io_task(task: F) -> impl std::future::Future> +where + F: std::future::Future> + Send + 'static, + T: Send + 'static, +{ + // Buffered streams may stop polling one I/O future while another future + // that would unblock it is waiting on the same connection pool. + tokio::spawn(task.in_current_span()).map(|task_result| { + task_result.map_err(|error| DataFusionError::External(Box::new(error)))? + }) +} + struct FragmentWithRange { fragment: FileFragment, range: Option>, @@ -444,7 +456,7 @@ impl LanceStream { let batches = if config.ordered_output { let readers = stream::iter(file_fragments) .map(move |file_fragment| { - Ok(open_file( + Ok(spawn_io_task(open_file( file_fragment, project_schema.clone(), FragReadConfig::default() @@ -456,7 +468,7 @@ impl LanceStream { .with_row_created_at_version(config.with_row_created_at_version), config.with_make_deletions_null, None, - )) + ))) }) .try_buffered(fragment_readahead); let tasks = readers.and_then(move |reader| async move { @@ -476,7 +488,7 @@ impl LanceStream { } else { let readers = stream::iter(file_fragments) .map(move |file_fragment| { - Ok(open_file( + Ok(spawn_io_task(open_file( file_fragment, project_schema.clone(), FragReadConfig::default() @@ -488,7 +500,7 @@ impl LanceStream { .with_row_created_at_version(config.with_row_created_at_version), config.with_make_deletions_null, None, - )) + ))) }) .try_buffered(fragment_readahead); let tasks = readers.and_then(move |reader| async move { @@ -844,6 +856,24 @@ mod tests { scan.execute(0, Arc::new(TaskContext::default())).unwrap(); } + #[tokio::test] + async fn spawned_io_task_progresses_without_outer_poll() { + let (ran_tx, ran_rx) = tokio::sync::oneshot::channel(); + let task = spawn_io_task(async move { + ran_tx.send(()).unwrap(); + Result::Ok(()) + }); + + tokio::time::timeout(std::time::Duration::from_secs(5), ran_rx) + .await + .expect("spawned I/O task did not make independent progress") + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), task) + .await + .expect("spawned I/O task did not finish") + .unwrap(); + } + /// Verify that executing with target_partitions=1 produces the same row count as the /// default context. Regression guard for the parallelism cap. #[tokio::test] From 7c683f406ca43c4743ce3f5c22b38728c11e9ada Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:08:27 +0000 Subject: [PATCH 2/2] fix: make legacy fragment opens cancellation safe --- rust/lance/src/io/exec/pushdown_scan.rs | 288 ++++++++++++++++++++++-- rust/lance/src/io/exec/scan.rs | 60 ++--- rust/lance/src/io/exec/utils.rs | 26 ++- 3 files changed, 313 insertions(+), 61 deletions(-) diff --git a/rust/lance/src/io/exec/pushdown_scan.rs b/rust/lance/src/io/exec/pushdown_scan.rs index b82434116b8..5992eb2b28f 100644 --- a/rust/lance/src/io/exec/pushdown_scan.rs +++ b/rust/lance/src/io/exec/pushdown_scan.rs @@ -47,7 +47,7 @@ use crate::{ }; use super::Planner; -use super::utils::InstrumentedRecordBatchStreamAdapter; +use super::utils::{InstrumentedRecordBatchStreamAdapter, buffered_fragment_opens}; #[derive(Debug, Clone)] pub struct ScanConfig { @@ -204,23 +204,24 @@ impl ExecutionPlan for LancePushdownScanExec { } }); - let batch_stream = fragment_stream.map(|(exec, fragment)| async move { - let frag_scanner = FragmentScanner::open( - fragment, - exec.dataset, - exec.projection, - exec.predicate_projection, - exec.predicate, - exec.config.clone(), - ) - .await?; - - frag_scanner.scan() - }); + let batch_stream = buffered_fragment_opens( + fragment_stream, + self.config.fragment_readahead, + |(exec, fragment)| async move { + let frag_scanner = FragmentScanner::open( + fragment, + exec.dataset, + exec.projection, + exec.predicate_projection, + exec.predicate, + exec.config.clone(), + ) + .await?; - let batch_stream = batch_stream - .buffered(self.config.fragment_readahead) - .try_flatten(); + frag_scanner.scan() + }, + ) + .try_flatten(); Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( self.schema(), @@ -697,6 +698,11 @@ impl FragmentScanner { #[cfg(test)] mod test { + use std::collections::HashSet; + use std::fmt::Display; + use std::sync::Mutex; + use std::time::Duration; + use arrow_array::{ ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int32Array, RecordBatchIterator, StringArray, StructArray, TimestampMicrosecondArray, UInt64Array, @@ -705,17 +711,265 @@ mod test { use arrow_ord::sort::sort_to_indices; use arrow_schema::{Field, TimeUnit}; use arrow_select::concat::concat_batches; + use async_trait::async_trait; use datafusion::prelude::{Column, SessionContext, lit}; + use futures::stream::BoxStream; use lance_arrow::{FixedSizeListArrayExt, SchemaExt}; use lance_core::utils::tempfile::TempStrDir; + use lance_datagen::{array, gen_batch}; use lance_file::version::LanceFileVersion; + use lance_io::object_store::WrappingObjectStore; + use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, + PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Result as ObjectStoreResult, path::Path, + }; use pretty_assertions::assert_eq; + use tokio::sync::{Semaphore, mpsc}; use crate::dataset::WriteParams; + use crate::io::exec::{LanceScanConfig, LanceScanExec}; + use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount}; use lance_datafusion::logical_expr::ExprExt; use super::*; + #[derive(Debug)] + struct BlockingDataFileReads { + release: Semaphore, + started_paths: Mutex>, + started_tx: mpsc::UnboundedSender<()>, + completed_tx: mpsc::UnboundedSender<()>, + } + + impl BlockingDataFileReads { + async fn wait_for_release(&self, location: &Path) -> bool { + let is_first_data_file_read = location.as_ref().ends_with(".lance") + && self.started_paths.lock().unwrap().insert(location.clone()); + if !is_first_data_file_read { + return false; + } + + self.started_tx.send(()).unwrap(); + self.release + .acquire() + .await + .expect("release semaphore was closed") + .forget(); + true + } + } + + #[derive(Debug, Clone)] + struct BlockingDataFileStoreWrapper { + reads: Arc, + } + + impl WrappingObjectStore for BlockingDataFileStoreWrapper { + fn wrap(&self, _prefix: &str, target: Arc) -> Arc { + Arc::new(BlockingDataFileStore { + target, + reads: self.reads.clone(), + }) + } + } + + #[derive(Debug)] + struct BlockingDataFileStore { + target: Arc, + reads: Arc, + } + + impl Display for BlockingDataFileStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BlockingDataFileStore({})", self.target) + } + } + + #[async_trait] + impl ObjectStore for BlockingDataFileStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + options: PutOptions, + ) -> ObjectStoreResult { + self.target.put_opts(location, payload, options).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + options: PutMultipartOptions, + ) -> ObjectStoreResult> { + self.target.put_multipart_opts(location, options).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> ObjectStoreResult { + let is_tracked_read = self.reads.wait_for_release(location).await; + let result = self.target.get_opts(location, options).await; + if is_tracked_read { + self.reads.completed_tx.send(()).unwrap(); + } + result + } + + fn delete_stream( + &self, + locations: BoxStream<'static, ObjectStoreResult>, + ) -> BoxStream<'static, ObjectStoreResult> { + self.target.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, ObjectStoreResult> { + self.target.list(prefix) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> ObjectStoreResult { + self.target.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> ObjectStoreResult<()> { + self.target.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> ObjectStoreResult<()> { + self.target.rename_opts(from, to, options).await + } + } + + #[derive(Debug)] + enum LegacyScanPath { + Regular, + Pushdown, + } + + #[tokio::test] + #[rstest::rstest] + #[case::regular(LegacyScanPath::Regular)] + #[case::pushdown(LegacyScanPath::Pushdown)] + async fn test_fragment_opens_progress_with_bounded_cancellation( + #[case] scan_path: LegacyScanPath, + ) { + const FRAGMENT_READAHEAD: usize = 2; + const ROWS_PER_FRAGMENT: u32 = 8; + + let dataset = gen_batch() + .col("x", array::step::()) + .into_ram_dataset_with_params( + FragmentCount::from(4), + FragmentRowCount::from(ROWS_PER_FRAGMENT), + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::Legacy), + max_rows_per_file: ROWS_PER_FRAGMENT as usize, + max_rows_per_group: ROWS_PER_FRAGMENT as usize, + ..WriteParams::default() + }), + ) + .await + .unwrap(); + dataset.session().file_metadata_cache().clear().await; + + let (started_tx, mut started_rx) = mpsc::unbounded_channel(); + let (completed_tx, mut completed_rx) = mpsc::unbounded_channel(); + let reads = Arc::new(BlockingDataFileReads { + release: Semaphore::new(0), + started_paths: Mutex::new(HashSet::new()), + started_tx, + completed_tx, + }); + let dataset = Arc::new(dataset.with_object_store_wrappers([Arc::new( + BlockingDataFileStoreWrapper { + reads: reads.clone(), + }, + ) + as Arc])); + let fragments = dataset.fragments().clone(); + let projection = Arc::new(dataset.schema().clone()); + + let exec: Arc = match scan_path { + LegacyScanPath::Regular => Arc::new(LanceScanExec::new( + dataset, + fragments, + None, + projection, + LanceScanConfig { + fragment_readahead: Some(FRAGMENT_READAHEAD), + ordered_output: true, + ..LanceScanConfig::default() + }, + )), + LegacyScanPath::Pushdown => Arc::new( + LancePushdownScanExec::try_new( + dataset, + fragments, + projection, + col("x").gt(lit(-1)), + ScanConfig { + fragment_readahead: FRAGMENT_READAHEAD, + ..ScanConfig::default() + }, + ) + .unwrap(), + ), + }; + let context = SessionContext::new(); + let mut output = exec.execute(0, context.task_ctx()).unwrap(); + + assert!(futures::poll!(output.next()).is_pending()); + for _ in 0..FRAGMENT_READAHEAD { + tokio::time::timeout(Duration::from_secs(5), started_rx.recv()) + .await + .expect("fragment open did not start") + .expect("fragment open start channel closed"); + } + tokio::task::yield_now().await; + assert!( + matches!(started_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)), + "fragment opens exceeded fragment_readahead" + ); + + // The child task must finish this request even though the output stream + // is not polled again. + reads.release.add_permits(1); + tokio::time::timeout(Duration::from_secs(5), completed_rx.recv()) + .await + .expect("fragment open did not progress independently") + .expect("fragment open completion channel closed"); + assert!( + matches!(started_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)), + "fragment opens exceeded fragment_readahead" + ); + + // Dropping the scan must abort the other in-flight open instead of + // detaching it from the cancelled query. + drop(output); + reads.release.add_permits(FRAGMENT_READAHEAD); + assert!( + tokio::time::timeout(Duration::from_secs(1), completed_rx.recv()) + .await + .is_err(), + "fragment open completed after scan cancellation" + ); + } + // TODO: test pushdown with nested column once https://github.com/apache/arrow-datafusion/pull/8256 // is released. diff --git a/rust/lance/src/io/exec/scan.rs b/rust/lance/src/io/exec/scan.rs index 1a2fa21c5b1..905658aed6a 100644 --- a/rust/lance/src/io/exec/scan.rs +++ b/rust/lance/src/io/exec/scan.rs @@ -43,7 +43,7 @@ use crate::dataset::scanner::{ }; use crate::datatypes::Schema; -use super::utils::IoMetrics; +use super::utils::{IoMetrics, buffered_fragment_opens}; async fn open_file( file_fragment: FileFragment, @@ -66,18 +66,6 @@ async fn open_file( Ok(reader) } -fn spawn_io_task(task: F) -> impl std::future::Future> -where - F: std::future::Future> + Send + 'static, - T: Send + 'static, -{ - // Buffered streams may stop polling one I/O future while another future - // that would unblock it is waiting on the same connection pool. - tokio::spawn(task.in_current_span()).map(|task_result| { - task_result.map_err(|error| DataFusionError::External(Box::new(error)))? - }) -} - struct FragmentWithRange { fragment: FileFragment, range: Option>, @@ -454,9 +442,11 @@ impl LanceStream { .collect::>(); let batches = if config.ordered_output { - let readers = stream::iter(file_fragments) - .map(move |file_fragment| { - Ok(spawn_io_task(open_file( + let readers = buffered_fragment_opens( + stream::iter(file_fragments), + fragment_readahead, + move |file_fragment| { + open_file( file_fragment, project_schema.clone(), FragReadConfig::default() @@ -468,9 +458,9 @@ impl LanceStream { .with_row_created_at_version(config.with_row_created_at_version), config.with_make_deletions_null, None, - ))) - }) - .try_buffered(fragment_readahead); + ) + }, + ); let tasks = readers.and_then(move |reader| async move { reader .read_all(config.batch_size as u32) @@ -486,9 +476,11 @@ impl LanceStream { .stream_in_current_span() .boxed() } else { - let readers = stream::iter(file_fragments) - .map(move |file_fragment| { - Ok(spawn_io_task(open_file( + let readers = buffered_fragment_opens( + stream::iter(file_fragments), + fragment_readahead, + move |file_fragment| { + open_file( file_fragment, project_schema.clone(), FragReadConfig::default() @@ -500,9 +492,9 @@ impl LanceStream { .with_row_created_at_version(config.with_row_created_at_version), config.with_make_deletions_null, None, - ))) - }) - .try_buffered(fragment_readahead); + ) + }, + ); let tasks = readers.and_then(move |reader| async move { reader .read_all(config.batch_size as u32) @@ -856,24 +848,6 @@ mod tests { scan.execute(0, Arc::new(TaskContext::default())).unwrap(); } - #[tokio::test] - async fn spawned_io_task_progresses_without_outer_poll() { - let (ran_tx, ran_rx) = tokio::sync::oneshot::channel(); - let task = spawn_io_task(async move { - ran_tx.send(()).unwrap(); - Result::Ok(()) - }); - - tokio::time::timeout(std::time::Duration::from_secs(5), ran_rx) - .await - .expect("spawned I/O task did not make independent progress") - .unwrap(); - tokio::time::timeout(std::time::Duration::from_secs(5), task) - .await - .expect("spawned I/O task did not finish") - .unwrap(); - } - /// Verify that executing with target_partitions=1 produces the same row count as the /// default context. Regression guard for the parallelism cap. #[tokio::test] diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index 44092f75459..fc77471afdd 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -18,6 +18,7 @@ use std::task::{Context, Poll}; use arrow_array::{RecordBatch, UInt64Array}; use arrow_schema::SchemaRef; use async_trait::async_trait; +use datafusion::common::runtime::SpawnedTask; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricValue, @@ -26,16 +27,39 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, }; use futures::stream::FuturesUnordered; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; use lance_core::error::{CloneableResult, Error}; use lance_core::utils::futures::{Capacity, SharedStreamExt}; use lance_core::{ROW_ID, Result}; use lance_index::prefilter::FilterLoader; use lance_select::{RowAddrMask, RowAddrTreeMap, result::IndexExprResult}; +use tracing::Instrument; use crate::Dataset; use crate::index::prefilter::DatasetPreFilter; +/// Open fragments on cancellation-safe tasks while preserving the stream's +/// ordering and readahead bound. +pub(crate) fn buffered_fragment_opens( + fragments: S, + fragment_readahead: usize, + mut open: Open, +) -> impl Stream> +where + S: Stream + Send, + Open: FnMut(S::Item) -> OpenFuture + Send, + OpenFuture: Future> + Send + 'static, + Reader: Send + 'static, +{ + fragments + .map(move |fragment| { + SpawnedTask::spawn(open(fragment).in_current_span()).map(|task_result| { + task_result.map_err(|error| DataFusionError::External(Box::new(error)))? + }) + }) + .buffered(fragment_readahead) +} + #[derive(Debug, Clone)] pub enum PreFilterSource { /// The prefilter input is an array of row ids that match the filter condition