diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 29d2115ed49..2c1089afdf5 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -7,10 +7,8 @@ //! * **write** runs the full default write pipeline — repartition into row //! blocks, zoned statistics, dictionary probe, coalesce, compress with the //! forced string scheme and its children, layout, serialize into a buffer; -//! * **read** opens that buffer and runs the scan, decoding each row split to -//! canonical `VarBinViewArray` form inside its own scan task and dropping it -//! before the next split runs — the shape production uses, where -//! `into_record_batch_stream` fuses the Arrow conversion into the split task. +//! * **read** opens that buffer and runs the scan, decoding each emitted chunk to +//! canonical `VarBinViewArray` form and dropping it before polling the next chunk. //! //! Both run on a current-thread runtime, so these are single-threaded CPU costs //! and exclude physical I/O. @@ -201,35 +199,20 @@ async fn write_serialized_file( } /// Time one complete read of a serialized Vortex buffer: open the file, then run -/// the scan with the canonical decode fused into each row split's task, dropping -/// each decoded chunk before the next split runs. -/// -/// The splits are awaited one at a time rather than through -/// `ScanBuilder::into_array_stream`, which spawns -/// `concurrency * available_parallelism()` of them at once. On a current-thread -/// runtime that read-ahead buys no parallelism; it only holds that many chunks in -/// memory and makes the result depend on the host's core count. Awaiting one at a -/// time keeps the per-split work identical to production while making the -/// measurement machine-independent. +/// the scan and decode each emitted chunk before polling the next one. async fn read_serialized_buffer(session: &VortexSession, data: Bytes) -> Result { let decode_session = session.clone(); let start = Instant::now(); let file = session.open_options().open_buffer(data)?; - let splits = file - .scan()? - .map(move |chunk: ArrayRef| { - let mut ctx = decode_session.create_execution_ctx(); - chunk.execute::(&mut ctx) - }) - .build()?; + let mut chunks = file.scan()?.into_stream()?; let mut rows = 0usize; - for split in splits { - if let Some(canonical) = split.await? { - rows += canonical.len(); - drop(black_box(canonical)); - } + while let Some(chunk) = chunks.try_next().await? { + let mut ctx = decode_session.create_execution_ctx(); + let canonical = chunk.execute::(&mut ctx)?; + rows += canonical.len(); + drop(black_box(canonical)); } black_box(rows); diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index c57bc91a65d..12e0150680f 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 vortex::array::ArrayRef; use vortex::array::Canonical; @@ -72,15 +73,17 @@ impl Dataset for TPCHLCommentChunked { let chunks: Vec<_> = file .scan()? .with_projection(projection) + .into_array_stream()? .map({ let ctx = ctx.clone(); move |a| { let mut ctx = ctx.clone(); - let canonical = a.execute::(&mut ctx)?; - Ok(canonical.into_array()) + a.and_then(|a| { + let canonical = a.execute::(&mut ctx)?; + Ok(canonical.into_array()) + }) } }) - .into_array_stream()? .try_collect() .await?; 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 f5fa2147100..88feae77d4b 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -397,9 +397,7 @@ impl FileOpener for VortexOpener { .transpose() .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; - if let Some(limit) = limit - && filter.is_none() - { + if let Some(limit) = limit { scan_builder = scan_builder.with_limit(limit); } @@ -445,28 +443,29 @@ 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_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(move |batch| { + .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) + .map(move |batch| -> DFResult { let batch = if projector.projection().as_ref().is_empty() { batch } else { @@ -549,10 +548,10 @@ impl NaturalSplits { } /// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. -fn natural_splits_for_file( +fn natural_splits_for_file( natural_splits: &DashMap>, path: &Path, - scan_builder: &ScanBuilder, + scan_builder: &ScanBuilder, total_size: u64, ) -> DFResult> { if let Some(splits) = natural_splits.get(path) { @@ -574,8 +573,8 @@ fn natural_splits_for_file( /// Walk the layout tree to compute the file's full natural split boundaries for the fields /// referenced by the scan's projection and filter. -fn compute_natural_splits( - scan_builder: &ScanBuilder, +fn compute_natural_splits( + scan_builder: &ScanBuilder, total_size: u64, ) -> DFResult> { let row_boundaries = scan_builder @@ -636,6 +635,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; @@ -668,6 +673,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; @@ -1086,6 +1092,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-file/src/file.rs b/vortex-file/src/file.rs index a235011b9c0..fd5e879f3f2 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; @@ -214,7 +213,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 f5c177c9cdf..54ce7512f2e 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1328,7 +1328,7 @@ async fn write_nullable_top_level_struct() { async fn round_trip( array: &ArrayRef, - f: impl FnOnce(ScanBuilder) -> VortexResult>, + f: impl FnOnce(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 94ce7acf277..c0c6c1808ba 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`]. @@ -168,6 +169,13 @@ impl DataSource for LayoutReaderDataSource { } } + // Only unordered scans share a limit across external partitions: reservation order is + // completion order, which an ordered scan cannot accept. Ordered partitions each apply + // the limit locally and the engine trims the concatenated result. + let row_limit = (!scan_request.ordered) + .then(|| scan_request.limit.map(RowLimit::new)) + .flatten(); + Ok(Box::new(LayoutReaderScan { reader: Arc::clone(&self.reader), session: self.session.clone(), @@ -175,6 +183,7 @@ impl DataSource for LayoutReaderDataSource { projection, filter, limit: scan_request.limit, + row_limit, selection: scan_request.selection, ordered: scan_request.ordered, metrics_registry: self.metrics_registry.clone(), @@ -196,6 +205,7 @@ struct LayoutReaderScan { projection: BoundExpression, filter: Option, limit: Option, + row_limit: Option, ordered: bool, selection: Selection, metrics_registry: Option>, @@ -236,6 +246,9 @@ 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 @@ -249,7 +262,8 @@ impl Stream for LayoutReaderScan { // 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. + // each split; a shared `row_limit` (unordered scans) caps the total, and otherwise + // the engine enforces the exact limit at the stream level. if this.filter.is_none() && let Some(ref mut limit) = this.limit { @@ -262,6 +276,7 @@ impl Stream for LayoutReaderScan { projection: this.projection.clone(), filter: this.filter.clone(), limit: split_limit, + row_limit: this.row_limit.clone(), ordered: this.ordered, row_range, selection: this.selection.clone(), @@ -274,7 +289,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; @@ -289,6 +307,7 @@ struct LayoutReaderSplit { projection: BoundExpression, filter: Option, limit: Option, + row_limit: Option, ordered: bool, row_range: Range, selection: Selection, @@ -311,7 +330,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) @@ -329,13 +348,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, @@ -402,3 +422,60 @@ impl Partition for Empty { ))) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use futures::StreamExt; + use futures::TryStreamExt; + use parking_lot::Mutex; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; + + use super::LayoutReaderDataSource; + use crate::scan::test::TestLayoutReader; + use crate::scan::test::collect_scan_values; + use crate::scan::test::session_with_handle; + + /// An unordered limit is shared by every partition of the scan, and is applied to each split's + /// mask, so the partitions together never project more rows than the limit can return. + #[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::>())?; + assert_eq!(partitions.len(), 6); + + let chunks = runtime.block_on( + futures::stream::iter(partitions) + .map(|partition| partition.execute()) + .try_flatten_unordered(Some(6)) + .try_collect::>(), + )?; + let values = collect_scan_values(chunks.into_iter().map(Ok))?; + + assert_eq!(values.len(), 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..63335a33d71 --- /dev/null +++ b/vortex-layout/src/scan/limit.rs @@ -0,0 +1,129 @@ +// 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 granted = self.take(mask.true_count()); + mask.limit(granted) + } + + /// Reserve up to `rows` rows, returning how many the remaining budget granted. + pub(crate) fn take(&self, rows: usize) -> usize { + let requested = u64::try_from(rows).unwrap_or(u64::MAX); + usize::try_from(self.reserve(requested)).unwrap_or(usize::MAX) + } + + 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 d9de768f271..a45f4859c0d 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -61,6 +61,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. @@ -306,11 +307,19 @@ impl DataSource for MultiLayoutDataSource { let request = BoundScanRequest::try_new(scan_request, &self.dtype)?; let dtype = request.projection.dtype().clone(); + // Only unordered scans share a limit across external partitions: reservation order is + // completion order, which an ordered scan cannot accept. Ordered partitions each apply the + // limit locally and the engine trims the concatenated result. + let row_limit = (!request.ordered) + .then(|| request.limit.map(RowLimit::new)) + .flatten(); + Ok(Box::new(MultiLayoutScan { session: self.session.clone(), source_dtype: self.dtype.clone(), dtype, request, + row_limit, ready, deferred, handle: self.session.handle(), @@ -369,6 +378,7 @@ struct MultiLayoutScan { source_dtype: DType, dtype: DType, request: BoundScanRequest, + row_limit: Option, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -395,6 +405,7 @@ impl DataSourceScan for MultiLayoutScan { source_dtype, dtype: _, request, + row_limit, ready, deferred, handle, @@ -449,9 +460,14 @@ 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(), &source_dtype, request.clone()) - } + Ok(reader) => reader_partition( + i, + reader, + session.clone(), + &source_dtype, + request.clone(), + row_limit.clone(), + ), Err(e) => stream::once(async move { Err(e) }).boxed(), }) .boxed() @@ -469,6 +485,7 @@ fn reader_partition( session: VortexSession, source_dtype: &DType, request: BoundScanRequest, + row_limit: Option, ) -> PartitionStream { if reader.dtype() != source_dtype { let error = vortex_err!( @@ -523,6 +540,7 @@ fn reader_partition( row_range: Some(row_range), ..request }, + row_limit, index: partition_idx, }) as PartitionRef) }) @@ -537,6 +555,7 @@ struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, request: BoundScanRequest, + row_limit: Option, index: usize, } @@ -564,7 +583,7 @@ impl Partition for MultiLayoutPartition { Ok(filter) => filter.is_some(), Err(_) => true, }; - if has_filter { + if has_filter || self.row_limit.is_some() { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -583,6 +602,7 @@ impl Partition for MultiLayoutPartition { .with_projection(request.projection) .with_some_filter(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 { @@ -590,7 +610,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, @@ -600,15 +620,27 @@ impl Partition for MultiLayoutPartition { #[cfg(test)] mod tests { + use std::sync::Arc; + + use async_trait::async_trait; + use futures::TryStreamExt; use rstest::rstest; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; use super::*; + use crate::scan::test::TestLayoutReader; + use crate::scan::test::collect_scan_values; use crate::scan::test::new_session; + use crate::scan::test::session_with_handle; struct NeverOpened; @@ -655,4 +687,51 @@ mod tests { assert!(request.filter.is_err()); Ok(()) } + + struct StaticReaderFactory { + reader: LayoutReaderRef, + } + + #[async_trait] + impl LayoutReaderFactory for StaticReaderFactory { + async fn open(&self) -> VortexResult> { + Ok(Some(Arc::clone(&self.reader))) + } + } + + /// An unordered limit is shared by every file of the scan, so the files together return no + /// more than the limit even though each one is scanned independently. + #[test] + fn unordered_limit_is_shared_across_readers() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let first: LayoutReaderRef = Arc::new(TestLayoutReader::new(2)); + let second: LayoutReaderRef = Arc::new(TestLayoutReader::new(2).with_base(10)); + 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: false, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 2); + + let mut values = Vec::new(); + for partition in partitions { + values.extend(collect_scan_values( + runtime.block_on_stream(partition.execute()?), + )?); + } + + // Three of the four rows, whichever partition reserved them first. + assert_eq!(values.len(), 3); + Ok(()) + } } diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 413761b8103..083a3cab4a3 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,18 @@ 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::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: BoundExpression, @@ -49,15 +60,164 @@ 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))) + } + + /// Build split tasks lazily, stopping as soon as `gate` has no budget left. + /// + /// `reserve` is the limit that each split reserves its filtered rows against before + /// projecting; pass `None` to leave the limit to the caller (see + /// [`ScheduledTaskStream::with_row_limit`]). + fn lazy( + handle: Handle, + split_ranges: Vec>, + selection: Selection, + ctx: Arc, + gate: RowLimit, + reserve: Option, + ) -> Self { + Self::new( + futures::stream::iter(split_ranges) + .take_while(move |_| future::ready(!gate.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 the buffered stream 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, reserve.clone()) + .unwrap_or_else(TaskFuture::recoverable) + }); + future::ready(task.map(|task| 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>, + /// A limit applied to the emitted arrays, for scans that could not reserve rows per split. + row_limit: 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), + row_limit: None, + } + } + + /// Trim emitted arrays to `row_limit`, ending the stream once its budget is spent. + /// + /// Used by ordered limited scans, where reserving per split would hand the budget to whichever + /// split filters first rather than to the earliest split. Because this stream yields in split + /// order, taking rows as they are emitted keeps the limit exact. + fn with_row_limit(mut self, row_limit: RowLimit) -> Self { + self.row_limit = Some(row_limit); + self + } + + /// Apply [`Self::row_limit`] to an emitted array, returning `None` once the budget is spent. + fn take_rows(&mut self, array: ArrayRef) -> Option> { + let Some(row_limit) = &self.row_limit else { + return Some(Ok(array)); + }; + + let granted = row_limit.take(array.len()); + let trimmed = granted < array.len(); + if row_limit.is_exhausted() { + // No later split can contribute now, so drop the queued and in-flight task handles. + // Their `Drop` implementations abort the work started for them. + self.tasks = None; + } + + match granted { + 0 => None, + _ if trimmed => Some(array.slice(0..granted)), + _ => Some(Ok(array)), + } + } +} + +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)) => { + if let Some(array) = self.take_rows(array) { + return Poll::Ready(Some(array)); + } + return Poll::Ready(None); + } + 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 +241,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: BoundExpression, @@ -99,8 +257,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 +271,28 @@ impl RepeatedScan { selection, splits, concurrency, - map_fn, limit, + row_limit, dtype, } } - pub fn execute( - &self, - row_range: Option>, - ) -> VortexResult>>>> { + fn has_filter(&self) -> bool { + self.filter.is_some() + } + + fn task_context(&self) -> Arc { + Arc::new(TaskContext { + filter: self + .filter + .clone() + .map(|filter| Arc::new(FilterExpr::new(filter))), + reader: Arc::clone(&self.layout_reader), + projection: self.projection.clone(), + }) + } + + 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 +305,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 +324,104 @@ 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; + pub(crate) fn execute( + &self, + row_range: Option>, + row_limit: Option, + ) -> VortexResult> { let mut tasks = Vec::new(); - let ctx = 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), - }); + let ctx = self.task_context(); + + 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 earlier work drains + // the budget, `take_while` prevents further task creation. + if self.has_filter() + && let Some(row_limit) = row_limit + { + // Reserving rows per split hands the budget to whichever split finishes filtering + // first, which an ordered scan cannot accept: it must return the *earliest* matching + // rows. Ordered scans therefore filter and project without reserving, and take rows + // from the (in-order) output instead. Unordered scans reserve before projecting, so + // rows they cannot return are never decoded. + let reserve = (!self.ordered).then(|| row_limit.clone()); + let tasks = TaskStream::lazy( + handle, + self.split_ranges(row_range), + self.selection.clone(), + self.task_context(), + row_limit.clone(), + reserve.clone(), + ); - let stream = if self.ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; + let stream = ScheduledTaskStream::new(tasks, self.ordered, concurrency); + return Ok(match reserve { + Some(_) => stream, + None => stream.with_row_limit(row_limit), + }); + } + + // 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`, + // which reserves in split order and so stays exact for ordered scans too. + let tasks = TaskStream::eager(handle, self.execute(row_range, row_limit)?); - Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) + let ordered = self.ordered; + Ok(ScheduledTaskStream::new(tasks, ordered, concurrency)) } } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index f7abbb21fb2..52297de66f4 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; @@ -20,32 +19,29 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::analysis::referenced_field_paths; 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_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_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::layouts::row_idx::RowIdx; 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: /// @@ -55,7 +51,7 @@ use crate::scan::splits::attempt_split_ranges; /// /// Projection and filter expressions must be bound against the reader dtype. 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: BoundExpression, @@ -74,19 +70,17 @@ pub struct ScanBuilder { natural_splits: Option>, /// 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 { let projection = BoundExpression::new_root(layout_reader.dtype().clone()); @@ -103,10 +97,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, } } @@ -134,7 +127,7 @@ impl ScanBuilder { } } -impl ScanBuilder { +impl ScanBuilder { /// Add a filter expression bound against the reader dtype. pub fn with_filter(mut self, filter: BoundExpression) -> Self { self.filter = Some(filter); @@ -262,6 +255,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 { Ok(self.projection.dtype().clone()) @@ -272,39 +271,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, - natural_splits: self.natural_splits, - 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; @@ -357,26 +327,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)) } @@ -384,81 +347,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(); - let task = handle - .spawn_cpu(move || builder.prepare().and_then(|scan| scan.execute(None))); - self.state = LazyScanState::Preparing(PreparingScan { - ordered, - concurrency, - handle, - task, + // IMPORTANT: Building the stream can synchronously walk the layout and + // register I/O for every split. Keep it with preparation in this CPU + // 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_cpu(move || { + 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)), } } @@ -507,6 +449,7 @@ mod test { use futures::Stream; use futures::task::noop_waker_ref; use parking_lot::Mutex; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; @@ -529,6 +472,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; @@ -538,6 +482,11 @@ mod test { use crate::RowSplits; use crate::SplitRange; use crate::scan::test::SCAN_SESSION; + use crate::scan::test::TestLayoutReader; + use crate::scan::test::collect_scan_values; + use crate::scan::test::drain_runtime; + use crate::scan::test::keep_all; + use crate::scan::test::keep_odd; use crate::scan::test::session_with_handle; fn nested_dtype() -> DType { @@ -719,6 +668,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 { @@ -728,8 +684,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 { @@ -782,6 +749,14 @@ mod test { _expr: &BoundExpression, _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) @@ -852,6 +827,206 @@ mod test { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + 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(()) + } + + #[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( + TestLayoutReader::new(row_count) + .with_split_size(2) + .with_keep_row(keep_row), + ); + let filter = root().bind(reader.dtype())?; + + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + .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(()) + } + + /// An unordered filtered limit reserves rows before projecting, so a huge split never decodes + /// more rows than the limit can return. + #[test] + fn unordered_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( + TestLayoutReader::new(100_000).with_projection_masks(Arc::clone(&projection_masks)), + ); + let filter = root().bind(reader.dtype())?; + + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + .with_limit(1) + .with_ordered(false) + .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(()) + } + + /// An ordered filtered limit cannot reserve per split (that would grant the budget to whichever + /// split filters first), so it trims the in-order output instead. + #[test] + fn ordered_filtered_limit_trims_the_emitted_rows() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new(TestLayoutReader::new(8).with_split_size(4)); + let filter = root().bind(reader.dtype())?; + + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + .with_limit(6) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + assert_eq!(values, [0, 1, 2, 3, 4, 5]); + 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( + TestLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_filter(), + ); + let filter = root().bind(reader.dtype())?; + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + .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(()) + } + + /// Rows reserved against a shared limit cannot be released back, so a projection failure after + /// reservation must end the scan rather than let a later split spend the freed budget. + #[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( + TestLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_projection(), + ); + let filter = root().bind(reader.dtype())?; + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + // 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) + .with_ordered(false) + .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 full_file_splits_ignore_row_range() -> VortexResult<()> { let calls = Arc::new(AtomicUsize::new(0)); @@ -862,6 +1037,52 @@ mod test { .full_file_splits()?; assert_eq!(splits, [0, 1, 2, 3, 4]); + + 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( + TestLayoutReader::new(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_projection_error(), + ); + let filter = root().bind(reader.dtype())?; + let stream = ScanBuilder::new(session, reader) + .with_filter(filter) + .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( + TestLayoutReader::new(8) + .with_split_size(2) + .with_keep_row(keep_odd), + ); + let filter = root().bind(reader.dtype())?; + + let scan = ScanBuilder::new(session, reader) + .with_filter(filter) + .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(()) } diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 218efb64a0d..50a28fa21c8 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,140 @@ use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::expr::BoundExpression; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; -use crate::LayoutReader; +use crate::ArrayFuture; +use crate::LayoutReaderRef; 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 { + pub(crate) 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)) + } + + /// Project a split whose mask needed no filtering. + /// + /// `terminal` marks that the mask already reserved rows against a limit, so a failure cannot + /// be reported without ending the scan. + pub(crate) 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), + } + }) + } + + /// Project the rows a filter matched, skipping projection entirely for an empty split. + /// + /// The projection is constructed by the caller, before the filter has run, so that the reader + /// can prefetch its I/O. + pub(crate) 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), + } + }) + } + + /// Filter, reserve the matching rows against `row_limit`, then project only what was granted. + /// + /// Projection work is constructed after reservation, so rows the limit cannot grant are never + /// decoded. Once rows have been reserved they cannot be released back to a concurrent limit, + /// so any projection failure is terminal. + 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. + let mask = row_limit.limit(mask); + if mask.all_false() { + return TaskResult::Array(None); + } + + 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 +164,142 @@ 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, + )) +} + +/// 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: &LayoutReaderRef, + 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); + } - let mapper = Arc::clone(&ctx.mapper); - let array_fut = async move { - let mask = filter_mask.await?; - if mask.all_false() { - return Ok(None); + // 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]; + + // 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(array_fut.boxed()) + 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 { +/// Row selection is evaluated before creating a split task so it's not included. +pub(crate) struct TaskContext { /// The shared filter expression. - pub filter: Option>, + pub(crate) filter: Option>, /// The layout reader. - pub reader: Arc, + pub(crate) reader: LayoutReaderRef, /// The projection expression to apply to gather the scanned rows. - pub projection: BoundExpression, - /// Function that maps into an A. - pub mapper: Arc VortexResult + Send + Sync>, + pub(crate) projection: BoundExpression, } diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 809639d169d..d2d52102eb1 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -1,17 +1,41 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; +use std::sync::Arc; use std::sync::LazyLock; +use std::task::Poll; +use parking_lot::Mutex; +use vortex_array::ArrayRef; +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::BoundExpression; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::Handle; +use vortex_io::runtime::single::SingleThreadRuntime; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_session::VortexSession; +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::RowSplits; +use crate::SplitRange; use crate::session::LayoutSession; pub fn new_session() -> VortexSession { - vortex_array::array_session() + array_session() .with::() .with::() } @@ -21,3 +45,216 @@ pub fn session_with_handle(handle: Handle) -> VortexSession { } pub static SCAN_SESSION: LazyLock = LazyLock::new(new_session); + +/// A configurable [`LayoutReader`] test double producing `base + row` for every selected row. +/// +/// `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 that a limit is applied before +/// projection. +#[derive(Debug)] +pub struct TestLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + base: i32, + split_size: Option, + keep_row: fn(u64) -> bool, + fail_first_filter: bool, + fail_first_projection: bool, + fail_projection: bool, + projection_masks: Option>>>, +} + +impl TestLayoutReader { + pub fn new(row_count: u64) -> Self { + Self { + name: Arc::from("test"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + base: 0, + split_size: None, + keep_row: keep_all, + fail_first_filter: false, + fail_first_projection: false, + fail_projection: false, + projection_masks: None, + } + } + + /// Offset every produced value, so tests can tell which reader a row came from. + pub fn with_base(mut self, base: i32) -> Self { + self.base = base; + self + } + + pub fn with_split_size(mut self, split_size: u64) -> Self { + self.split_size = Some(split_size); + self + } + + pub fn with_keep_row(mut self, keep_row: fn(u64) -> bool) -> Self { + self.keep_row = keep_row; + self + } + + pub fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = Some(projection_masks); + self + } + + pub fn with_fail_first_filter(mut self) -> Self { + self.fail_first_filter = true; + self + } + + pub fn with_fail_first_projection(mut self) -> Self { + self.fail_first_projection = true; + self + } + + pub fn with_projection_error(mut self) -> Self { + self.fail_projection = true; + self + } +} + +impl LayoutReader for TestLayoutReader { + 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: &BoundExpression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &BoundExpression, + 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: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let base = self.base; + let projection_masks = self.projection_masks.clone(); + let fail = self.fail_projection || (self.fail_first_projection && row_range.start == 0); + + Ok(Box::pin(async move { + let mask = mask.await?; + if let Some(projection_masks) = projection_masks { + 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).map(|value| base + value)) + .into_array() + .filter(mask) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +pub fn keep_all(_: u64) -> bool { + true +} + +pub fn keep_odd(row: u64) -> bool { + row % 2 == 1 +} + +/// Canonicalize every chunk of a scan into a flat list of `i32` values. +pub 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) +} + +/// Let already-spawned scan tasks run to completion after the stream has been dropped. +pub 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 + } + })); + } +} diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 121cfce52e9..4348cf16165 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -231,7 +231,7 @@ fn scan_builder( indices: Option, batch_size: Option, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult { let projection = projection .unwrap_or_else(root) .optimize_recursive(vxf.dtype())? 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 3aba97cda8a..ea009a1bdd4 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -142,8 +142,12 @@ pub struct ScanRequest { /// 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. pub ordered: bool, - /// Optional limit on the number of rows returned by scan. Limits are applied after all - /// filtering and row selection. + /// Optional limit on the number of rows returned, applied after filtering and row selection. + /// + /// For an unordered scan the limit is global: partitions share it, and each one trims its + /// selection mask before projection so that rows which cannot be returned are never decoded. + /// An ordered scan cannot share a budget whose reservation order is completion order, so each + /// partition applies the limit locally and the caller must trim the concatenated result. pub limit: Option, }