diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index ebe6ff2e96a..049837b56df 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use anyhow::Result; use async_trait::async_trait; +use futures::StreamExt; use futures::TryStreamExt; use glob::glob; use vortex::array::ArrayRef; @@ -78,15 +79,17 @@ impl Dataset for TPCHLCommentChunked { let file_chunks: Vec<_> = file .scan()? .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable)) + .into_stream()? .map({ let ctx = ctx.clone(); - move |a| { + move |result| { let mut ctx = ctx.clone(); - let canonical = a.execute::(&mut ctx)?; - Ok(canonical.into_array()) + result.and_then(|a| { + let canonical = a.execute::(&mut ctx)?; + Ok(canonical.into_array()) + }) } }) - .into_array_stream()? .try_collect() .await?; chunks.extend(file_chunks); diff --git a/vortex-datafusion/src/persistent/access_plan.rs b/vortex-datafusion/src/persistent/access_plan.rs index ad7bf941a2c..0e14b6730b0 100644 --- a/vortex-datafusion/src/persistent/access_plan.rs +++ b/vortex-datafusion/src/persistent/access_plan.rs @@ -58,10 +58,7 @@ impl VortexAccessPlan { /// /// This is used internally by the file opener after it has translated a /// `PartitionedFile` into a Vortex scan. - pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder - where - A: 'static + Send, - { + pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder { let Self { selection } = self; if let Some(selection) = selection { diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..f142bf486cd 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -420,9 +420,7 @@ impl FileOpener for VortexOpener { }) .transpose()?; - if let Some(limit) = limit - && filter.is_none() - { + if let Some(limit) = limit { scan_builder = scan_builder.with_limit(limit); } @@ -431,29 +429,30 @@ impl FileOpener for VortexOpener { } let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + let file_location = file.object_meta.location.clone(); let stream = scan_builder .with_metrics_registry(metrics_registry) .with_projection(scan_projection) .with_some_filter(filter) .with_ordered(has_output_ordering) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) + // Convert to Arrow inline on the polling thread: DataFusion sources are expected + // to do their CPU work inside `poll_next`, and spawning this onto the blocking + // pool oversubscribes the CPU. + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) }) + .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) .map(move |batch| { let batch = if projector.projection().as_ref().is_empty() { batch @@ -572,6 +571,12 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") } +fn vortex_file_read_error(path: &Path, error: VortexError) -> DataFusionError { + DataFusionError::External(Box::new( + error.with_context(format!("Failed to read Vortex file: {path}")), + )) +} + #[cfg(test)] mod tests { use std::fmt; @@ -604,6 +609,7 @@ mod tests { use datafusion_physical_expr::expressions as df_expr; use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::projection::ProjectionExpr; + use futures::TryStreamExt; use insta::assert_snapshot; use itertools::Itertools; use object_store::ObjectStore; @@ -988,6 +994,48 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_applies_limit_after_filtering() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "filtered-limit/file.vortex"; + let batch = record_batch!(( + "a", + Int32, + vec![Some(1), Some(2), Some(3), Some(4), Some(5), Some(6)] + )) + .unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let file = PartitionedFile::new(file_path.to_string(), data_size); + let table_schema = TableSchema::from_file_schema(batch.schema()); + // `a > 3` excludes the first three rows, so a limit applied *before* filtering would take + // rows [1, 2, 3] and filter them all out (yielding nothing), whereas a limit applied + // *after* filtering yields the first three matching rows [4, 5, 6]. Asserting the values + // (not just the count) is what makes this test able to detect a pre-filter regression. + let filter = logical2physical(&col("a").gt(lit(3_i32)), table_schema.table_schema()); + + let mut opener = make_opener(object_store, table_schema, Some(filter)); + opener.limit = Some(3); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("projected column should be Int32") + .values() + .to_vec() + }) + .collect::>(); + + assert_eq!(values, [4, 5, 6]); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; diff --git a/vortex-datafusion/src/persistent/stream.rs b/vortex-datafusion/src/persistent/stream.rs index af2fbc8693e..9038ba81db8 100644 --- a/vortex-datafusion/src/persistent/stream.rs +++ b/vortex-datafusion/src/persistent/stream.rs @@ -17,14 +17,14 @@ use futures::stream::BoxStream; /// [`PartitionedFile`]: datafusion_datasource::PartitionedFile pub(crate) struct PrunableStream { file_pruner: FilePruner, - stream: BoxStream<'static, DFResult>, + stream: Option>>, } impl PrunableStream { pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult>) -> Self { Self { file_pruner, - stream, + stream: Some(stream), } } } @@ -34,9 +34,13 @@ impl Stream for PrunableStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.as_mut().file_pruner.should_prune()? { + self.stream.take(); Poll::Ready(None) } else { - self.stream.poll_next_unpin(cx) + match self.stream.as_mut() { + Some(stream) => stream.poll_next_unpin(cx), + None => Poll::Ready(None), + } } } } diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 325f8eae92f..c802a1c0502 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -97,12 +97,19 @@ use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; use futures::TryStreamExt; use futures::future::try_join_all; +use futures::stream::BoxStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; +use vortex::array::stream::SendableArrayStream; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::Nullability; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_err; use vortex::expr::Expression; use vortex::expr::and as vx_and; use vortex::expr::get_item; @@ -110,6 +117,8 @@ use vortex::expr::pack; use vortex::expr::root; use vortex::expr::stats::Precision; use vortex::expr::transform::replace; +use vortex::io::runtime::Handle; +use vortex::io::runtime::JoinOutcome; use vortex::io::session::RuntimeSessionExt; use vortex::scan::DataSourceRef; use vortex::scan::ScanRequest; @@ -402,6 +411,7 @@ impl DataSource for VortexDataSource { )); let session = self.session.clone(); let num_partitions = self.num_partitions; + let ordered = self.ordered; // Pre-build the leftover projector (if any) so we can apply it after batch conversion. let leftover_projector = self @@ -418,17 +428,18 @@ impl DataSource for VortexDataSource { .map_err(|e| DataFusionError::External(Box::new(e)))?; // Each split.execute() returns a lazy stream whose early polls do preparation - // work (expression resolution, layout traversal, first I/O spawns). We use - // try_flatten_unordered to poll multiple split streams concurrently so that - // the next split is already warm when the current one finishes. + // work (expression resolution, layout traversal, first I/O spawns). Both ordering + // modes flatten with cross-partition I/O look-ahead; the ordered path additionally + // preserves partition order so an ordered global limit cannot observe later rows first. let scan_streams = scan.partitions().map(|split_result| { let split = split_result?; split.execute() }); let handle = session.handle(); - let stream = scan_streams - .try_flatten_unordered(Some(num_partitions * 2)) + let chunks = + flatten_scan_streams(scan_streams, ordered, num_partitions * 2, handle.clone()); + let stream = chunks .map(move |result| { let session = session.clone(); let target_field = Arc::clone(&projected_target_field); @@ -654,6 +665,85 @@ impl DataSource for VortexDataSource { } } +fn flatten_scan_streams( + scan_streams: S, + ordered: bool, + concurrency: usize, + handle: Handle, +) -> BoxStream<'static, VortexResult> +where + S: futures::Stream> + Send + 'static, +{ + if !ordered { + return scan_streams + .try_flatten_unordered(Some(concurrency)) + .boxed(); + } + + // Ordered output must preserve global row order, but later partitions should still warm up + // (start their I/O) while an earlier one is draining. `buffered` invokes this map while + // filling its window, so creating the channel and spawning the drain here starts each window + // member before `try_flatten` begins draining the first receiver. Each drain has a bounded + // channel for back-pressure. It can hold `CHUNK_BUFFER_CAPACITY` chunks plus one pending + // send, and at most `lookahead` drains are live at once. + const CHUNK_BUFFER_CAPACITY: usize = 2; + let lookahead = concurrency.max(1); + scan_streams + .map(move |stream_result| { + let handle = handle.clone(); + let receiver = stream_result.map(|mut stream| { + let (tx, rx) = mpsc::channel(CHUNK_BUFFER_CAPACITY); + // Keep the drain task's handle (do not `detach` it) so that a panic inside the + // source stream is surfaced as a terminal error. A detached task's panic would be + // discarded; `tx` would drop, the receiver would end cleanly, and `try_flatten` + // would treat the truncated partition as a normal completion — silently returning + // incomplete results as success. + let drain = handle.spawn(async move { + loop { + tokio::select! { + _ = tx.closed() => break, + item = stream.next() => { + let Some(item) = item else { + break; + }; + if tx.send(item).await.is_err() { + break; + } + } + } + } + }); + // Once the receiver is drained, join the drain task and turn a panic into a + // terminal error item. A clean completion or a benign runtime abort ends the + // partition without appending an extra item. + let finalizer = futures::stream::once(async move { + let mut drain = drain; + match futures::future::poll_fn(|cx| drain.poll_join(cx)).await { + JoinOutcome::Completed(()) | JoinOutcome::Aborted => None, + JoinOutcome::Panicked(payload) => Some(Err(panic_to_vortex_err(payload))), + } + }) + .filter_map(|item| async move { item }); + ReceiverStream::new(rx).chain(finalizer) + }); + futures::future::ready(receiver) + }) + .buffered(lookahead) + .try_flatten() + .boxed() +} + +/// Convert a caught panic payload from an ordered drain task into a terminal [`VortexError`], so +/// the scan surfaces the failure instead of silently truncating its output. +fn panic_to_vortex_err(payload: Box) -> VortexError { + let message = payload + .downcast_ref::<&'static str>() + .map(|s| (*s).to_string()) + .or_else(|| payload.downcast_ref::().cloned()) + .unwrap_or_else(|| "unknown panic".to_string()); + vortex_err!("ordered scan partition drain task panicked: {message}") +} + /// Convert a Vortex [`Option`] to a DataFusion /// [`DataFusionPrecision`]. /// @@ -665,3 +755,202 @@ fn estimate_to_df_precision(est: &Precision) -> DFPrecision { Precision::Absent => DFPrecision::Absent, } } + +#[cfg(test)] +mod tests { + use std::pin::Pin; + use std::sync::Arc; + use std::task::Context; + use std::task::Poll; + use std::time::Duration; + + use futures::Stream; + use futures::TryStreamExt; + use tokio::sync::Notify; + use tokio::sync::oneshot; + use vortex::array::ArrayRef; + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::array_session; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::stream::ArrayStreamAdapter; + use vortex::array::stream::ArrayStreamExt; + use vortex::dtype::DType; + use vortex::dtype::Nullability; + use vortex::dtype::PType; + use vortex::error::VortexError; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + use vortex::io::runtime::tokio::TokioRuntime; + + use super::flatten_scan_streams; + + fn i32_stream(dtype: &DType, chunks: Vec>) -> super::SendableArrayStream { + let items = chunks + .into_iter() + .map(|values| Ok(PrimitiveArray::from_iter(values).into_array())) + .collect::>>(); + ArrayStreamAdapter::new(dtype.clone(), futures::stream::iter(items)).boxed() + } + + fn collect_i32(chunks: Vec) -> VortexResult> { + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for chunk in chunks { + let primitive = chunk.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + Ok(values) + } + + struct DropNotifier(Option>); + + impl Drop for DropNotifier { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + struct DropNotifyingPendingStream { + started: Arc, + _drop_notifier: DropNotifier, + } + + impl Stream for DropNotifyingPendingStream { + type Item = VortexResult; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.started.notify_one(); + Poll::Pending + } + } + + #[tokio::test] + async fn ordered_flatten_preserves_partition_order() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let first = i32_stream(&dtype, vec![vec![0, 1], vec![2]]); + let second = i32_stream(&dtype, vec![vec![10, 11]]); + + let streams = + futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); + let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let chunks = flattened.try_collect::>().await?; + + assert_eq!(collect_i32(chunks)?, [0, 1, 2, 10, 11]); + Ok(()) + } + + #[tokio::test] + async fn ordered_flatten_warms_later_partition_while_first_pending() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let first_started = Arc::new(Notify::new()); + let first_started_for_stream = Arc::clone(&first_started); + let release_first = Arc::new(Notify::new()); + let release_first_for_stream = Arc::clone(&release_first); + let first = ArrayStreamAdapter::new( + dtype.clone(), + futures::stream::once(async move { + first_started_for_stream.notify_one(); + release_first_for_stream.notified().await; + Ok(PrimitiveArray::from_iter([0i32]).into_array()) + }), + ) + .boxed(); + + let first_started_for_second = Arc::clone(&first_started); + let release_first_for_second = Arc::clone(&release_first); + let second = ArrayStreamAdapter::new( + dtype, + futures::stream::once(async move { + first_started_for_second.notified().await; + release_first_for_second.notify_one(); + Ok(PrimitiveArray::from_iter([100i32]).into_array()) + }), + ) + .boxed(); + + let streams = + futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); + let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let chunks = + tokio::time::timeout(Duration::from_secs(1), flattened.try_collect::>()) + .await + .map_err(|_| vortex_err!("ordered flatten did not warm the later partition"))??; + + assert_eq!(collect_i32(chunks)?, [0, 100]); + Ok(()) + } + + struct PanickingStream; + + impl Stream for PanickingStream { + type Item = VortexResult; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + panic!("boom") + } + } + + #[tokio::test] + async fn ordered_flatten_surfaces_a_partition_panic_as_an_error() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let panicking = ArrayStreamAdapter::new(dtype, PanickingStream).boxed(); + + let streams = futures::stream::iter([Ok::<_, VortexError>(panicking)]); + let flattened = flatten_scan_streams(streams, true, 1, TokioRuntime::current()); + let result = flattened.try_collect::>().await; + + assert!( + result.is_err(), + "a panicking ordered partition must surface as an error, not silent truncation" + ); + Ok(()) + } + + #[tokio::test] + async fn ordered_flatten_cancels_all_warmed_drains_when_dropped() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let started = Arc::new(Notify::new()); + let mut dropped_recvs = Vec::new(); + let mut streams = Vec::new(); + for _ in 0..3 { + let (dropped_send, dropped_recv) = oneshot::channel(); + dropped_recvs.push(dropped_recv); + let pending = ArrayStreamAdapter::new( + dtype.clone(), + DropNotifyingPendingStream { + started: Arc::clone(&started), + _drop_notifier: DropNotifier(Some(dropped_send)), + }, + ) + .boxed(); + streams.push(Ok::<_, VortexError>(pending)); + } + + // A concurrency of 4 yields a lookahead of 4, so all three partitions warm (spawn their + // drains) concurrently before the first one is drained. + let streams = futures::stream::iter(streams); + let mut flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let mut next = Box::pin(futures::StreamExt::next(&mut flattened)); + tokio::select! { + _ = started.notified() => {} + _ = &mut next => return Err(vortex_err!("pending drain unexpectedly produced a chunk")), + } + drop(next); + drop(flattened); + + for dropped_recv in dropped_recvs { + let dropped = tokio::time::timeout(Duration::from_secs(1), dropped_recv) + .await + .map_err(|_| { + vortex_err!("dropping the flattened stream did not cancel every warmed drain") + })?; + dropped + .map_err(|_| vortex_err!("pending source dropped without notifying the test"))?; + } + Ok(()) + } +} diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index c9a71f85c55..94a3b70494a 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -11,7 +11,6 @@ use std::sync::Arc; use std::sync::OnceLock; use itertools::Itertools; -use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::Expression; @@ -175,7 +174,7 @@ impl VortexFile { /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and /// execution options. - pub fn scan(&self) -> VortexResult> { + pub fn scan(&self) -> VortexResult { Ok(ScanBuilder::new( self.session.clone(), self.layout_reader()?, diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..b7550949c8c 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1278,7 +1278,7 @@ async fn write_nullable_top_level_struct() { async fn round_trip( array: &ArrayRef, - f: impl Fn(ScanBuilder) -> VortexResult>, + f: impl Fn(ScanBuilder) -> VortexResult, ) -> VortexResult { let mut writer = vec![]; SESSION diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 663b29d9501..223c5962441 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -10,6 +10,7 @@ use arrow_schema::ArrowError; use arrow_schema::Field; use arrow_schema::SchemaRef; use futures::Stream; +use futures::StreamExt; use futures::TryStreamExt; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -17,10 +18,12 @@ use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::session::RuntimeSessionExt; +use vortex_utils::parallelism::get_available_parallelism; use crate::scan::scan_builder::ScanBuilder; -impl ScanBuilder { +impl ScanBuilder { /// Creates a new `RecordBatchReader` from the scan builder. /// /// The `schema` parameter is used to define the schema of the resulting record batches. In @@ -35,11 +38,11 @@ impl ScanBuilder { let session = self.session().clone(); let iter = self + .into_iter(runtime)? .map(move |chunk| { let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) + chunk.and_then(|chunk| to_record_batch(chunk, &struct_field, &mut ctx)) }) - .into_iter(runtime)? .map(|result| result.map_err(|e| ArrowError::ExternalError(Box::new(e)))); Ok(RecordBatchIteratorAdapter { iter, schema }) @@ -49,15 +52,29 @@ impl ScanBuilder { self, schema: SchemaRef, ) -> VortexResult> + Send + 'static> { - let struct_field = Field::new_struct("", schema.fields().clone(), false); + let struct_field = Arc::new(Field::new_struct("", schema.fields().clone(), false)); let session = self.session().clone(); + let handle = session.handle(); + let concurrency = get_available_parallelism().unwrap_or(1); let stream = self + .into_stream()? .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) + let session = session.clone(); + let handle = handle.clone(); + let struct_field = Arc::clone(&struct_field); + async move { + handle + .spawn_blocking(move || { + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + to_record_batch(chunk, struct_field.as_ref(), &mut ctx) + }) + }) + .await + } }) - .into_stream()? + .buffered(concurrency) .map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(stream) diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 0b1f150c67e..8d0feb95e37 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -40,6 +40,7 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; use crate::LayoutReaderRef; +use crate::scan::limit::RowLimit; use crate::scan::scan_builder::ScanBuilder; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. @@ -157,6 +158,19 @@ impl DataSource for LayoutReaderDataSource { } } + // An ordered limit must see earlier filtered rows before later external partitions are + // allowed to reserve any budget. Emit one partition for that path; its inner scan keeps + // the limit at mask level. Unordered partitions share a completion-order budget instead. + let ordered_limit = scan_request.ordered && scan_request.limit.is_some(); + let row_limit = (!scan_request.ordered) + .then(|| scan_request.limit.map(RowLimit::new)) + .flatten(); + let split_size = if ordered_limit { + row_range.end - row_range.start + } else { + self.split_max_row_count + }; + Ok(Box::new(LayoutReaderScan { reader: Arc::clone(&self.reader), session: self.session.clone(), @@ -164,12 +178,13 @@ impl DataSource for LayoutReaderDataSource { projection: scan_request.projection, filter: scan_request.filter, limit: scan_request.limit, + row_limit, selection: scan_request.selection, ordered: scan_request.ordered, metrics_registry: self.metrics_registry.clone(), next_row: row_range.start, end_row: row_range.end, - split_size: self.split_max_row_count, + split_size, })) } @@ -185,6 +200,7 @@ struct LayoutReaderScan { projection: Expression, filter: Option, limit: Option, + row_limit: Option, ordered: bool, selection: Selection, metrics_registry: Option>, @@ -225,32 +241,22 @@ impl Stream for LayoutReaderScan { if this.limit.is_some_and(|limit| limit == 0) { return Poll::Ready(None); } + if this.row_limit.as_ref().is_some_and(RowLimit::is_exhausted) { + return Poll::Ready(None); + } let split_end = this .next_row .saturating_add(this.split_size) .min(this.end_row); let row_range = this.next_row..split_end; - let split_rows = split_end - this.next_row; - - let split_limit = this.limit; - // Only decrement the remaining limit when there is no filter. With a filter, - // the actual output row count is unknown (could be anywhere from 0 to split_rows), - // so decrementing by split_rows would be too aggressive and could stop producing - // splits before the limit is reached. Instead, pass the full remaining limit to - // each split and let the engine enforce the exact limit at the stream level. - if this.filter.is_none() - && let Some(ref mut limit) = this.limit - { - *limit = limit.saturating_sub(split_rows); - } - let split = Box::new(LayoutReaderSplit { reader: Arc::clone(&this.reader), session: this.session.clone(), projection: this.projection.clone(), filter: this.filter.clone(), - limit: split_limit, + limit: this.limit, + row_limit: this.row_limit.clone(), ordered: this.ordered, row_range, selection: this.selection.clone(), @@ -263,7 +269,10 @@ impl Stream for LayoutReaderScan { } fn size_hint(&self) -> (usize, Option) { - if self.next_row >= self.end_row { + if self.next_row >= self.end_row + || self.limit.is_some_and(|limit| limit == 0) + || self.row_limit.as_ref().is_some_and(RowLimit::is_exhausted) + { return (0, Some(0)); } let remaining_rows = self.end_row - self.next_row; @@ -278,6 +287,7 @@ struct LayoutReaderSplit { projection: Expression, filter: Option, limit: Option, + row_limit: Option, ordered: bool, row_range: Range, selection: Selection, @@ -300,7 +310,7 @@ impl Partition for LayoutReaderSplit { let row_count = self.selection.row_count(row_count); let row_count = self.limit.map_or(row_count, |limit| row_count.min(limit)); - if self.filter.is_some() { + if self.filter.is_some() || self.row_limit.is_some() { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -318,13 +328,14 @@ impl Partition for LayoutReaderSplit { .with_projection(self.projection) .with_some_filter(self.filter) .with_some_limit(self.limit) + .with_some_row_limit(self.row_limit) .with_some_metrics_registry(self.metrics_registry) .with_ordered(self.ordered); let dtype = builder.dtype()?; // Use into_stream() which creates a LazyScanStream that spawns individual I/O // tasks onto the runtime, enabling parallel execution across executor threads. - let stream = builder.into_stream()?; + let stream = builder.into_stream()?.boxed(); Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, @@ -391,3 +402,316 @@ impl Partition for Empty { ))) } } + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::Arc; + use std::task::Poll; + + use futures::StreamExt; + use futures::TryStreamExt; + use parking_lot::Mutex; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_mask::Mask; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; + + use super::LayoutReaderDataSource; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; + use crate::scan::test::session_with_handle; + + /// A configurable [`LayoutReader`] test double. Splits come from the data source's + /// `with_split_max_row_count`; the filter passes rows through (optionally delaying the first + /// split to exercise concurrent prefetch), and projection masks/ranges plus evaluated filter + /// ranges are recorded for assertions. + #[derive(Debug)] + struct TestLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + split_size: Option, + delay_first_filter: bool, + projection_masks: Option>>>, + projection_ranges: Option>>>>, + filter_ranges: Arc>>>, + } + + impl TestLayoutReader { + fn new(row_count: u64) -> Self { + Self { + name: Arc::from("test"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + split_size: None, + delay_first_filter: false, + projection_masks: None, + projection_ranges: None, + filter_ranges: Arc::new(Mutex::new(Vec::new())), + } + } + + fn with_split_size(mut self, split_size: u64) -> Self { + self.split_size = Some(split_size); + self + } + + fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = Some(projection_masks); + self + } + + fn with_projection_ranges( + mut self, + projection_ranges: Arc>>>, + ) -> Self { + self.projection_ranges = Some(projection_ranges); + self + } + + fn with_delayed_first_filter(mut self) -> Self { + self.delay_first_filter = true; + self + } + + /// Row ranges whose filter was actually evaluated (recorded when the filter future runs, + /// not when it is merely scheduled). + fn filter_ranges(&self) -> Arc>>> { + Arc::clone(&self.filter_ranges) + } + } + + impl LayoutReader for TestLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + if let Some(size) = self.split_size { + let mut boundary = row_range.start + size; + while boundary < row_range.end { + splits.push(split_range.row_offset() + boundary); + boundary += size; + } + } + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + self.filter_ranges.lock().push(row_range.clone()); + if self.delay_first_filter && row_range.start == 0 { + let len = mask.len(); + return Ok(MaskFuture::new(len, async move { + let mut yielded = false; + futures::future::poll_fn(move |cx| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + }) + .await; + mask.await + })); + } + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let projection_masks = self.projection_masks.clone(); + let projection_ranges = self.projection_ranges.clone(); + + Ok(Box::pin(async move { + let mask = mask.await?; + if let Some(projection_masks) = projection_masks { + projection_masks.lock().push(mask.true_count()); + } + if let Some(projection_ranges) = projection_ranges { + projection_ranges.lock().push(row_range.clone()); + } + let start = i32::try_from(row_range.start)?; + let end = i32::try_from(row_range.end)?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask) + })) + } + } + + #[test] + fn filtered_limit_is_global_across_scan_partitions() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let source = LayoutReaderDataSource::new(Arc::new(TestLayoutReader::new(6)), session) + .with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + assert_eq!(values, [0, 1, 2]); + Ok(()) + } + + #[test] + fn ordered_filtered_limit_evaluates_later_split_filter_concurrently() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_ranges = Arc::new(Mutex::new(Vec::new())); + let reader = TestLayoutReader::new(4) + .with_split_size(2) + .with_projection_ranges(Arc::clone(&projection_ranges)) + .with_delayed_first_filter(); + let filter_ranges = reader.filter_ranges(); + let source = + LayoutReaderDataSource::new(Arc::new(reader), session).with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(1), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + // Ordered LIMIT is still exact: only the first split's earliest row is projected. + assert_eq!(values, [0]); + let projection_ranges = projection_ranges.lock(); + assert_eq!(projection_ranges.len(), 1); + assert_eq!(projection_ranges[0], 0..2); + drop(projection_ranges); + + // But the later split's filter still runs while the delayed first split reserves, so + // prefetch is not disabled (serializing to concurrency=1 would only ever filter 0..2). + let filter_ranges = filter_ranges.lock(); + assert!( + filter_ranges.contains(&(0..2)) && filter_ranges.contains(&(2..4)), + "expected both splits' filters to be evaluated, got {filter_ranges:?}" + ); + Ok(()) + } + + #[test] + fn unordered_limit_never_projects_more_than_the_global_budget() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let source = LayoutReaderDataSource::new( + Arc::new( + TestLayoutReader::new(12).with_projection_masks(Arc::clone(&projection_masks)), + ), + session, + ) + .with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: false, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + let chunks = runtime.block_on( + futures::stream::iter(partitions) + .map(|partition| partition.execute()) + .try_flatten_unordered(Some(6)) + .try_collect::>(), + )?; + + let mut ctx = array_session().create_execution_ctx(); + let values = chunks + .into_iter() + .map(|chunk| { + chunk + .execute::(&mut ctx) + .map(|primitive| primitive.into_buffer::()) + }) + .collect::>>()?; + let row_count = values.iter().map(|values| values.len()).sum::(); + + assert_eq!(row_count, 3); + assert_eq!(projection_masks.lock().iter().sum::(), 3); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs new file mode 100644 index 00000000000..7242e002c38 --- /dev/null +++ b/vortex-layout/src/scan/limit.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use vortex_mask::Mask; + +/// A cloneable row limit shared by all work that can contribute rows to one scan. +/// +/// Rows are reserved from a selection mask before projection work is constructed. This keeps +/// rows that cannot be returned out of projection evaluation entirely. When a limit is shared by +/// concurrent unordered partitions, reservation order is completion order, so callers may return +/// any matching rows. Ordered limited scans instead serialize their external partitions before +/// sharing a `RowLimit`, preserving the first matching rows in scan order. +#[derive(Clone)] +pub(crate) struct RowLimit(Arc); + +impl RowLimit { + pub(crate) fn new(limit: u64) -> Self { + Self(Arc::new(AtomicU64::new(limit))) + } + + /// Reserve rows selected by `mask` and retain only the earliest granted rows in that mask. + pub(crate) fn limit(&self, mask: Mask) -> Mask { + let requested = u64::try_from(mask.true_count()).unwrap_or(u64::MAX); + let granted = self.reserve(requested); + let granted = usize::try_from(granted).unwrap_or(usize::MAX); + mask.limit(granted) + } + + pub(crate) fn is_exhausted(&self) -> bool { + self.0.load(Ordering::Relaxed) == 0 + } + + fn reserve(&self, requested: u64) -> u64 { + let mut remaining = self.0.load(Ordering::Relaxed); + loop { + let granted = remaining.min(requested); + match self.0.compare_exchange_weak( + remaining, + remaining - granted, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return granted, + Err(actual) => remaining = actual, + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; + use std::thread; + + use vortex_mask::Mask; + + use super::RowLimit; + + #[test] + fn reserve_grants_up_to_the_remaining_budget() { + let limit = RowLimit::new(5); + assert_eq!(limit.reserve(3), 3); + assert!(!limit.is_exhausted()); + // Only two rows remain, so a larger request saturates at what is left. + assert_eq!(limit.reserve(10), 2); + assert!(limit.is_exhausted()); + // Once exhausted, further requests grant nothing. + assert_eq!(limit.reserve(1), 0); + } + + #[test] + fn limit_keeps_the_earliest_granted_rows() { + let limit = RowLimit::new(2); + // Rows 0, 2, 3, 5 are selected; only the first two survive the budget of 2. + let mask = Mask::from_iter([true, false, true, true, false, true]); + let limited = limit.limit(mask); + + assert_eq!(limited.true_count(), 2); + assert!(limited.value(0)); + assert!(limited.value(2)); + assert!(!limited.value(3)); + assert!(!limited.value(5)); + assert!(limit.is_exhausted()); + } + + #[test] + fn concurrent_reservations_never_exceed_the_budget() { + const THREADS: usize = 8; + const PER_THREAD: u64 = 10_000; + const LIMIT: u64 = 25_000; + + let limit = RowLimit::new(LIMIT); + let granted_total = Arc::new(AtomicU64::new(0)); + let barrier = Arc::new(Barrier::new(THREADS)); + + thread::scope(|scope| { + for _ in 0..THREADS { + let limit = limit.clone(); + let granted_total = Arc::clone(&granted_total); + let barrier = Arc::clone(&barrier); + scope.spawn(move || { + // Start all threads together to maximize contention on the atomic. + barrier.wait(); + let mut local = 0; + for _ in 0..PER_THREAD { + local += limit.reserve(1); + } + granted_total.fetch_add(local, Ordering::Relaxed); + }); + } + }); + + // Total requested (THREADS * PER_THREAD = 80_000) exceeds the budget, so exactly the + // budget is granted across all threads — no double-grant, over-grant, or lost reservation. + assert_eq!(granted_total.load(Ordering::Relaxed), LIMIT); + assert!(limit.is_exhausted()); + } +} diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index 98fd1918a42..e08e4cb3f31 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -4,6 +4,7 @@ pub mod arrow; mod filter; pub mod layout; +mod limit; pub mod multi; pub mod repeated_scan; pub mod scan_builder; diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index cb516d2500d..31d0c13e2ee 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -57,6 +57,7 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; +use crate::scan::limit::RowLimit; use crate::scan::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. @@ -87,6 +88,7 @@ pub struct MultiLayoutDataSource { concurrency: usize, } +#[derive(Clone)] pub enum MultiLayoutChild { Opened { reader: LayoutReaderRef, @@ -286,6 +288,19 @@ impl DataSource for MultiLayoutDataSource { } async fn scan(&self, scan_request: ScanRequest) -> VortexResult { + let dtype = scan_request.projection.return_dtype(&self.dtype)?; + + if scan_request.ordered && scan_request.limit.is_some() { + // Ordered global limits must consume complete files in order. A single composite + // partition owns a local mask-level limit and scans each selected file sequentially. + return Ok(Box::new(OrderedMultiLayoutScan { + session: self.session.clone(), + dtype, + request: scan_request, + children: self.children.iter().cloned().enumerate().collect(), + })); + } + let mut ready = VecDeque::new(); let mut deferred = VecDeque::new(); @@ -298,12 +313,14 @@ impl DataSource for MultiLayoutDataSource { } } - let dtype = scan_request.projection.return_dtype(&self.dtype)?; + // Only unordered scans share a completion-order limit across external partitions. + let row_limit = scan_request.limit.map(RowLimit::new); Ok(Box::new(MultiLayoutScan { session: self.session.clone(), dtype, request: scan_request, + row_limit, ready, deferred, handle: self.session.handle(), @@ -320,6 +337,7 @@ struct MultiLayoutScan { session: VortexSession, dtype: DType, request: ScanRequest, + row_limit: Option, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -345,6 +363,7 @@ impl DataSourceScan for MultiLayoutScan { session, dtype: _, request, + row_limit, ready, deferred, handle, @@ -399,13 +418,139 @@ impl DataSourceScan for MultiLayoutScan { .chain(deferred_stream) .enumerate() .flat_map(move |(i, reader_result)| match reader_result { - Ok(reader) => reader_partition(i, reader, session.clone(), request.clone()), + Ok(reader) => reader_partition( + i, + reader, + session.clone(), + request.clone(), + row_limit.clone(), + ), Err(e) => stream::once(async move { Err(e) }).boxed(), }) .boxed() } } +/// An ordered scan with a global limit. It produces one partition so a later file cannot reserve +/// rows before every earlier selected file has been accounted for. +struct OrderedMultiLayoutScan { + session: VortexSession, + dtype: DType, + request: ScanRequest, + children: VecDeque<(usize, MultiLayoutChild)>, +} + +impl DataSourceScan for OrderedMultiLayoutScan { + fn dtype(&self) -> &DType { + &self.dtype + } + + fn partition_count(&self) -> Precision { + Precision::exact(1usize) + } + + fn partitions(self: Box) -> PartitionStream { + let Self { + session, + dtype, + request, + children, + } = *self; + + stream::once(async move { + Ok(Box::new(OrderedMultiLayoutPartition { + session, + dtype, + request, + children, + }) as PartitionRef) + }) + .boxed() + } +} + +/// A composite partition that opens and scans selected readers one at a time. +struct OrderedMultiLayoutPartition { + session: VortexSession, + dtype: DType, + request: ScanRequest, + children: VecDeque<(usize, MultiLayoutChild)>, +} + +impl Partition for OrderedMultiLayoutPartition { + fn as_any(&self) -> &dyn Any { + self + } + + fn index(&self) -> usize { + 0 + } + + fn row_count(&self) -> Precision { + Precision::inexact(self.request.limit.unwrap_or_default()) + } + + fn byte_size(&self) -> Precision { + Precision::Absent + } + + fn execute(self: Box) -> VortexResult { + let Self { + session, + dtype, + request, + children, + } = *self; + let Some(limit) = request.limit else { + vortex_bail!("ordered multi-layout partitions require a row limit"); + }; + let row_limit = RowLimit::new(limit); + + let stream = async_stream::try_stream! { + for (index, child) in children { + if row_limit.is_exhausted() { + break; + } + if !partition_is_selected(index, &request) { + continue; + } + + let reader = match child { + MultiLayoutChild::Opened { reader, .. } => reader, + MultiLayoutChild::Deferred { factory, .. } => { + let Some(reader) = factory + .open() + .instrument(tracing::info_span!("LayoutReaderFactory::open")) + .await? + else { + continue; + }; + reader + } + }; + + let mut partitions = reader_partition( + index, + reader, + session.clone(), + request.clone(), + Some(row_limit.clone()), + ); + while let Some(partition) = partitions.next().await { + let mut chunks = partition?.execute()?; + while let Some(chunk) = chunks.next().await { + yield chunk?; + } + } + } + }; + + Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( + dtype, stream, + ))) + } +} + /// Generates a partition stream for a single layout reader. /// /// Checks file-level pruning first (via `pruning_evaluation`). If the filter proves no rows @@ -416,29 +561,14 @@ fn reader_partition( reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, + row_limit: Option, ) -> PartitionStream { let row_count = reader.row_count(); let row_range = request.row_range.clone().unwrap_or(0..row_count); - let partition_idx_u64: u64 = partition_idx as u64; - if let Some(range) = &request.partition_range - && !range.contains(&partition_idx_u64) - { + if !partition_is_selected(partition_idx, &request) { return stream::empty().boxed(); - }; - match &request.partition_selection { - Selection::IncludeByIndex(buffer) => { - if buffer.as_slice().binary_search(&partition_idx_u64).is_err() { - return stream::empty().boxed(); - } - } - Selection::ExcludeByIndex(buffer) => { - if buffer.as_slice().binary_search(&partition_idx_u64).is_ok() { - return stream::empty().boxed(); - } - } - _ => {} - }; + } // Check file-level pruning: if the filter can be proven false for the entire row range // using file-level statistics, skip this reader entirely. @@ -461,12 +591,32 @@ fn reader_partition( row_range: Some(row_range), ..request }, + row_limit, index: partition_idx, }) as PartitionRef) }) .boxed() } +fn partition_is_selected(partition_idx: usize, request: &ScanRequest) -> bool { + let partition_idx = partition_idx as u64; + if let Some(range) = &request.partition_range + && !range.contains(&partition_idx) + { + return false; + } + + match &request.partition_selection { + Selection::IncludeByIndex(buffer) => { + buffer.as_slice().binary_search(&partition_idx).is_ok() + } + Selection::ExcludeByIndex(buffer) => { + buffer.as_slice().binary_search(&partition_idx).is_err() + } + _ => true, + } +} + /// A partition backed by a single [`LayoutReaderRef`] and a row range. /// /// On `execute()`, creates a [`ScanBuilder`] over the row range, enabling @@ -475,6 +625,7 @@ struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, + row_limit: Option, index: usize, } @@ -498,7 +649,7 @@ impl Partition for MultiLayoutPartition { .limit .map_or(row_count, |limit| row_count.min(limit)); - if self.request.filter.is_some() { + if self.request.filter.is_some() || self.row_limit.is_some() { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -516,6 +667,7 @@ impl Partition for MultiLayoutPartition { .with_projection(request.projection) .with_some_filter(request.filter) .with_some_limit(request.limit) + .with_some_row_limit(self.row_limit) .with_ordered(request.ordered); if let Some(row_range) = request.row_range { @@ -523,7 +675,7 @@ impl Partition for MultiLayoutPartition { } let dtype = builder.dtype()?; - let stream = builder.into_stream()?; + let stream = builder.into_stream()?.boxed(); Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, @@ -533,11 +685,36 @@ impl Partition for MultiLayoutPartition { #[cfg(test)] mod tests { + use std::ops::Range; + use std::sync::Arc; + + use async_trait::async_trait; + use futures::TryStreamExt; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_mask::Mask; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; use super::*; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; use crate::scan::test::new_session; + use crate::scan::test::session_with_handle; struct NeverOpened; @@ -569,4 +746,133 @@ mod tests { fn byte_size_precision(#[case] sizes: Vec>, #[case] expected: Precision) { assert_eq!(deferred_source(sizes).byte_size(), expected); } + + #[derive(Debug)] + struct TestLayoutReader { + name: Arc, + dtype: DType, + base: i32, + row_count: u64, + } + + impl TestLayoutReader { + fn new(name: &'static str, base: i32, row_count: u64) -> Self { + Self { + name: Arc::from(name), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + base, + row_count, + } + } + } + + impl LayoutReader for TestLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let base = self.base; + + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start)?; + let end = i32::try_from(row_range.end)?; + PrimitiveArray::from_iter((start..end).map(|value| base + value)) + .into_array() + .filter(mask.await?) + })) + } + } + + struct StaticReaderFactory { + reader: LayoutReaderRef, + } + + #[async_trait] + impl LayoutReaderFactory for StaticReaderFactory { + async fn open(&self) -> VortexResult> { + Ok(Some(Arc::clone(&self.reader))) + } + } + + #[test] + fn ordered_limit_scans_multiple_readers_sequentially() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let first: LayoutReaderRef = Arc::new(TestLayoutReader::new("first", 0, 2)); + let second: LayoutReaderRef = Arc::new(TestLayoutReader::new("second", 10, 2)); + let source = MultiLayoutDataSource::new_with_first( + first, + vec![Arc::new(StaticReaderFactory { reader: second })], + vec![], + &session, + ); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + assert_eq!(values, [0, 1, 10]); + Ok(()) + } } diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 681f33639bc..7c3ea7973b4 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -4,10 +4,16 @@ use std::cmp; use std::iter; use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; use futures::Stream; -use futures::future::BoxFuture; +use futures::StreamExt; +use futures::future; +use futures::stream::BoxStream; use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; @@ -20,6 +26,8 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::runtime::Handle; +use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_scan::selection::Selection; use vortex_session::VortexSession; @@ -27,15 +35,20 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; +use crate::scan::limit::RowLimit; use crate::scan::splits::Splits; use crate::scan::tasks::TaskContext; +use crate::scan::tasks::TaskFuture; +use crate::scan::tasks::TaskResult; +use crate::scan::tasks::filter_split; +use crate::scan::tasks::project_split; use crate::scan::tasks::split_exec; /// A projected subset (by indices, range, and filter) of rows from a Vortex data source. /// /// The method of this struct enable, possibly concurrent, scanning of multiple row ranges of this /// data source. -pub struct RepeatedScan { +pub struct RepeatedScan { session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -49,15 +62,172 @@ pub struct RepeatedScan { splits: Splits, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, - /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - map_fn: Arc VortexResult + Send + Sync>, - /// Maximal number of rows to read (after filtering) + /// Maximal number of rows to read (after filtering). + /// + /// When no shared [`RowLimit`] (`row_limit`) is set, this limit is applied independently to + /// each `execute_*` call: a `RepeatedScan` executed over several row ranges may therefore + /// return up to `limit` rows *per call*. Supply a shared `row_limit` to cap the total across + /// calls and sibling partitions. limit: Option, + /// An optional row limit shared with sibling external partitions. + row_limit: Option, /// The dtype of the projected arrays. dtype: DType, } -impl RepeatedScan { +/// A source of split tasks that has not yet applied task concurrency or output error handling. +#[must_use = "task streams must be scheduled"] +struct TaskStream { + inner: BoxStream<'static, Task>, +} + +impl TaskStream { + fn new(stream: impl Stream> + Send + 'static) -> Self { + Self { + inner: stream.boxed(), + } + } + + fn eager(handle: Handle, tasks: Vec) -> Self { + Self::new(futures::stream::iter(tasks).map(move |task| handle.spawn(task))) + } + + fn unordered_filtered_limit( + handle: Handle, + split_ranges: Vec>, + selection: Selection, + ctx: Arc, + row_limit: RowLimit, + ) -> Self { + let take_limit = row_limit.clone(); + Self::new( + futures::stream::iter(split_ranges) + .take_while(move |_| future::ready(!take_limit.is_exhausted())) + .filter_map(move |range| { + // Build the row mask and split task synchronously so the I/O system sees the + // split's ranges as soon as `buffer_unordered` pulls it, without cloning + // `selection`. + let row_mask = selection.row_mask(&range); + let task = (!row_mask.mask().all_false()).then(|| { + // A synchronous split-construction failure happens before any row is + // reserved, so it is recoverable (yielded as a stream error) rather than + // terminal (which would abort the whole scan). See the `TaskResult` docs. + split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) + .unwrap_or_else(TaskFuture::recoverable) + }); + future::ready(task.map(|task| handle.spawn(task))) + }), + ) + } + + fn ordered_filtered_limit( + handle: Handle, + split_ranges: Vec>, + selection: Selection, + ctx: Arc, + row_limit: RowLimit, + concurrency: usize, + ) -> Self { + // Stage one: evaluate the filter for each split concurrently while preserving split + // order. + let filter_tasks = { + let filter_ctx = Arc::clone(&ctx); + let filter_handle = handle.clone(); + let take_limit = row_limit.clone(); + futures::stream::iter(split_ranges) + .take_while(move |_| future::ready(!take_limit.is_exhausted())) + .filter_map(move |range| { + // Build the row mask synchronously so the I/O system sees the split's ranges + // as soon as `buffered` pulls it, without cloning `selection`. + let row_mask = selection.row_mask(&range); + let task = (!row_mask.mask().all_false()).then(|| { + filter_handle.spawn(filter_split(Arc::clone(&filter_ctx), row_mask)) + }); + future::ready(task) + }) + .buffered(concurrency) + }; + + // Seam + stage two: reserve in split order (this map runs sequentially as the ordered + // stage-one stream is consumed), then spawn projection work for the reserved mask. The + // second gate drops results that were already filtering when an earlier split exhausted + // the budget, including any later filter errors. + let stage_two_limit = row_limit.clone(); + Self::new( + filter_tasks + .take_while(move |_| future::ready(!stage_two_limit.is_exhausted())) + .map(move |result| { + let task = match result { + Ok((row_range, mask)) => { + let mask = row_limit.limit(mask); + if mask.all_false() { + TaskFuture::empty() + } else { + project_split(Arc::clone(&ctx), row_range, mask) + } + } + Err(error) => TaskFuture::recoverable(error), + }; + handle.spawn(task) + }), + ) + } +} + +impl Stream for TaskStream { + type Item = Task; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_next(cx) + } +} + +/// A buffered task stream that exposes arrays and applies the task error policy. +pub(crate) struct ScheduledTaskStream { + tasks: Option>, +} + +impl ScheduledTaskStream { + fn new(tasks: TaskStream, ordered: bool, concurrency: usize) -> Self { + let tasks = if ordered { + tasks.buffered(concurrency).boxed() + } else { + tasks.buffer_unordered(concurrency).boxed() + }; + Self { tasks: Some(tasks) } + } +} + +impl Stream for ScheduledTaskStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let Some(tasks) = self.tasks.as_mut() else { + return Poll::Ready(None); + }; + let result = ready!(tasks.as_mut().poll_next(cx)); + let Some(result) = result else { + self.tasks = None; + return Poll::Ready(None); + }; + match result { + TaskResult::Array(Some(array)) => return Poll::Ready(Some(Ok(array))), + TaskResult::Array(None) => {} + TaskResult::Recoverable(error) => return Poll::Ready(Some(Err(error))), + TaskResult::Terminal(error) => { + // Drop queued and in-flight task handles before yielding the error. Their + // `Drop` implementations abort work that can no longer contribute to this + // limited scan. + self.tasks = None; + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + +impl RepeatedScan { pub fn dtype(&self) -> &DType { &self.dtype } @@ -81,15 +251,13 @@ impl RepeatedScan { let stream = self.execute_stream(row_range)?; Ok(ArrayStreamAdapter::new(dtype, stream)) } -} -impl RepeatedScan { /// Constructor just to allow `scan_builder` to create a `RepeatedScan`. #[expect( clippy::too_many_arguments, reason = "all arguments are needed for scan construction" )] - pub fn new( + pub(crate) fn new( session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -99,8 +267,8 @@ impl RepeatedScan { selection: Selection, splits: Splits, concurrency: usize, - map_fn: Arc VortexResult + Send + Sync>, limit: Option, + row_limit: Option, dtype: DType, ) -> Self { Self { @@ -113,16 +281,13 @@ impl RepeatedScan { selection, splits, concurrency, - map_fn, limit, + row_limit, dtype, } } - pub fn execute( - &self, - row_range: Option>, - ) -> VortexResult>>>> { + fn split_ranges(&self, row_range: Option>) -> Vec> { let selection_range: Option> = match &self.selection { Selection::IncludeByIndex(buf) if !buf.is_empty() => { Some(buf[0]..buf[buf.len() - 1] + 1) @@ -135,14 +300,14 @@ impl RepeatedScan { let row_range = intersect_ranges(self.row_range.as_ref(), row_range); let row_range = intersect_ranges(row_range.as_ref(), selection_range); - let ranges = match &self.splits { + match &self.splits { Splits::Natural(vec) => { debug_assert!(vec.is_sorted()); let splits_iter = match row_range { None => Either::Left(vec.iter().copied()), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } let lo = vec.partition_point(|&x| x < range.start); let hi = vec.partition_point(|&x| x < range.end); @@ -154,66 +319,135 @@ impl RepeatedScan { } }; - Either::Left(splits_iter.tuple_windows().map(|(start, end)| start..end)) + splits_iter + .tuple_windows() + .map(|(start, end)| start..end) + .collect() } - Splits::Ranges(ranges) => Either::Right(match row_range { - None => Either::Left(ranges.iter().cloned()), + Splits::Ranges(ranges) => match row_range { + None => ranges.to_vec(), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } - Either::Right(ranges.iter().filter_map(move |r| { - let start = cmp::max(r.start, range.start); - let end = cmp::min(r.end, range.end); - (start < end).then_some(start..end) - })) + ranges + .iter() + .filter_map(move |r| { + let start = cmp::max(r.start, range.start); + let end = cmp::min(r.end, range.end); + (start < end).then_some(start..end) + }) + .collect() } - }), - }; + }, + } + } - let mut limit = self.limit; - let mut tasks = Vec::new(); - let ctx = Arc::new(TaskContext { + fn task_context(&self) -> Arc { + Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), reader: Arc::clone(&self.layout_reader), projection: self.projection.clone(), - mapper: Arc::clone(&self.map_fn), - }); + }) + } + + pub(crate) fn execute( + &self, + row_range: Option>, + row_limit: Option, + ) -> VortexResult> { + let ctx = self.task_context(); + + let mut tasks = Vec::new(); + + for range in self.split_ranges(row_range) { + if range.start >= range.end { + continue; + } + if row_limit.as_ref().is_some_and(RowLimit::is_exhausted) { + break; + } - for range in ranges { let row_mask = self.selection.row_mask(&range); if row_mask.mask().all_false() { continue; } - tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); - if limit.is_some_and(|l| l == 0) { - break; - } + tasks.push(split_exec(Arc::clone(&ctx), row_mask, row_limit.clone())?); } Ok(tasks) } - pub fn execute_stream( + pub(crate) fn execute_stream( &self, row_range: Option>, - ) -> VortexResult> + Send + 'static + use> { - use futures::StreamExt; + ) -> VortexResult { let num_workers = get_available_parallelism().unwrap_or(1); + let row_limit = self + .row_limit + .clone() + .or_else(|| self.limit.map(RowLimit::new)); let concurrency = self.concurrency * num_workers; let handle = self.session.handle(); - let stream = - futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); + // With both a filter and a limit we cannot know each split's output row count ahead of + // time, so split tasks are built lazily as the stream is polled. Once another task drains + // the shared budget, `take_while` prevents further task creation; already-prefetched + // tasks observe the exhausted budget after filtering and return `None` without projection. + if self.filter.is_some() + && let Some(row_limit) = row_limit.clone() + { + // An ordered LIMIT would be violated if a later split reserved rows before an earlier + // one, so ordered scans use a two-stage pipeline that keeps I/O concurrent while + // reserving strictly in split order. + if self.ordered { + return Ok(self.ordered_filtered_limit_stream(row_range, row_limit, concurrency)); + } - let stream = if self.ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; + let tasks = TaskStream::unordered_filtered_limit( + handle, + self.split_ranges(row_range), + self.selection.clone(), + self.task_context(), + row_limit, + ); + return Ok(ScheduledTaskStream::new(tasks, false, concurrency)); + } + + // No filter (or no limit): build every task eagerly so the IO system sees all split + // ranges up front. A no-filter limit is applied to each selection mask inside `execute`. + let tasks = TaskStream::eager(handle, self.execute(row_range, row_limit)?); + + let ordered = self.ordered; + Ok(ScheduledTaskStream::new(tasks, ordered, concurrency)) + } - Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) + /// Ordered filtered scan with a shared row limit. + /// + /// I/O concurrency is decoupled from limit reservation so that an ordered `LIMIT` does not + /// force serial execution. Filters for a window of splits are evaluated concurrently (stage + /// one), rows are reserved against the shared limit strictly in split order on the seam + /// between the stages, and only the reserved masks are projected (stage two, also concurrent). + /// Because reservation happens in split order, the earliest matching rows always win the + /// budget; because projection only runs for reserved masks, no rows past the limit are + /// decoded. `take_while` stops feeding stage one once the limit is exhausted, bounding + /// speculative filter I/O to the pipeline depth. + fn ordered_filtered_limit_stream( + &self, + row_range: Option>, + row_limit: RowLimit, + concurrency: usize, + ) -> ScheduledTaskStream { + let tasks = TaskStream::ordered_filtered_limit( + self.session.handle(), + self.split_ranges(row_range), + self.selection.clone(), + self.task_context(), + row_limit, + concurrency, + ); + ScheduledTaskStream::new(tasks, true, concurrency) } } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..6b952a601c5 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -10,7 +10,6 @@ use std::task::ready; use futures::Stream; use futures::StreamExt; -use futures::future::BoxFuture; use futures::stream::BoxStream; use itertools::Itertools; use vortex_array::ArrayRef; @@ -21,31 +20,28 @@ use vortex_array::expr::analysis::referenced_field_paths; use vortex_array::expr::root; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; -use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_io::runtime::BlockingRuntime; -use vortex_io::runtime::Handle; use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_metrics::MetricsRegistry; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::layouts::row_idx::RowIdxLayoutReader; +use crate::scan::limit::RowLimit; use crate::scan::repeated_scan::RepeatedScan; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; -/// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. +/// Builder for scanning a [`LayoutReader`] into arrays, streams, or iterators. /// /// A scan has three independent row restriction mechanisms: /// @@ -56,7 +52,7 @@ use crate::scan::splits::attempt_split_ranges; /// Projection and filter expressions are optimized against the reader dtype during /// [`prepare`](Self::prepare). Work is divided by the configured [`SplitBy`] strategy or by /// explicit selection ranges. -pub struct ScanBuilder { +pub struct ScanBuilder { session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -72,19 +68,17 @@ pub struct ScanBuilder { split_by: SplitBy, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, - /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, - /// Should we try to prune the file (using stats) on open. - file_stats: Option>, - /// Maximal number of rows to read (after filtering) + /// Maximal number of rows to read after filtering. limit: Option, + /// A row limit shared with sibling external partitions, when the caller owns one. + row_limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, /// but not by the scan [`Selection`] which remains relative. row_offset: u64, } -impl ScanBuilder { +impl ScanBuilder { /// Create a scan builder over `layout_reader` using `session` for runtime and execution state. pub fn new(session: VortexSession, layout_reader: Arc) -> Self { Self { @@ -99,10 +93,9 @@ impl ScanBuilder { // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. concurrency: 4, - map_fn: Arc::new(Ok), metrics_registry: None, - file_stats: None, limit: None, + row_limit: None, row_offset: 0, } } @@ -130,7 +123,7 @@ impl ScanBuilder { } } -impl ScanBuilder { +impl ScanBuilder { /// Add a filter expression evaluated against the projected row ranges. pub fn with_filter(mut self, filter: Expression) -> Self { self.filter = Some(filter); @@ -227,6 +220,12 @@ impl ScanBuilder { self } + /// Use a row limit supplied by the enclosing data source instead of creating a local one. + pub(crate) fn with_some_row_limit(mut self, row_limit: Option) -> Self { + self.row_limit = row_limit; + self + } + /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { self.projection.return_dtype(self.layout_reader.dtype()) @@ -237,38 +236,10 @@ impl ScanBuilder { &self.session } - /// Map each split of the scan. The function will be run on the spawned task. - pub fn map( - self, - map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, - ) -> ScanBuilder { - let old_map_fn = self.map_fn; - ScanBuilder { - session: self.session, - layout_reader: self.layout_reader, - projection: self.projection, - filter: self.filter, - ordered: self.ordered, - row_range: self.row_range, - selection: self.selection, - split_by: self.split_by, - concurrency: self.concurrency, - metrics_registry: self.metrics_registry, - file_stats: self.file_stats, - limit: self.limit, - row_offset: self.row_offset, - map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), - } - } - /// Optimize expressions, compute split ranges, and return an executable repeated scan. - pub fn prepare(self) -> VortexResult> { + pub fn prepare(self) -> VortexResult { let dtype = self.dtype()?; - if self.filter.is_some() && self.limit.is_some() { - vortex_bail!("Vortex doesn't support scans with both a filter and a limit") - } - // Spin up the root layout reader, and wrap it in a FilterLayoutReader to perform // conjunction splitting if a filter is provided. let mut layout_reader = self.layout_reader; @@ -319,26 +290,19 @@ impl ScanBuilder { self.selection, splits, self.concurrency, - self.map_fn, self.limit, + self.row_limit, dtype, )) } - /// Constructs a task per row split of the scan, returned as a vector of futures. - pub fn build(self) -> VortexResult>>>> { - // The ultimate short circuit - if self.limit.is_some_and(|l| l == 0) { - return Ok(vec![]); - } - - self.prepare()?.execute(None) - } - /// Returns a [`Stream`] with tasks spawned onto the session's runtime handle. + /// + /// Preparation and initial stream construction begin on the first poll. Errors from either + /// step are returned as the stream's next item. pub fn into_stream( self, - ) -> VortexResult> + Send + 'static + use> { + ) -> VortexResult> + Send + 'static> { Ok(LazyScanStream::new(self)) } @@ -346,82 +310,60 @@ impl ScanBuilder { pub fn into_iter( self, runtime: &B, - ) -> VortexResult> + 'static> { + ) -> VortexResult> + 'static> { let stream = self.into_stream()?; Ok(runtime.block_on_stream(stream)) } } -enum LazyScanState { - Builder(Option>>), - Preparing(PreparingScan), - Stream(BoxStream<'static, VortexResult>), +enum LazyScanState { + Builder(Option>), + Preparing(PreparingScan), + Stream(BoxStream<'static, VortexResult>), Error(Option), } -type PreparedScanTasks = Vec>>>; - -struct PreparingScan { - ordered: bool, - concurrency: usize, - handle: Handle, - task: Task>>, +struct PreparingScan { + task: Task>>>, } -struct LazyScanStream { - state: LazyScanState, +struct LazyScanStream { + state: LazyScanState, } -impl LazyScanStream { - fn new(builder: ScanBuilder) -> Self { +impl LazyScanStream { + fn new(builder: ScanBuilder) -> Self { Self { state: LazyScanState::Builder(Some(Box::new(builder))), } } } -impl Unpin for LazyScanStream {} +impl Unpin for LazyScanStream {} -impl Stream for LazyScanStream { - type Item = VortexResult; +impl Stream for LazyScanStream { + type Item = VortexResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { match &mut self.state { LazyScanState::Builder(builder) => { let builder = builder.take().vortex_expect("polled after completion"); - let ordered = builder.ordered; - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = builder.concurrency * num_workers; let handle = builder.session.handle(); + // IMPORTANT: Building the stream can synchronously walk the layout and + // register I/O for every split. Keep it with preparation in this blocking + // task: poll_next must only wait for and poll an already-constructed stream. + // This also keeps construction errors on the Preparing -> Error path rather + // than running construction on the caller's executor. let task = handle.spawn_blocking(move || { - builder.prepare().and_then(|scan| scan.execute(None)) - }); - self.state = LazyScanState::Preparing(PreparingScan { - ordered, - concurrency, - handle, - task, + let scan = builder.prepare()?; + Ok(scan.execute_stream(None)?.boxed()) }); + self.state = LazyScanState::Preparing(PreparingScan { task }); } LazyScanState::Preparing(preparing) => { match ready!(Pin::new(&mut preparing.task).poll(cx)) { - Ok(tasks) => { - let ordered = preparing.ordered; - let concurrency = preparing.concurrency; - let handle = preparing.handle.clone(); - let stream = - futures::stream::iter(tasks).map(move |task| handle.spawn(task)); - let stream = if ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; - let stream = stream - .filter_map(|chunk| async move { chunk.transpose() }) - .boxed(); - self.state = LazyScanState::Stream(stream); - } + Ok(stream) => self.state = LazyScanState::Stream(stream), Err(err) => self.state = LazyScanState::Error(Some(err)), } } @@ -464,8 +406,11 @@ mod test { use std::time::Duration; use futures::Stream; + use futures::channel::oneshot; use futures::task::noop_waker_ref; use parking_lot::Mutex; + use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; @@ -487,6 +432,7 @@ mod test { use vortex_error::vortex_err; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_io::runtime::tokio::TokioRuntime; use vortex_mask::Mask; use super::ScanBuilder; @@ -647,6 +593,13 @@ mod test { dtype: DType, row_count: u64, register_splits_calls: Arc, + blocking_projection: Option, + } + + #[derive(Debug)] + struct BlockingProjection { + started: mpsc::Sender<()>, + gate: Arc>, } impl SplittingLayoutReader { @@ -656,8 +609,19 @@ mod test { dtype: DType::Primitive(PType::I32, Nullability::NonNullable), row_count: 4, register_splits_calls, + blocking_projection: None, } } + + fn with_blocking_projection( + register_splits_calls: Arc, + gate: Arc>, + started: mpsc::Sender<()>, + ) -> Self { + let mut reader = Self::new(register_splits_calls); + reader.blocking_projection = Some(BlockingProjection { started, gate }); + reader + } } impl LayoutReader for SplittingLayoutReader { @@ -710,6 +674,14 @@ mod test { _expr: &Expression, _mask: MaskFuture, ) -> VortexResult { + if let Some(blocking_projection) = &self.blocking_projection { + blocking_projection + .started + .send(()) + .map_err(|_| vortex_err!("test projection-start receiver dropped"))?; + let _guard = blocking_projection.gate.lock(); + } + let start = usize::try_from(row_range.start) .map_err(|_| vortex_err!("row_range.start must fit in usize"))?; let end = usize::try_from(row_range.end) @@ -752,6 +724,545 @@ mod test { Ok(()) } + #[tokio::test] + async fn into_stream_constructs_tasks_off_the_poller() -> VortexResult<()> { + let gate = Arc::new(Mutex::new(())); + let guard = gate.lock(); + let calls = Arc::new(AtomicUsize::new(0)); + let (started_send, started_recv) = mpsc::channel(); + let reader = Arc::new(SplittingLayoutReader::with_blocking_projection( + Arc::clone(&calls), + Arc::clone(&gate), + started_send, + )); + + let runtime = TokioRuntime::new(tokio::runtime::Handle::current()); + let session = session_with_handle(runtime.handle()); + let mut stream = ScanBuilder::new(session, reader).into_stream()?; + + let (poll_send, poll_recv) = mpsc::channel(); + let (release_send, release_recv) = mpsc::channel(); + let join = std::thread::spawn(move || { + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + let poll = Pin::new(&mut stream).poll_next(&mut cx); + let _ = poll_send.send(matches!(poll, Poll::Pending)); + let _ = release_recv.recv(); + }); + + let poll_result = poll_recv.recv_timeout(Duration::from_secs(1)); + let projection_started = started_recv.recv_timeout(Duration::from_secs(1)); + + // Release the task and join its caller before reporting a failed assertion. + drop(guard); + let _ = release_send.send(()); + drop(join.join()); + + assert!( + poll_result.is_ok_and(|poll_pending| poll_pending), + "first poll must return while scan task construction is blocked" + ); + projection_started + .map_err(|_| vortex_err!("stream construction did not begin in the background"))?; + assert_eq!(calls.load(Ordering::Relaxed), 1); + + Ok(()) + } + + #[tokio::test] + async fn into_stream_reports_stream_construction_errors() -> VortexResult<()> { + let range_start = i32::MAX as u64 + 1; + let reader = Arc::new(SplittingLayoutReader::new(Arc::new(AtomicUsize::new(0)))); + let session = session_with_handle(TokioRuntime::current()); + let mut stream = ScanBuilder::new(session, reader) + .with_row_range(range_start..range_start + 1) + .into_stream()?; + + assert!(matches!( + futures::StreamExt::next(&mut stream).await, + Some(Err(_)) + )); + assert!(futures::StreamExt::next(&mut stream).await.is_none()); + + Ok(()) + } + + /// A configurable [`LayoutReader`] test double that replaces several near-identical mocks. + /// + /// `split_size` controls the split layout (`None` is a single split), `keep_row` filters rows, + /// and the `fail_*` flags inject filter/projection errors. Every projection records its mask's + /// true-count into `projection_masks`, letting tests assert the limit is applied before + /// projection. + struct MockLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + split_size: Option, + keep_row: fn(u64) -> bool, + fail_first_filter: bool, + fail_first_projection: bool, + fail_projection: bool, + projection_masks: Arc>>, + } + + impl MockLayoutReader { + fn new(row_count: u64) -> Self { + Self { + name: Arc::from("mock"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + split_size: None, + keep_row: |_| true, + fail_first_filter: false, + fail_first_projection: false, + fail_projection: false, + projection_masks: Arc::new(Mutex::new(Vec::new())), + } + } + + fn with_split_size(mut self, split_size: u64) -> Self { + self.split_size = Some(split_size); + self + } + + fn with_keep_row(mut self, keep_row: fn(u64) -> bool) -> Self { + self.keep_row = keep_row; + self + } + + fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = projection_masks; + self + } + + fn with_fail_first_filter(mut self) -> Self { + self.fail_first_filter = true; + self + } + + fn with_fail_first_projection(mut self) -> Self { + self.fail_first_projection = true; + self + } + + fn with_projection_error(mut self) -> Self { + self.fail_projection = true; + self + } + } + + impl LayoutReader for MockLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + if let Some(size) = self.split_size { + let mut boundary = row_range.start + size; + while boundary < row_range.end { + splits.push(split_range.row_offset() + boundary); + boundary += size; + } + } + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + if self.fail_first_filter && row_range.start == 0 { + let len = mask.len(); + return Ok(MaskFuture::new(len, async move { + Err(vortex_err!("first split filter failed")) + })); + } + + let row_range = row_range.clone(); + let keep_row = self.keep_row; + let row_count = usize::try_from(row_range.end - row_range.start) + .map_err(|_| vortex_err!("row range must fit in usize"))?; + + Ok(MaskFuture::new(row_count, async move { + let input_mask = mask.await?; + Ok(Mask::from_iter( + (row_range.start..row_range.end) + .enumerate() + .map(|(idx, row)| input_mask.value(idx) && keep_row(row)), + )) + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let projection_masks = Arc::clone(&self.projection_masks); + let fail = self.fail_projection || (self.fail_first_projection && row_range.start == 0); + + Ok(Box::pin(async move { + let mask = mask.await?; + projection_masks.lock().push(mask.true_count()); + if fail { + return Err(vortex_err!("projection failed")); + } + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + struct MatchThenErrorLayoutReader { + name: Arc, + dtype: DType, + release_first_filter: Mutex>>, + wait_for_later_filter: Mutex>>, + } + + impl MatchThenErrorLayoutReader { + fn new() -> Self { + let (release_first_filter, wait_for_later_filter) = oneshot::channel(); + Self { + name: Arc::from("match-then-error"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + release_first_filter: Mutex::new(Some(release_first_filter)), + wait_for_later_filter: Mutex::new(Some(wait_for_later_filter)), + } + } + } + + impl LayoutReader for MatchThenErrorLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 2 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + splits.push(split_range.row_offset() + row_range.start + 1); + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let len = mask.len(); + if row_range.start == 0 { + let wait_for_later_filter = self + .wait_for_later_filter + .lock() + .take() + .ok_or_else(|| vortex_err!("first split filter was evaluated twice"))?; + return Ok(MaskFuture::new(len, async move { + wait_for_later_filter + .await + .map_err(|_| vortex_err!("later split filter was cancelled"))?; + mask.await + })); + } + + self.release_first_filter + .lock() + .take() + .ok_or_else(|| vortex_err!("later split filter was evaluated twice"))? + .send(()) + .map_err(|_| vortex_err!("first split filter was cancelled"))?; + Ok(MaskFuture::new(len, async move { + Err(vortex_err!("later split filter failed")) + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask.await?) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + fn collect_scan_values(iter: I) -> VortexResult> + where + I: IntoIterator>, + { + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for chunk in iter { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + Ok(values) + } + + fn drain_runtime(runtime: &SingleThreadRuntime) { + for _ in 0..4 { + let mut yielded = false; + runtime.block_on(futures::future::poll_fn(move |cx| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + })); + } + } + + fn keep_all(_: u64) -> bool { + true + } + + fn keep_odd(row: u64) -> bool { + row % 2 == 1 + } + + #[rstest] + #[case::limit_below_matches(8, keep_all, 3, &[0, 1, 2])] + #[case::limit_zero(8, keep_all, 0, &[])] + #[case::limit_exceeds_matches(8, keep_odd, 100, &[1, 3, 5, 7])] + #[case::empty_input(0, keep_all, 3, &[])] + fn filtered_limit_yields_expected_rows( + #[case] row_count: u64, + #[case] keep_row: fn(u64) -> bool, + #[case] limit: u64, + #[case] expected: &[i32], + ) -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new( + MockLayoutReader::new(row_count) + .with_split_size(2) + .with_keep_row(keep_row), + ); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(limit) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + assert_eq!(values.as_slice(), expected); + Ok(()) + } + + #[test] + fn filtered_limit_limits_projection_mask_before_projection() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new( + MockLayoutReader::new(100_000).with_projection_masks(Arc::clone(&projection_masks)), + ); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + + assert_eq!(values, [0]); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + + #[test] + fn ordered_filtered_limit_drops_prefetched_filter_errors_after_the_budget_is_full() + -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new(MatchThenErrorLayoutReader::new()); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + + // The first split waits until the later split has already produced its error. Once the + // first matching row fills the limit, that speculative error must be dropped. + let values = collect_scan_values(runtime.block_on_stream(stream))?; + + assert_eq!(values, [0]); + Ok(()) + } + + #[test] + fn filter_errors_are_stream_items_and_do_not_consume_the_limit() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new( + MockLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_filter(), + ); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + let Some(chunk) = iter.next() else { + return Err(vortex_err!( + "matching split was not polled after the filter error" + )); + }; + let mut ctx = array_session().create_execution_ctx(); + let primitive = chunk?.execute::(&mut ctx)?; + + assert_eq!(primitive.into_buffer::().as_slice(), [1]); + assert!(iter.next().is_none()); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + + #[test] + fn projection_error_after_reservation_terminates_the_limited_scan() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new( + MockLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_projection(), + ); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + // A budget of two leaves room for the second matching split. Continuing after the + // first projection failure would therefore yield a second stream item. + .with_limit(2) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + assert!(iter.next().is_none()); + assert!(projection_masks.lock().contains(&1)); + Ok(()) + } + + #[test] + fn projection_errors_are_stream_items() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new( + MockLayoutReader::new(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_projection_error(), + ); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + assert!(iter.next().is_none()); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + + #[test] + fn prepared_scan_limits_filtered_results() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new( + MockLayoutReader::new(8) + .with_split_size(2) + .with_keep_row(keep_odd), + ); + + let scan = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(3) + .prepare()?; + let values = collect_scan_values(scan.execute_array_iter(None, &runtime)?)?; + drain_runtime(&runtime); + + assert_eq!(values, [1, 3, 5]); + Ok(()) + } + #[derive(Debug)] struct BlockingSplitsLayoutReader { name: Arc, diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index a86546e15ef..7a987555291 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -3,8 +3,13 @@ //! Split scanning task implementation. +use std::future::Future; use std::ops::BitAnd; +use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; use bit_vec::BitVec; use futures::FutureExt; @@ -12,14 +17,144 @@ use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::expr::Expression; +use vortex_error::VortexError; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; +use crate::ArrayFuture; use crate::LayoutReader; use crate::scan::filter::FilterExpr; +use crate::scan::limit::RowLimit; -pub type TaskFuture = BoxFuture<'static, VortexResult>; +/// The result of a split task. +/// +/// Filter errors happen before a row limit reserves any rows, so callers may report them and +/// continue with later splits. Projection errors after reservation cannot safely release rows back +/// to a concurrent limit, so callers must report them and terminate the limited scan. +pub(crate) enum TaskResult { + /// A completed projection, or an empty split. + Array(Option), + /// An error that occurred before a row limit reserved rows. + Recoverable(VortexError), + /// An error that occurred after a row limit reserved rows. + Terminal(VortexError), +} + +/// A future that executes one split and classifies any failure by whether it happened before or +/// after a row-limit reservation. +#[must_use = "split tasks must be scheduled or awaited"] +pub(crate) struct TaskFuture { + inner: BoxFuture<'static, TaskResult>, +} + +impl TaskFuture { + fn new(future: impl Future + Send + 'static) -> Self { + Self { + inner: future.boxed(), + } + } + + fn ready(result: TaskResult) -> Self { + Self::new(futures::future::ready(result)) + } + + pub(crate) fn empty() -> Self { + Self::ready(TaskResult::Array(None)) + } + + pub(crate) fn recoverable(error: VortexError) -> Self { + Self::ready(TaskResult::Recoverable(error)) + } + + pub(crate) fn terminal(error: VortexError) -> Self { + Self::ready(TaskResult::Terminal(error)) + } + + fn projection(projection: ArrayFuture, terminal: bool) -> Self { + Self::new(async move { + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) if terminal => TaskResult::Terminal(error), + Err(error) => TaskResult::Recoverable(error), + } + }) + } + + fn filtered_projection(filter_mask: MaskFuture, projection: ArrayFuture) -> Self { + Self::new(async move { + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(error) => return TaskResult::Recoverable(error), + }; + if mask.all_false() { + return TaskResult::Array(None); + } + + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => TaskResult::Recoverable(error), + } + }) + } + + fn limited_filtered_projection( + ctx: Arc, + row_range: Range, + filter_mask: MaskFuture, + row_limit: RowLimit, + ) -> Self { + Self::new(async move { + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(error) => return TaskResult::Recoverable(error), + }; + // A filter error above returns before reserving any rows. Once filtering has + // succeeded, reserve only the matching rows and defer projection construction until + // this task is polled by the runtime. + let mask = row_limit.limit(mask); + if mask.all_false() { + return TaskResult::Array(None); + } + + run_deferred_projection(ctx, row_range, mask).await + }) + } + + fn deferred_projection(ctx: Arc, row_range: Range, mask: Mask) -> Self { + Self::new(run_deferred_projection(ctx, row_range, mask)) + } +} + +/// Construct and run the projection for an already-reserved mask, classifying any failure as +/// terminal (rows have been reserved and cannot be released back to a concurrent limit). +async fn run_deferred_projection( + ctx: Arc, + row_range: Range, + mask: Mask, +) -> TaskResult { + let projection = + match ctx + .reader + .projection_evaluation(&row_range, &ctx.projection, MaskFuture::ready(mask)) + { + Ok(projection) => projection, + Err(error) => return TaskResult::Terminal(error), + }; + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => TaskResult::Terminal(error), + } +} + +impl Future for TaskFuture { + type Output = TaskResult; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.inner.as_mut().poll(cx) + } +} /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this @@ -33,133 +168,353 @@ pub type TaskFuture = BoxFuture<'static, VortexResult>; /// The intersected row range is then further reduced via expression-based pruning. After pruning /// has eliminated more blocks, the full filter is executed over the remainder of the split. /// -/// This mask is then provided to the reader to perform a filtered projection over the split data, -/// finally mapping the Vortex columnar record batches into some result type `A`. -pub fn split_exec( - ctx: Arc>, +/// The final mask is limited before it is given to the reader to perform a filtered projection +/// over the split data, yielding the projected array (or `None` when the split selects no rows). +/// Limiting before projection prevents decode work for rows that the scan cannot return. +pub(crate) fn split_exec( + ctx: Arc, read_mask: RowMask, - limit: Option<&mut u64>, -) -> VortexResult>> { + row_limit: Option, +) -> VortexResult { let row_range = read_mask.row_range(); let row_mask = read_mask.mask().clone(); - let filter_mask = match ctx.filter.as_ref() { - // No filter == immediate mask - None => { - let row_mask = match limit { - Some(l) if *l == 0 => Mask::new_false(row_mask.len()), - Some(l) => { - let true_count = row_mask.true_count(); - let mask_limit = usize::try_from(*l) - .map(|l| l.min(true_count)) - .unwrap_or(true_count); - let row_mask = row_mask.limit(mask_limit); - *l -= mask_limit as u64; - row_mask - } - None => row_mask, - }; - - MaskFuture::ready(row_mask) - } - Some(filter) => { - // NOTE: it's very important that the pruning and filter evaluations are built OUTSIDE - // the future. Registering these row ranges eagerly is a hint to the IO system that - // we want to start prefetching the IO for this split. - let reader = Arc::clone(&ctx.reader); - let filter = Arc::clone(filter); - let row_range = row_range.clone(); - - MaskFuture::new(row_mask.len(), async move { - let mut mask = row_mask; - let mut dynamic_versions = vec![None; filter.conjuncts().len()]; - - // TODO(ngates): we could use FuturedUnordered to intersect the masks in parallel. - for (idx, conjunct) in filter.conjuncts().iter().enumerate() { - if mask.all_false() { - return Ok(mask); - } - - // Store the latest version of the dynamic expression prior to pruning. - // We will re-run the pruning later if the version has changed in the meantime. - dynamic_versions[idx] = filter.dynamic_updates(idx).map(|du| du.version()); - - let conjunct_mask = reader - .pruning_evaluation(&row_range, conjunct, mask.clone())? - .await?; - mask = mask.bitand(&conjunct_mask); - } - - // Now we loop through the conjuncts in the preferred order and evaluate them. - let mut remaining = BitVec::from_elem(filter.conjuncts().len(), true); - while let Some(idx) = filter.next_conjunct(&remaining) { - remaining.set(idx, false); - if mask.all_false() { - return Ok(mask); - } - - let conjunct = &filter.conjuncts()[idx]; - - // If the dynamic expression has changed since pruning, re-run the pruning. - // Store the dynamic update once to avoid TOCTOU race condition - let current_version = filter.dynamic_updates(idx).map(|du| du.version()); - if let Some(dv) = current_version - && dynamic_versions[idx].is_none_or(|v| v < dv) - { - // The dynamic expression has been updated, re-run the pruning. - dynamic_versions[idx] = Some(dv); - let conjunct_mask = reader - .pruning_evaluation(&row_range, conjunct, mask.clone())? - .await?; - mask = mask.bitand(&conjunct_mask); - } - if mask.all_false() { - return Ok(mask); - } - - let conjunct_mask = reader - .filter_evaluation(&row_range, conjunct, MaskFuture::ready(mask))? - .await?; - filter.report_selectivity(idx, conjunct_mask.density()); - - // Filter evaluations return a mask already intersected with the input mask. - mask = conjunct_mask; - } - - Ok(mask) - }) + let Some(filter) = ctx.filter.as_ref() else { + let limited = row_limit.is_some(); + let row_mask = if let Some(limit) = row_limit { + limit.limit(row_mask) + } else { + row_mask + }; + if row_mask.all_false() { + return Ok(TaskFuture::empty()); } + + // With no filter, limit the selection before constructing projection work. + let projection = match ctx.reader.projection_evaluation( + &row_range, + &ctx.projection, + MaskFuture::ready(row_mask), + ) { + Ok(projection) => projection, + Err(err) if limited => return Ok(TaskFuture::terminal(err)), + Err(err) => return Err(err), + }; + return Ok(TaskFuture::projection(projection, limited)); }; - // Step 4: execute the projection, only at the mask for rows which match the filter - let projection_future = - ctx.reader - .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; + let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); + + let Some(row_limit) = row_limit else { + // Without a limit, retain the existing eager projection setup so readers can prefetch + // projection work while the filter is being evaluated. + let projection = + ctx.reader + .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; + return Ok(TaskFuture::filtered_projection(filter_mask, projection)); + }; + + Ok(TaskFuture::limited_filtered_projection( + ctx, + row_range, + filter_mask, + row_limit, + )) +} + +/// Evaluate the filter for a split and return the matching mask together with its row range. +/// +/// This is the first stage of the ordered filtered-limit pipeline. It performs the same pruning +/// and filter I/O as [`split_exec`] (registered eagerly, outside the returned future, so the IO +/// system can prefetch while earlier splits are still reserving), but it neither reserves against +/// the limit nor projects. The caller reserves the returned mask in split order and then projects +/// it via [`project_split`], which keeps ordered `LIMIT` semantics without serializing I/O. +pub(crate) fn filter_split( + ctx: Arc, + read_mask: RowMask, +) -> BoxFuture<'static, VortexResult<(Range, Mask)>> { + let row_range = read_mask.row_range(); + let row_mask = read_mask.mask().clone(); + let filter = ctx + .filter + .as_ref() + .vortex_expect("filter_split requires a filtered scan"); + let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); - let mapper = Arc::clone(&ctx.mapper); - let array_fut = async move { + async move { let mask = filter_mask.await?; - if mask.all_false() { - return Ok(None); + Ok((row_range, mask)) + } + .boxed() +} + +/// Project an already-reserved, non-empty mask for a split. +/// +/// This is the second stage of the ordered filtered-limit pipeline, run after the caller has +/// reserved rows against the limit in split order (see [`filter_split`]). +pub(crate) fn project_split( + ctx: Arc, + row_range: Range, + mask: Mask, +) -> TaskFuture { + TaskFuture::deferred_projection(ctx, row_range, mask) +} + +/// Build the filtered mask for a split. +/// +/// The pruning and filter evaluations are constructed OUTSIDE the returned future on purpose: +/// registering these row ranges eagerly is a hint to the IO system that we want to start +/// prefetching the IO for this split. +fn build_filter_mask( + reader: &Arc, + filter: &Arc, + row_range: &Range, + row_mask: Mask, +) -> MaskFuture { + let reader = Arc::clone(reader); + let filter = Arc::clone(filter); + let filter_row_range = row_range.clone(); + MaskFuture::new(row_mask.len(), async move { + let mut mask = row_mask; + let mut dynamic_versions = vec![None; filter.conjuncts().len()]; + + // TODO(ngates): we could use FuturedUnordered to intersect the masks in parallel. + for (idx, conjunct) in filter.conjuncts().iter().enumerate() { + if mask.all_false() { + return Ok(mask); + } + + // Store the latest version of the dynamic expression prior to pruning. + // We will re-run the pruning later if the version has changed in the meantime. + dynamic_versions[idx] = filter.dynamic_updates(idx).map(|du| du.version()); + + let conjunct_mask = reader + .pruning_evaluation(&filter_row_range, conjunct, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); } - let array = projection_future.await?; - mapper(array).map(Some) - }; + // Now we loop through the conjuncts in the preferred order and evaluate them. + let mut remaining = BitVec::from_elem(filter.conjuncts().len(), true); + while let Some(idx) = filter.next_conjunct(&remaining) { + remaining.set(idx, false); + if mask.all_false() { + return Ok(mask); + } + + let conjunct = &filter.conjuncts()[idx]; - Ok(array_fut.boxed()) + // If the dynamic expression has changed since pruning, re-run the pruning. + // Store the dynamic update once to avoid TOCTOU race condition. + let current_version = filter.dynamic_updates(idx).map(|du| du.version()); + if let Some(dv) = current_version + && dynamic_versions[idx].is_none_or(|v| v < dv) + { + // The dynamic expression has changed, re-run the pruning. + dynamic_versions[idx] = Some(dv); + let conjunct_mask = reader + .pruning_evaluation(&filter_row_range, conjunct, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); + } + if mask.all_false() { + return Ok(mask); + } + + let conjunct_mask = reader + .filter_evaluation(&filter_row_range, conjunct, MaskFuture::ready(mask))? + .await?; + filter.report_selectivity(idx, conjunct_mask.density()); + + // Filter evaluations return a mask already intersected with the input mask. + mask = conjunct_mask; + } + + Ok(mask) + }) } /// Information needed to execute a single split task. /// /// Row selection is evaluated before creating a split task so it's not included -pub struct TaskContext { +pub(crate) struct TaskContext { /// The shared filter expression. - pub filter: Option>, + pub(crate) filter: Option>, /// The layout reader. - pub reader: Arc, + pub(crate) reader: Arc, /// The projection expression to apply to gather the scanned rows. - pub projection: Expression, - /// Function that maps into an A. - pub mapper: Arc VortexResult + Send + Sync>, + pub(crate) projection: Expression, +} + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::Arc; + use std::sync::mpsc; + use std::thread; + use std::time::Duration; + + use parking_lot::Mutex; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::TaskContext; + use super::TaskResult; + use super::project_split; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; + + struct BlockingProjectionReader { + name: Arc, + dtype: DType, + started: mpsc::Sender<()>, + gate: Arc>, + } + + impl LayoutReader for BlockingProjectionReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 1 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + _mask: MaskFuture, + ) -> VortexResult { + self.started + .send(()) + .map_err(|_| vortex_err!("test projection-start receiver dropped"))?; + let _guard = self.gate.lock(); + let array = PrimitiveArray::from_iter([0_i32]).into_array(); + Ok(Box::pin(async move { Ok(array) })) + } + } + + #[test] + fn project_split_defers_projection_construction_until_task_poll() -> VortexResult<()> { + let gate = Arc::new(Mutex::new(())); + let guard = gate.lock(); + let (started_send, started_recv) = mpsc::channel(); + let reader: Arc = Arc::new(BlockingProjectionReader { + name: Arc::from("blocking-projection"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + started: started_send, + gate: Arc::clone(&gate), + }); + let ctx = Arc::new(TaskContext { + filter: None, + reader, + projection: root(), + }); + + let (task_send, task_recv) = mpsc::channel(); + let construction = thread::spawn(move || { + let task = project_split(ctx, 0..1, Mask::new_true(1)); + drop(task_send.send(task)); + }); + + let task = match task_recv.recv_timeout(Duration::from_secs(1)) { + Ok(task) => task, + Err(_) => { + drop(guard); + drop(construction.join()); + return Err(vortex_err!( + "constructing a projection task blocked before it was polled" + )); + } + }; + if construction.join().is_err() { + return Err(vortex_err!("projection-task construction panicked")); + } + match started_recv.try_recv() { + Err(mpsc::TryRecvError::Empty) => {} + Ok(()) => { + return Err(vortex_err!( + "projection construction started before the task was polled" + )); + } + Err(mpsc::TryRecvError::Disconnected) => { + return Err(vortex_err!( + "projection-start sender disconnected unexpectedly" + )); + } + } + + let (result_send, result_recv) = mpsc::channel(); + let execution = thread::spawn(move || { + let result = futures::executor::block_on(task); + drop(result_send.send(result)); + }); + + if started_recv.recv_timeout(Duration::from_secs(1)).is_err() { + drop(guard); + drop(execution.join()); + return Err(vortex_err!( + "projection construction did not start when polled" + )); + } + drop(guard); + let result = result_recv + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| vortex_err!("projection task did not complete"))?; + if execution.join().is_err() { + return Err(vortex_err!("projection task panicked")); + } + + assert!(matches!(result, TaskResult::Array(Some(_)))); + Ok(()) + } } diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 9550067b956..214b1a6244b 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -208,7 +208,7 @@ fn scan_builder( indices: Option, batch_size: Option, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult { let mut builder = vxf .scan()? .with_some_filter(expr) diff --git a/vortex-python/src/scan.rs b/vortex-python/src/scan.rs index cfdf77a6b8b..8ee95f7f8d5 100644 --- a/vortex-python/src/scan.rs +++ b/vortex-python/src/scan.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; -use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; use vortex::error::VortexResult; use vortex::layout::scan::repeated_scan::RepeatedScan; @@ -30,7 +29,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { #[pyclass(name = "RepeatedScan", module = "vortex", frozen)] pub struct PyRepeatedScan { - pub scan: Arc>, + pub scan: Arc, pub row_count: u64, } diff --git a/vortex-scan/src/lib.rs b/vortex-scan/src/lib.rs index 2f2e6e1c429..58877399422 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -139,10 +139,13 @@ pub struct ScanRequest { /// Partition range to scan, which allows readers to skip unwanted partitions. pub partition_range: Option>, /// Whether the scan should preserve row order. If false, the scan may produce rows in any - /// order, for example to enable parallel execution across partitions. + /// order, for example to enable parallel execution across partitions. An ordered global + /// limit intentionally serializes external partitions so the returned rows are the first + /// matching rows in scan order. pub ordered: bool, - /// Optional limit on the number of rows returned by scan. Limits are applied after all - /// filtering and row selection. + /// Optional global limit on the number of rows returned by the scan. Limits are applied after + /// filtering and row selection, but implementations must trim the resulting selection mask + /// before projection so they do not decode rows that cannot be returned. pub limit: Option, }